[mlir] Add gpu async integration test.

[aartbik]

Just curious, are these tests candidates for mlir/integration_test rather than mlir/test?

[csigg]

Rebase.

[csigg]

Trying to upload again with D94399 as base.

[csigg]

In D94421#2490657, @aartbik wrote:

Just curious, are these tests candidates for mlir/integration_test rather than mlir/test?

The cuda-runner and rocm-runner tests have traditionally been in the test directory.
But yeah, they might fit better in the integration_test directory.
I would move them in a separate change though.
That said, is there any specific reason we have two test directories?

[aartbik]

In D94421#2499080, @csigg wrote:

I would move them in a separate change though.
That said, is there any specific reason we have two test directories?

Of course, I was not suggesting to make that part of this change.

The test directory is for fast running, light weight unit tests. The integration test directory is for (potentially) slightly longer running, heavier weight tests. Invoking full runners (like we do for mlir-cpu-runner) always seem good candidates. The latter is run less often to save resources.

It is of course, up to the test "owners" ultimately to make the call

[herhut]

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

300

Long term we have to find a better way to do these things. For now this is fine (just writing this here for documentation).

mlir/test/mlir-cuda-runner/async.mlir

15

Is this needed?

60

Why not explicitly copy back, too?

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

115

Just wondering whether all of this could also become a textual pass specification. Then the cuda runner would essentially only contain of a special pass that wraps the creation of cubin blobs (the createConvertGPUKernelToBlobPass) and otherwise uses a specified pipeline. Then we would not need to bake in the async lowering here.

[csigg]

mlir/test/mlir-cuda-runner/async.mlir

15

It is not required for correctness if we copy the result as you suggested below.

But without page-locking, the memcpy becomes synchronous and the second and third memcpy would schedule (and execute) after the first has completed even without synchronization.

60

I had that initially, but I wanted to keep the number of GPU operations minimal.

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

115

Good idea, I will try that.

[csigg]

Rebase.

[herhut]

Thanks for the explanations!

[csigg]

Rebase.

mlir/lib/Conversion/GPUCommon/CMakeLists.txt

add_mlir_conversion_library(MLIRGPUToGPURuntimeTransforms
 
   LINK_COMPONENTS
   Core
   MC
   ${AMDGPU_LIBS}
   ${NVPTX_LIBS}
 
   LINK_LIBS PUBLIC
+  MLIRAsyncToLLVM
   MLIRGPU
   MLIRIR
   MLIRLLVMIR
   MLIRPass
   MLIRSupport
   MLIRStandardToLLVM
 )

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

 // ABI, this pass uses a slim runtime layer that builds on top of the public
 // API from GPU runtime headers.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 
 #include "../PassDetail.h"
