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
26 package hat.backend.ffi;
27
28 import hat.ComputeContext;
29 import hat.Config;
30 import hat.NDRange;
31 import hat.annotations.Kernel;
32 import hat.annotations.Preformatted;
33 import hat.annotations.TypeDef;
34 import hat.buffer.ArgArray;
35 import hat.buffer.DispatchContext;
36 import hat.callgraph.IfaceDataDag;
37 import hat.callgraph.KernelCallGraph;
38 import hat.callgraph.MethodCallDag;
39 import hat.codebuilders.C99HATKernelBuilder;
40 import hat.codebuilders.C99VecAndMatHandler;
41 import hat.device.DeviceSchema;
42 import hat.device.NonMappableIface;
43 import hat.types.BF16;
44 import hat.types.F16;
45 import jdk.incubator.code.CodeTransformer;
46 import optkl.ifacemapper.BoundSchema;
47 import optkl.ifacemapper.Buffer;
48 import optkl.ifacemapper.BufferState;
49 import optkl.ifacemapper.BufferTracker;
50 import optkl.ifacemapper.MappableIface;
51 import optkl.ifacemapper.Schema;
52
53 import java.lang.foreign.Arena;
54 import java.lang.invoke.MethodHandles;
55 import java.util.Arrays;
56 import java.util.HashMap;
57 import java.util.HashSet;
58 import java.util.List;
59 import java.util.Map;
60 import java.util.Set;
61
62 public abstract class C99FFIBackend extends FFIBackendDriver implements BufferTracker {
63
64 protected C99FFIBackend(Arena arena, MethodHandles.Lookup lookup, String libName, Config config) {
65 super(arena, lookup, libName, config);
66 }
67
68 public static class CompiledKernel {
69 public final C99FFIBackend c99FFIBackend;
70 public final KernelCallGraph kernelCallGraph;
71 public final BackendBridge.CompilationUnitBridge.KernelBridge kernelBridge;
72 public final ArgArray argArray;
73 public final DispatchContext dispatchContext;
74
75 public CompiledKernel(C99FFIBackend c99FFIBackend, KernelCallGraph kernelCallGraph, BackendBridge.CompilationUnitBridge.KernelBridge kernelBridge, Object[] dispatchContextAndArgs) {
76 this.c99FFIBackend = c99FFIBackend;
77 this.kernelCallGraph = kernelCallGraph;
78 this.kernelBridge = kernelBridge;
79 this.dispatchContext = (DispatchContext) dispatchContextAndArgs[0];
80 this.argArray = ArgArray.create(kernelCallGraph.computeCallGraph.computeContext.accelerator(), kernelCallGraph, dispatchContextAndArgs);
81 }
82
83 public void dispatch( NDRange ndRange, Object[] dispatchContextAndArgs) {
84 dispatchContext.gsy(1);
85 dispatchContext.gsz(1);
86 switch (ndRange.global()) {
87 case NDRange.Global1D global1D -> {
88 dispatchContext.gsx(global1D.x());
89 dispatchContext.dimensions(global1D.dimension());
90 }
91 case NDRange.Global2D global2D -> {
92 dispatchContext.gsx(global2D.x());
93 dispatchContext.gsy(global2D.y());
94 dispatchContext.dimensions(global2D.dimension());
95 }
96 case NDRange.Global3D global3D -> {
97 dispatchContext.gsx(global3D.x());
98 dispatchContext.gsy(global3D.y());
99 dispatchContext.gsz(global3D.z());
100 dispatchContext.dimensions(global3D.dimension());
101 }
102 case null, default ->
103 throw new IllegalArgumentException("Unknown global range " + ndRange.global().getClass());
104 }
105
106 if (ndRange.hasLocal()) {
107 dispatchContext.lsy(1);
108 dispatchContext.lsz(1);
109 switch (ndRange.local()) {
110 case NDRange.Local1D local1D -> {
111 dispatchContext.lsx(local1D.x());
112 dispatchContext.dimensions(local1D.dimension());
113 }
114 case NDRange.Local2D local2D -> {
115 dispatchContext.lsx(local2D.x());
116 dispatchContext.lsy(local2D.y());
117 dispatchContext.dimensions(local2D.dimension());
118 }
119 case NDRange.Local3D local3D -> {
120 dispatchContext.lsx(local3D.x());
121 dispatchContext.lsy(local3D.y());
122 dispatchContext.lsz(local3D.z());
123 dispatchContext.dimensions(local3D.dimension());
124 }
125 case null, default ->
126 throw new IllegalArgumentException("Unknown global range " + ndRange.local().getClass());
127 }
128 } else {
129 dispatchContext.lsx(0);
130 dispatchContext.lsy(0);
131 dispatchContext.lsz(0);
132 }
133
134 // Set Tile
135
136 dispatchContext.tlx(0);
137 dispatchContext.tly(0);
138 dispatchContext.tlz(0);
139 if (ndRange.hasTile()) {
140 switch (ndRange.tile()) {
141 case NDRange.Tile1D tile1D -> {
142 dispatchContext.tlx(tile1D.x());
143 }
144 case NDRange.Tile2D tile2D -> {
145 dispatchContext.tlx(tile2D.x());
146 dispatchContext.tly(tile2D.y());
147 }
148 case NDRange.Tile3D tile3D -> {
149 dispatchContext.tlx(tile3D.x());
150 dispatchContext.tly(tile3D.y());
151 dispatchContext.tlz(tile3D.z());
152 }
153 case null, default ->
154 throw new IllegalArgumentException("Unknown global range " + ndRange.tile().getClass());
155 }
156 }
157
158 // Set warp
159 dispatchContext.wsx(0);
160 dispatchContext.wsy(0);
161 dispatchContext.wsz(0);
162 if (ndRange.hasWarp()) {
163 switch (ndRange.warp()) {
164 case NDRange.Warp1D warp1D -> {
165 dispatchContext.wsx(warp1D.x() ? 1 : 0);
166 }
167 case NDRange.Warp2D warp2D -> {
168 dispatchContext.wsx(warp2D.x() ? 1 : 0);
169 dispatchContext.wsy(warp2D.y() ? 1 : 0);
170 }
171 case NDRange.Warp3D warp3D -> {
172 dispatchContext.wsx(warp3D.x() ? 1 : 0);
173 dispatchContext.wsy(warp3D.y() ? 1 : 0);
174 dispatchContext.wsz(warp3D.z() ? 1 : 0);
175 }
176 case null, default ->
177 throw new IllegalArgumentException("Unknown global range " + ndRange.warp().getClass());
178 }
179 }
180 // Set the tile Model
181 dispatchContext.type(dispatchContext.type());
182
183 dispatchContextAndArgs[0] = this.dispatchContext;
184 ArgArray.update(argArray, kernelCallGraph, dispatchContextAndArgs);
185 kernelBridge.ndRange(this.argArray);
186 }
187 }
188
189
190 public Map<KernelCallGraph, CompiledKernel> kernelCallGraphCompiledCodeMap = new HashMap<>();
191
192 public abstract String createCode(KernelCallGraph kernelCallGraph, Object... justArgs);
193
194 @Override
195 public void computeContextHandoff(ComputeContext computeContext) {
196 computeContext.computeCallGraph().callDag.entryPoint.funcOp(injectBufferTracking(config(), computeContext.lookup(), computeContext.computeCallGraph().callDag.entryPoint.funcOp()));
197 }
198
199 @Override
200 public void dispatchCompute(ComputeContext computeContext, Object... args) {
201 backendBridge.computeStart();
202 computeContext.invokeWithArgs(args);
203 backendBridge.computeEnd();
204 }
205
206 @Override
207 public final void dispatchKernel(KernelCallGraph kernelCallGraph, NDRange ndRange, Object... dispatchContextAndArgs) {
208 CompiledKernel compiledKernel = kernelCallGraphCompiledCodeMap.computeIfAbsent(kernelCallGraph, (_) -> {
209 String code = createCode(kernelCallGraph,Arrays.copyOfRange(dispatchContextAndArgs,1,dispatchContextAndArgs.length));
210 if (config().showCode()) {
211 IO.println(code);
212 }
213 var compilationUnit = backendBridge.compile(code, 0);
214 if (compilationUnit.ok()) {
215 var kernel = compilationUnit.getKernel(kernelCallGraph.callDag.entryPoint.method().getName());
216 return new CompiledKernel(this, kernelCallGraph, kernel, dispatchContextAndArgs);
217 } else {
218 throw new IllegalStateException("backend failed to compile ");
219 }
220 });
221 compiledKernel.dispatch(ndRange, dispatchContextAndArgs);
222 }
223
224 public <T extends C99HATKernelBuilder<T>> String createCode(KernelCallGraph kernelCallGraph, T builder, Object... args) {
225 builder.defines().types();
226
227 var preformattedAnnotation = kernelCallGraph.callDag.entryPoint.method().getAnnotation(Preformatted.class);
228 var visitedAlready = new HashSet<Schema.IfaceType>();
229 Arrays.stream(args)
230 .filter(arg -> arg instanceof Buffer)
231 .map(arg -> (Buffer) arg)
232 .forEach(ifaceBuffer -> {
233 BoundSchema<?> boundSchema = MappableIface.getBoundSchema(ifaceBuffer);
234 boundSchema.schema().rootIfaceType.visitUniqueTypes(t -> {
235 if (visitedAlready.add(t)) { // true first time we see this type
236 if (preformattedAnnotation == null) {
237 builder.typedef(boundSchema, t);
238 }
239 }
240 });
241 });
242
243
244 var kernelAnnotation = kernelCallGraph.callDag.entryPoint.method().getAnnotation(Kernel.class);
245 if (kernelAnnotation != null) {
246 // If we find a kernelAnnotation we can't trust the data in kernelCallGraph's state.
247 kernelCallGraph.setUsesAtomics(true);
248 kernelCallGraph.accessedFP16Classes.addAll(List.of(F16.class, BF16.class));
249 kernelCallGraph.setUsesBarrier(true);
250
251 var typedefAnnotation = kernelCallGraph.callDag.entryPoint.method().getAnnotation(TypeDef.class);
252 if (typedefAnnotation != null) {
253 builder.lineComment("Preformatted typedef body from @Typedef annotation");
254 builder.typedefStruct(typedefAnnotation.name(), _ -> builder.preformatted(typedefAnnotation.body())).semicolon().nl();
255 }
256 if (preformattedAnnotation != null) {
257 builder.lineComment("Preformatted text from @Preformatted annotation");
258 builder.preformatted(preformattedAnnotation.value());
259 }
260 builder.lineComment("Preformatted code body from @Kernel annotation");
261 builder.preformatted(kernelAnnotation.value());
262 } else {
263 Set<Class<?>> typeDefined = new HashSet<>();
264 typeDefined.add(F16.class);
265 typeDefined.add(BF16.class);
266 kernelCallGraph.accessedNonMappableIfaceClasses.stream()
267 .filter(c->!typeDefined.contains(c))
268 .map(c->(Class<NonMappableIface>) c) // why do we need to do this.
269 .forEach(c -> {
270 // We create a dag of iface references rooted at c
271 var ifaceDataDag = new IfaceDataDag<NonMappableIface>(dag -> {
272 var entryPoint = dag.getNode(c);
273 dag.methodsWithIfaceReturnTypes(c).forEach(ifaceInfo ->
274 dag.addEdge(entryPoint, dag.getNode(ifaceInfo.clazz())) // this recurses with each added class
275 );
276 });
277 // Now we can generate typedefs in rankOrder (so inner typedefs first)
278 if (ifaceDataDag.isDag()) {
279 ifaceDataDag.rankOrdered.stream()
280 .filter(ifaceInfo -> !typeDefined.contains(ifaceInfo.clazz()))
281 .forEach(ifaceInfo -> typeDefined.add(
282 DeviceSchema.getDeviceSchemaOrThrow(ifaceInfo.clazz()).typedef(builder).clazz()
283 )
284 );
285 } else {
286 typeDefined.add(DeviceSchema.getDeviceSchemaOrThrow(c).typedef(builder).clazz());
287 }
288 });
289
290 // This is a slight hack for Shader support.
291 if (!kernelCallGraph.accessedVecClasses.isEmpty()) {
292 C99VecAndMatHandler.createVecFunctions(builder);
293 }
294
295 kernelCallGraph.callDag.rankOrdered.stream()
296 .filter(m -> m instanceof MethodCallDag.OtherMethodCall)
297 .forEach(m -> builder.nl().kernelMethod( m.funcOp()).nl());
298
299 builder.nl().kernelEntrypoint().nl();
300
301 if (config().showKernelModel()) {
302 IO.println("Non Lowered");
303 IO.println(kernelCallGraph.callDag.entryPoint.funcOp().toText());
304 }
305 if (config().showLoweredKernelModel()) {
306 IO.println("Lowered");
307 IO.println(kernelCallGraph.callDag.entryPoint.funcOp().transform(CodeTransformer.LOWERING_TRANSFORMER).toText());
308 }
309 }
310 return builder.toString();
311 }
312
313 @Override
314 public void preMutate(MappableIface b) {
315 switch (b.getState()) {
316 case BufferState.NO_STATE:
317 case BufferState.NEW_STATE:
318 case BufferState.HOST_OWNED:
319 case BufferState.DEVICE_VALID_HOST_HAS_COPY: {
320 if (config().showState()) {
321 System.out.println("in preMutate state = " + b.getStateString() + " no action to take");
322 }
323 break;
324 }
325 case BufferState.DEVICE_OWNED: {
326 backendBridge.getBufferFromDeviceIfDirty(b);// calls through FFI and might block when fetching from device
327 if (config().showState()) {
328 System.out.print("in preMutate state = " + b.getStateString() + " we pulled from device ");
329 }
330 b.setState(BufferState.DEVICE_VALID_HOST_HAS_COPY);
331 if (config().showState()) {
332 System.out.println("and switched to " + b.getStateString());
333 }
334 break;
335 }
336 default:
337 throw new IllegalStateException("Not expecting this state ");
338 }
339 }
340
341 @Override
342 public void postMutate(MappableIface b) {
343 if (config().showState()) {
344 System.out.print("in postMutate state = " + b.getStateString() + " no action to take ");
345 }
346 if (b.getState() != BufferState.NEW_STATE) {
347 b.setState(BufferState.HOST_OWNED);
348 }
349 if (config().showState()) {
350 System.out.println("and switched to (or stayed on) " + b.getStateString());
351 }
352 }
353
354 @Override
355 public void preAccess(MappableIface b) {
356 switch (b.getState()) {
357 case BufferState.NO_STATE:
358 case BufferState.NEW_STATE:
359 case BufferState.HOST_OWNED:
360 case BufferState.DEVICE_VALID_HOST_HAS_COPY: {
361 if (config().showState()) {
362 System.out.println("in preAccess state = " + b.getStateString() + " no action to take");
363 }
364 break;
365 }
366 case BufferState.DEVICE_OWNED: {
367 backendBridge.getBufferFromDeviceIfDirty(b);// calls through FFI and might block when fetching from device
368
369 if (config().showState()) {
370 System.out.print("in preAccess state = " + b.getStateString() + " we pulled from device ");
371 }
372 b.setState(BufferState.DEVICE_VALID_HOST_HAS_COPY);
373 if (config().showState()) {
374 System.out.println("and switched to " + b.getStateString());
375 }
376 break;
377 }
378 default:
379 throw new IllegalStateException("Not expecting this state ");
380 }
381 }
382
383
384 @Override
385 public void postAccess(MappableIface b) {
386 if (config().showState()) {
387 System.out.println("in postAccess state = " + b.getStateString());
388 }
389 }
390 }