1 /*
2 * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25 package hat;
26
27 import hat.buffer.DispatchContext;
28 import optkl.OpHelper;
29 import optkl.ifacemapper.Buffer;
30 import optkl.util.carriers.ArenaAndLookupCarrier;
31 import optkl.ifacemapper.BufferTracker;
32 import hat.callgraph.ComputeCallGraph;
33 import hat.callgraph.KernelCallGraph;
34 import optkl.ifacemapper.MappableIface;
35 import jdk.incubator.code.dialect.core.CoreOp.FuncOp;
36 import jdk.incubator.code.Op;
37 import jdk.incubator.code.Quoted;
38 import jdk.incubator.code.dialect.java.JavaOp;
39 import jdk.incubator.code.dialect.java.MethodRef;
40
41 import java.lang.foreign.Arena;
42 import java.lang.invoke.MethodHandles;
43 import java.lang.reflect.Method;
44 import java.lang.reflect.Parameter;
45 import java.util.ArrayList;
46 import java.util.List;
47 import java.util.Map;
48 import java.util.Optional;
49 import java.util.concurrent.ConcurrentHashMap;
50
51 import static optkl.OpHelper.Invoke.invoke;
52
53 /**
54 * A ComputeContext is created by an Accelerator to capture and control compute and kernel
55 * callgraphs for the work to be performed by the backend.
56 * <p/>
57 * The Compute closure is created first, by walking the code model of the entrypoint, then transitively
58 * visiting all conventional code reachable from this entrypoint.
59 * <p/>
60 * Generally all user defined methods reachable from the entrypoint (and the entrypoint itself) must be static methods of the same
61 * enclosing classes.
62 * <p/>
63 * We do allow calls on the ComputeContext itself, and on the mapped interface buffers holding non uniform kernel data.
64 * <p/>
65 * Each request to dispatch a kernel discovered in the compute graph, results in a new Kernel call graph
66 * being created with the dispatched kernel as it's entrypoint.
67 * <p/>
68 * When the ComputeContext is finalized, it is passed to the backend via <a href="Backend.computeClosureHandoff(ComputeContext)"></a>
69 *
70 * @author Gary Frost
71 */
72 public class ComputeContext implements ArenaAndLookupCarrier, BufferTracker {
73
74 @Override
75 public Arena arena() {
76 return accelerator.arena();
77 }
78
79 @Override
80 public MethodHandles.Lookup lookup() {
81 return accelerator.lookup();
82 }
83
84
85 public Config config() {
86 return accelerator().config();
87 }
88
89 public void invokeWithArgs(Object[] args) {
90 computeCallGraph.invokeWithArgs(args);
91 }
92
93 /**
94 * Suggest copy out on exit of the compute layer for the specified buffers. This
95 * is mainly used when HAT is run with minimized copies (MC) flag.
96 *
97 * <p>This call has no effect if minimized copies is not used.</p>
98 */
99 public void copyOutOnExit(Buffer ... buffers) {
100 if (buffers != null && buffers.length > 0) {
101 for (Buffer buffer : buffers) {
102 accelerator.preAccess(buffer);
103 }
104 }
105 }
106
107 public enum WRAPPER {
108 MUTATE("Mutate"), ACCESS("Access");
109 public final MethodRef pre;
110 public final MethodRef post;
111
112 WRAPPER(String name) {
113 this.pre = MethodRef.method(ComputeContext.class, "pre" + name, void.class, MappableIface.class);
114 this.post = MethodRef.method(ComputeContext.class, "post" + name, void.class, MappableIface.class);
115 }
116 }
117
118 private final Accelerator accelerator;
119
120 public final Accelerator accelerator(){
121 return accelerator;
122 }
123
124 private final ComputeCallGraph computeCallGraph;
125
126 public final ComputeCallGraph computeCallGraph(){
127 return computeCallGraph;
128 }
129
130 /**
131 * Called by the Accelerator when the accelerator is passed a compute entrypoint.
132 * <p>
133 * So given a ComputeClass such as..
134 * <pre>
135 * public class MyComputeClass {
136 * @Reflect
137 * public static void addDeltaKernel(KernelContext kc, S32Array arrayOfInt, int delta) {
138 * arrayOfInt.array(kc.x, arrayOfInt.array(kc.x)+delta);
139 * }
140 *
141 * @Reflect
142 * static public void doSomeWork(final ComputeContext cc, S32Array arrayOfInt) {
143 * cc.dispatchKernel(KernelContext kc -> addDeltaKernel(kc,arrayOfInt.length(), 5, arrayOfInt);
144 * }
145 * }
146 * </pre>
147 *
148 * @param accelerator
149 * @param computeMethod
150 */
151
152 protected ComputeContext(Accelerator accelerator, Method computeMethod) {
153 this.accelerator = accelerator;
154 Optional<FuncOp> funcOp = Op.ofMethod(computeMethod);
155 if (funcOp.isEmpty()) {
156 throw new RuntimeException("Failed to create ComputeCallGraph (did you miss @Reflect annotation?).");
157 }
158 this.computeCallGraph = new ComputeCallGraph(this, computeMethod, funcOp.get());
159 this.accelerator.backend.computeContextHandoff(this);
160 }
161
162 public record KernelCallSite(Quoted<JavaOp.LambdaOp> quoted, JavaOp.LambdaOp lambdaOp, MethodRef methodRef, KernelCallGraph kernelCallGraph, Object[] capturedArgs, DispatchContext dispatchContext) {}
163
164 private record ConstantArgument(int paramIndex, Class<?> type, Object value) {
165 public static ConstantArgument of(int i, Class<?> type, Object capturedValue) {
166 // We only specialized for a small set of types
167 if (type == int.class && capturedValue instanceof Number value) {
168 return new ConstantArgument(i, type, value);
169 } else if (type == float.class && capturedValue instanceof Float value) {
170 return new ConstantArgument(i, type, value);
171 }
172 throw new IllegalStateException("Input constant of type: " + type.getName() + " not supported");
173 }
174 }
175
176 private record SpecializationKey(List<ConstantArgument> arguments) {
177 public SpecializationKey {
178 arguments = List.copyOf(arguments);
179 }
180
181 public static SpecializationKey of(Method kernelMethod, Object[] quotedCapturedValues) {
182 Parameter[] parameters = kernelMethod.getParameters();
183 List<ConstantArgument> arguments = new ArrayList<>();
184 for (int i = 0; i < quotedCapturedValues.length; i++) {
185 if (parameters[i].getType().isPrimitive()) {
186 arguments.add(ConstantArgument.of(i, parameters[i].getType(), quotedCapturedValues[i]));
187 }
188 }
189 if (arguments.isEmpty()) {
190 return empty();
191 }
192 return new SpecializationKey(arguments);
193 }
194
195 public static SpecializationKey empty() {
196 return new SpecializationKey(List.of());
197 }
198 }
199
200 private final Map<Op.Location, Map<SpecializationKey, KernelCallSite>> kernelCallSiteCache = new ConcurrentHashMap<>();
201
202 static OpHelper.Invoke getTargetInvoke(MethodHandles.Lookup lookup, JavaOp.LambdaOp lambdaOp) {
203 return lambdaOp.body().entryBlock().ops().stream()
204 .filter(ce -> ce instanceof JavaOp.InvokeOp)
205 .map(ce -> (OpHelper.Invoke)invoke(lookup, ce))
206 .filter(invoke->!invoke.refIs(ComputeContext.class))
207 .findFirst()
208 .orElseThrow();
209 }
210
211 private static class Dispatcher {
212 private final Runnable kernelType;
213
214 private Dispatcher(Runnable kernelType) {
215 this.kernelType = kernelType;
216 }
217
218 private void dispatch(Map<Op.Location, Map<SpecializationKey, KernelCallSite>> kernelCallSiteCache, MethodHandles.Lookup lookup, ComputeCallGraph computeCallGraph, Accelerator accelerator, NDRange ndRange) {
219 Quoted<JavaOp.LambdaOp> quoted = Op.ofLambda(kernelType).orElseThrow();
220 JavaOp.LambdaOp lambdaOp = quoted.op();
221 var location = quoted.op().location();
222
223 OpHelper.Invoke targetInvoke = getTargetInvoke(lookup, lambdaOp);
224 MethodRef methodRef = targetInvoke.op().invokeReference();
225 Object[] quotedCapturedValues = targetInvoke.getCapturedValuesForOperands(quoted.capturedValues());
226 KernelCallSite kernelCallSite;
227
228 try {
229 SpecializationKey key = SpecializationKey.of(methodRef.resolveToMethod(lookup), quotedCapturedValues);
230 var m = methodRef.resolveToMethod(lookup);
231 if (kernelCallSiteCache.containsKey(location) && kernelCallSiteCache.get(location).containsKey(key)) {
232 var oldKernelCallSite = kernelCallSiteCache.get(location).get(key);
233 kernelCallSite = new KernelCallSite(quoted, oldKernelCallSite.lambdaOp(), oldKernelCallSite.methodRef(), oldKernelCallSite.kernelCallGraph(), quotedCapturedValues, oldKernelCallSite.dispatchContext);
234 } else {
235 kernelCallSite = kernelCallSiteCache.computeIfAbsent(location, k -> new ConcurrentHashMap<>())
236 .computeIfAbsent(key, _ -> {
237 KernelCallGraph kernelCallGraph = computeCallGraph.kernelCallGraphMap.get(methodRef);
238 if (kernelCallGraph == null) {
239 throw new IllegalStateException("Failed to create KernelCallGraph (did you miss @Reflect annotation?).");
240 }
241 // Create a new KernelCallGraph starting from the original method
242 KernelCallGraph kcg = new KernelCallGraph(kernelCallGraph.computeCallGraph, m, kernelCallGraph.getOriginalKernelFunction());
243
244 // Compilation happens here!
245 kcg.compile(quotedCapturedValues);
246 DispatchContext dispatchContext;
247 if (kernelType instanceof Kernel) {
248 dispatchContext = DispatchContext.createDefaultContext(kernelCallGraph.computeCallGraph.computeContext.accelerator());
249 } else if (kernelType instanceof TileKernel) {
250 dispatchContext = DispatchContext.createTileContext(kernelCallGraph.computeCallGraph.computeContext.accelerator());
251 } else {
252 throw new IllegalStateException("Unknown KernelType " + kernelType.getClass());
253 }
254
255 return new KernelCallSite(quoted, lambdaOp, methodRef, kcg, quotedCapturedValues, dispatchContext);
256 });
257 }
258 } catch (ReflectiveOperationException e) {
259 throw new RuntimeException(e);
260 }
261
262 Object[] dispatchContextAndArgs = new Object[kernelCallSite.capturedArgs.length + 1];
263 dispatchContextAndArgs[0] = kernelCallSite.dispatchContext;
264 System.arraycopy(kernelCallSite.capturedArgs(), 0, dispatchContextAndArgs, 1, kernelCallSite.capturedArgs().length);
265 if (kernelType instanceof Kernel) {
266 accelerator.backend.dispatchKernel(kernelCallSite.kernelCallGraph(), ndRange, dispatchContextAndArgs);
267 } else if (kernelType instanceof TileKernel) {
268 accelerator.backend.dispatchTile(kernelCallSite.kernelCallGraph(), ndRange, dispatchContextAndArgs);
269 } else {
270 throw new IllegalStateException("Unknown KernelType: " + kernelType.getClass());
271 }
272 }
273 }
274
275
276 /** Creating the kernel callsite involves
277 walking the code model of the lambda
278 analyzing the callgraph and transforming to HATDialect
279 So we cache the callsite against the location from the lambdaop.
280 */
281 public void dispatchKernel(NDRange ndRange, Kernel kernel) {
282 Dispatcher dispatcher = new Dispatcher(kernel);
283 dispatcher.dispatch(kernelCallSiteCache, lookup(), computeCallGraph, accelerator, ndRange);
284 }
285
286 /**
287 * Function to dispatch a TileKernel in HAT. The dispatch takes the following parameters:
288 *
289 * @param ndRange A Tile Range that specified the total number of tiles and the tile-size
290 * @param tileKernel The tile kernel of offload and run on the hardware accelerator
291 */
292 public void dispatchTile(NDRange ndRange, TileKernel tileKernel) {
293 Dispatcher dispatcher = new Dispatcher(tileKernel);
294 dispatcher.dispatch(kernelCallSiteCache, lookup(), computeCallGraph, accelerator, ndRange);
295 }
296
297 @Override
298 public void preMutate(MappableIface b) {
299 if (accelerator.backend instanceof BufferTracker bufferTracker) {
300 bufferTracker.preMutate(b);
301 }
302 }
303
304 @Override
305 public void postMutate(MappableIface b) {
306 if (accelerator.backend instanceof BufferTracker bufferTracker) {
307 bufferTracker.postMutate(b);
308 }
309
310 }
311
312 @Override
313 public void preAccess(MappableIface b) {
314 if (accelerator.backend instanceof BufferTracker bufferTracker) {
315 bufferTracker.preAccess(b);
316 }
317
318 }
319
320 @Override
321 public void postAccess(MappableIface b) {
322 if (accelerator.backend instanceof BufferTracker bufferTracker) {
323 bufferTracker.postAccess(b);
324 }
325 }
326
327 @FunctionalInterface
328 public interface Kernel extends Runnable { }
329
330 @FunctionalInterface
331 public interface TileKernel extends Runnable { }
332
333 }