+#include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 
private:
   matchAndRewrite(gpu::MemcpyOp memcpyOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 void GpuToLLVMConversionPass::runOnOperation() {
   LLVMTypeConverter converter(&getContext());
   OwningRewritePatternList patterns;
+  LLVMConversionTarget target(getContext());
+
   populateStdToLLVMConversionPatterns(converter, patterns);
+  populateAsyncStructuralTypeConversionsAndLegality(&getContext(), converter,
+                                                    patterns, target);
   populateGpuToLLVMConversionPatterns(converter, patterns, gpuBinaryAnnotation);
 
-  LLVMConversionTarget target(getContext());
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
 }
 
 LLVM::CallOp FunctionCallBuilder::create(Location loc, OpBuilder &builder,
                                          ArrayRef<Value> arguments) const {
   auto module = builder.getBlock()->getParent()->getParentOfType<ModuleOp>();

mlir/test/mlir-cuda-runner/async.mlir

+// RUN: mlir-cuda-runner %s --entry-point-result=void -O0 \
+// RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
+// RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_async_runtime%shlibext \
+// RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
+// RUN: | FileCheck %s
+
+func @main() {
+  %c0    = constant 0 : index
+  %c1    = constant 1 : index
+  %count = constant 2 : index
+
+  // initialize h0 on host
+  %h0 = alloc(%count) : memref<?xi32>
+  %h0_unranked = memref_cast %h0 : memref<?xi32> to memref<*xi32>
+  gpu.host_register %h0_unranked : memref<*xi32>
+
+  %v0 = constant 42 : i32
+  store %v0, %h0[%c0] : memref<?xi32>
+  store %v0, %h0[%c1] : memref<?xi32>
+
+  // copy h0 to b0 on device.
+  %t0, %f0 = async.execute () -> !async.value<memref<?xi32>> {
+    %b0 = gpu.alloc(%count) : memref<?xi32>
+    gpu.memcpy %b0, %h0 : memref<?xi32>, memref<?xi32>
+    async.yield %b0 : memref<?xi32>
+  }
+
+  // copy h0 to b1 and b2 (fork)
+  %t1, %f1 = async.execute [%t0] (
+    %f0 as %b0 : !async.value<memref<?xi32>>
+  ) -> !async.value<memref<?xi32>> {
+    %b1 = gpu.alloc(%count) : memref<?xi32>
+    gpu.memcpy %b1, %b0 : memref<?xi32>, memref<?xi32>
+    async.yield %b1 : memref<?xi32>
+  }
+  %t2, %f2 = async.execute [%t0] (
+    %f0 as %b0 : !async.value<memref<?xi32>>
+  ) -> !async.value<memref<?xi32>> {
+    %b2 = gpu.alloc(%count) : memref<?xi32>
+    gpu.memcpy %b2, %b0 : memref<?xi32>, memref<?xi32>
+    async.yield %b2 : memref<?xi32>
+  }
+
+  // h0 = b1 + b2 (join).
+  %t3 = async.execute [%t1, %t2] (
+    %f1 as %b1 : !async.value<memref<?xi32>>,
+    %f2 as %b2 : !async.value<memref<?xi32>>
+  ) {
+    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
+               threads(%tx, %ty, %tz) in (%block_x = %count, %block_y = %c1, %block_z = %c1) {
+      %v1 = load %b1[%tx] : memref<?xi32>
+      %v2 = load %b2[%tx] : memref<?xi32>
+      %sum = addi %v1, %v2 : i32
+      store %sum, %h0[%tx] : memref<?xi32>
+      gpu.terminator
+    }
+    async.yield
+  }
+
+  async.await %t3 : !async.token
+  // CHECK: [84, 84]
+  call @print_memref_i32(%h0_unranked) : (memref<*xi32>) -> ()
+  return
+}
+
+func private @print_memref_i32(memref<*xi32>)

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

 //===- mlir-cuda-runner.cpp - MLIR CUDA Execution Driver-------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is a command line utility that executes an MLIR file on the GPU by
 // translating MLIR to NVVM/LVVM IR before JIT-compiling and executing the
 // latter.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/STLExtras.h"
 
+#include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/ExecutionEngine/JitRunner.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/InitAllDialects.h"
OwnedBlob compilePtxToCubin(const std::string ptx, Location loc,
 
   return result;
 }
 
 static LogicalResult runMLIRPasses(ModuleOp m) {
   PassManager pm(m.getContext());
   applyPassManagerCLOptions(pm);
 
   const char gpuBinaryAnnotation[] = "nvvm.cubin";
   pm.addPass(createGpuKernelOutliningPass());
   auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
   kernelPm.addPass(createStripDebugInfoPass());
   kernelPm.addPass(createLowerGpuOpsToNVVMOpsPass());
   kernelPm.addPass(createConvertGPUKernelToBlobPass(
       translateModuleToNVVMIR, compilePtxToCubin, "nvptx64-nvidia-cuda",
       "sm_35", "+ptx60", gpuBinaryAnnotation));
+  auto &funcPm = pm.nest<FuncOp>();
+  funcPm.addPass(createGpuAsyncRegionPass());
+  funcPm.addPass(createAsyncRefCountingPass());
   pm.addPass(createGpuToLLVMConversionPass(gpuBinaryAnnotation));
+  pm.addPass(createAsyncToAsyncRuntimePass());
+  pm.addPass(createConvertAsyncToLLVMPass());
+  mlir::LowerToLLVMOptions lower_to_llvm_opts;
+  pm.addPass(mlir::createLowerToLLVMPass(lower_to_llvm_opts));
 
   return pm.run(m);
 }
 
 int main(int argc, char **argv) {
   registerPassManagerCLOptions();
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();