Opening Hat Dialect for Handling Barriers and Shared/Private Memory (#560)

Author: jjfumero

hat/backends/ffi/cuda/src/main/java/hat/backend/ffi/PTXHATKernelBuilder.java

@@ -24,6 +24,7 @@
  */
 package hat.backend.ffi;
 
+import hat.dialect.HatOP;
 import hat.ifacemapper.BoundSchema;
 import hat.optools.*;
 import hat.codebuilders.CodeBuilder;
@@ -163,7 +164,7 @@ public void functionEpilogue() {
         cbrace();
     }
 
-    public static class PTXPtrOp extends Op {
+    public static class PTXPtrOp extends HatOP {
         public String fieldName;
         public static final String NAME = "ptxPtr";
         final TypeElement resultType;

hat/backends/ffi/cuda/src/main/native/cpp/cuda_backend_queue.cpp

@@ -178,13 +178,25 @@ void CudaBackend::CudaQueue::dispatch(KernelContext *kernelContext, CompilationU
         std::cout << "dispatch()  thread=" <<thread_id<< " != "<< streamCreationThread<< std::endl;
     }
 
+    // // CUDA events for timing
+    // cudaEvent_t start, stop;
+    // cuEventCreate(&start, cudaEventDefault);
+    // cuEventCreate(&stop, cudaEventDefault);
+    // cuEventRecord(start, 0);
+
     const auto status = cuLaunchKernel(cudaKernel->function, //
                                  blocksPerGridX, blocksPerGridY, blocksPerGridZ, //
                                  threadsPerBlockX, threadsPerBlockY, threadsPerBlockZ, //
                                  0, //
                                  cuStream, //
                                  cudaKernel->argslist, //
                                  nullptr);
+    // cuEventRecord(stop, 0);
+    // cuEventSynchronize(stop); // Wait for completion
+    //
+    // float elapsedTimeMs = 0.0f;
+    // cuEventElapsedTime(&elapsedTimeMs, start, stop);
+    // std::cout << "Kernel Elapsed Time: " << elapsedTimeMs << " ms\n";
 
     CUDA_CHECK(status, "cuLaunchKernel");
 }

hat/backends/ffi/shared/src/main/java/hat/backend/ffi/C99FFIBackend.java

@@ -36,15 +36,13 @@
 import hat.buffer.KernelContext;
 import hat.callgraph.KernelCallGraph;
 import hat.codebuilders.ScopedCodeBuilderContext;
+import hat.dialect.HatMemoryOp;
 import hat.ifacemapper.BoundSchema;
 import hat.ifacemapper.BufferState;
 import hat.ifacemapper.Schema;
 import hat.optools.OpTk;
 import jdk.incubator.code.CopyContext;
 import jdk.incubator.code.Op;
-import jdk.incubator.code.dialect.java.ClassType;
-import jdk.incubator.code.dialect.java.JavaOp;
-import jdk.incubator.code.dialect.java.JavaType;
 
 import java.lang.invoke.MethodHandles;
 import java.lang.reflect.InvocationTargetException;
@@ -142,16 +140,15 @@ public void dispatch(NDRange ndRange, Object[] args) {
 
     public Map<KernelCallGraph, CompiledKernel> kernelCallGraphCompiledCodeMap = new HashMap<>();
 
-    private void updateListOfSchemas(Op op, MethodHandles.Lookup lookup, List<String> localIfaceList) {
-        if (Objects.requireNonNull(op) instanceof JavaOp.InvokeOp invokeOp) {
-            if (OpTk.isIfaceAccessor(lookup, invokeOp)) {
-                String klassName = invokeOp.resultType().toString();
-                localIfaceList.add(klassName);
-            }
+    private void updateListOfSchemas(Op op, List<String> localIfaceList) {
+        if (Objects.requireNonNull(op) instanceof HatMemoryOp hatMemoryOp) {
+            String klassName = hatMemoryOp.invokeType().toString();
+            localIfaceList.add(klassName);
         }
     }
 
     public <T extends C99HATKernelBuilder<T>> String createCode(KernelCallGraph kernelCallGraph, T builder, Object... args) {
+
         builder.defines().pragmas().types();
         Set<Schema.IfaceType> already = new LinkedHashSet<>();
         Arrays.stream(args)
@@ -171,9 +168,9 @@ public <T extends C99HATKernelBuilder<T>> String createCode(KernelCallGraph kern
         // Traverse the list of reachable functions and append the intrinsics functions found for each of the functions
         if (kernelCallGraph.moduleOp != null) {
             kernelCallGraph.moduleOp.functionTable()
-                    .forEach((_, funcOp) -> {
-                        funcOp.transform(CopyContext.create(), (blockBuilder, op) -> {
-                            updateListOfSchemas(op, kernelCallGraph.computeContext.accelerator.lookup, localIFaceList);
+                    .forEach((entryName, f) -> {
+                        f.transform(CopyContext.create(), (blockBuilder, op) -> {
+                            updateListOfSchemas(op, localIFaceList);
                             blockBuilder.op(op);
                             return blockBuilder;
                         });
@@ -183,7 +180,7 @@ public <T extends C99HATKernelBuilder<T>> String createCode(KernelCallGraph kern
             // this else-branch will be deleted.
             kernelCallGraph.kernelReachableResolvedStream().forEach((kernel) -> {
                 kernel.funcOp().transform(CopyContext.create(), (blockBuilder, op) -> {
-                    updateListOfSchemas(op, kernelCallGraph.computeContext.accelerator.lookup, localIFaceList);
+                    updateListOfSchemas(op, localIFaceList);
                     blockBuilder.op(op);
                     return blockBuilder;
                 });
@@ -193,7 +190,7 @@ public <T extends C99HATKernelBuilder<T>> String createCode(KernelCallGraph kern
         // Traverse the main kernel and append the intrinsics functions found in the main kernel
         kernelCallGraph.entrypoint.funcOp()
                 .transform(CopyContext.create(), (blockBuilder, op) -> {
-                    updateListOfSchemas(op, kernelCallGraph.computeContext.accelerator.lookup, localIFaceList);
+                    updateListOfSchemas(op, localIFaceList);
                     blockBuilder.op(op);
                     return blockBuilder;
                 });

hat/core/src/main/java/hat/ComputeContext.java

@@ -29,17 +29,22 @@
 import hat.buffer.BufferTracker;
 import hat.callgraph.ComputeCallGraph;
 import hat.callgraph.KernelCallGraph;
+import hat.dialect.HatBarrierOp;
 import hat.ifacemapper.BoundSchema;
 import hat.ifacemapper.SegmentMapper;
 import hat.optools.OpTk;
+import jdk.incubator.code.Block;
+import jdk.incubator.code.CopyContext;
 import jdk.incubator.code.Op;
 import jdk.incubator.code.Quotable;
 import jdk.incubator.code.Quoted;
+import jdk.incubator.code.Value;
 import jdk.incubator.code.dialect.core.CoreOp;
 import jdk.incubator.code.dialect.java.JavaOp;
 import jdk.incubator.code.dialect.java.MethodRef;
 
 import java.lang.reflect.Method;
+import java.util.List;
 import java.util.function.Consumer;
 
 /**
@@ -132,13 +137,39 @@ public void dispatchKernel(ComputeRange computeRange, QuotableKernelContextConsu
         dispatchKernelWithComputeRange(computeRange, quotableKernelContextConsumer);
     }
 
+    private boolean isMethodFromHatKernelContext(JavaOp.InvokeOp invokeOp) {
+        String kernelContextCanonicalName = hat.KernelContext.class.getName();
+        return invokeOp.invokeDescriptor().refType().toString().equals(kernelContextCanonicalName);
+    }
+
+    private boolean isMethod(JavaOp.InvokeOp invokeOp, String methodName) {
+        return invokeOp.invokeDescriptor().name().equals(methodName);
+    }
+
+    private void createBarrierNodeOp(CopyContext context, JavaOp.InvokeOp invokeOp, Block.Builder blockBuilder) {
+        List<Value> inputOperands = invokeOp.operands();
+        List<Value> outputOperands = context.getValues(inputOperands);
+        HatBarrierOp hatBarrierOp = new HatBarrierOp(outputOperands);
+        Op.Result outputResult = blockBuilder.op(hatBarrierOp);
+        Op.Result inputResult = invokeOp.result();
+        context.mapValue(inputResult, outputResult);
+    }
+
     record CallGraph(Quoted quoted, JavaOp.LambdaOp lambdaOp, MethodRef methodRef, KernelCallGraph kernelCallGraph) {}
 
     private CallGraph buildKernelCallGraph(QuotableKernelContextConsumer quotableKernelContextConsumer) {
         Quoted quoted = Op.ofQuotable(quotableKernelContextConsumer).orElseThrow();
         JavaOp.LambdaOp lambdaOp = (JavaOp.LambdaOp) quoted.op();
-        MethodRef methodRef =OpTk.getQuotableTargetInvokeOpWrapper( lambdaOp).invokeDescriptor();
+        MethodRef methodRef = OpTk.getQuotableTargetInvokeOpWrapper( lambdaOp).invokeDescriptor();
         KernelCallGraph kernelCallGraph = computeCallGraph.kernelCallGraphMap.get(methodRef);
+        // Analysis : dialect
+        // NOTE: Keep the following boolean until we have the config available/reachable
+        // from this class
+        boolean useDialect = true;
+        if (useDialect) {
+            //System.out.println("[INFO] Using Hat Dialect?: " + useDialect);
+            kernelCallGraph.dialectifyToHat();
+        }
         return new CallGraph(quoted, lambdaOp, methodRef, kernelCallGraph);
     }
 
@@ -157,6 +188,7 @@ private void dispatchKernel(int rangeX, int rangeY, int rangeZ, int dimNumber, Q
             accelerator.backend.dispatchKernel(cg.kernelCallGraph, ndRange, args);
         } catch (Throwable t) {
             System.out.print("what?" + cg.methodRef + " " + t);
+            t.printStackTrace();
             throw t;
         }
     }
@@ -170,6 +202,7 @@ private void dispatchKernelWithComputeRange(ComputeRange computeRange, QuotableK
             accelerator.backend.dispatchKernel(cg.kernelCallGraph, ndRange, args);
         } catch (Throwable t) {
             System.out.print("what?" + cg.methodRef + " " + t);
+            t.printStackTrace();
             throw t;
         }
     }

hat/core/src/main/java/hat/buffer/Buffer.java

@@ -25,7 +25,6 @@
 package hat.buffer;
 
 
-import hat.Space;
 import hat.ifacemapper.BoundSchema;
 import hat.ifacemapper.BufferState;
 import hat.ifacemapper.MappableIface;

hat/core/src/main/java/hat/callgraph/ComputeCallGraph.java

@@ -229,9 +229,11 @@ public void closeWithModuleOp(ComputeReachableResolvedMethodCall computeReachabl
 
     @Override
     public boolean filterCalls(CoreOp.FuncOp f, JavaOp.InvokeOp invokeOp, Method method, MethodRef methodRef, Class<?> javaRefTypeClass) {
-        if (entrypoint.method.getDeclaringClass().equals(OpTk.javaRefClassOrThrow(computeContext.accelerator.lookup,invokeOp)) && isKernelDispatch(computeContext.accelerator.lookup,method, f)) {
+        if (entrypoint.method.getDeclaringClass().equals(OpTk.javaRefClassOrThrow(computeContext.accelerator.lookup,invokeOp))
+                && isKernelDispatch(computeContext.accelerator.lookup,method, f)) {
             kernelCallGraphMap.computeIfAbsent(methodRef, _ ->
-                    new KernelCallGraph(this, methodRef, method, f).closeWithModuleOp()
+                    new KernelCallGraph(this, methodRef, method, f)
+                            .closeWithModuleOp()
             );
         } else if (ComputeContext.class.isAssignableFrom(javaRefTypeClass)) {
             computeContextMethodCall = new ComputeContextMethodCall(this, methodRef, method);