Clean up cuda-runtime-wrappers API.

[csigg]

Summary

Do not return error code, instead return created resource handles or void. Error reporting is done by the library function.

[csigg]

Merge from master.

[csigg]

Fix test.

[csigg]

Rebase from master, fix clang-tidy errors.

[herhut]

Thanks. I only have a minor nit.

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

68–69

should this be void?

[csigg]

Fix.

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

 
 using namespace mlir;
 
 // To avoid name mangling, these are defined in the mini-runtime file.
 static constexpr const char *kGpuModuleLoadName = "mgpuModuleLoad";
 static constexpr const char *kGpuModuleGetFunctionName =
     "mgpuModuleGetFunction";
 static constexpr const char *kGpuLaunchKernelName = "mgpuLaunchKernel";
-static constexpr const char *kGpuGetStreamHelperName = "mgpuGetStreamHelper";
+static constexpr const char *kGpuStreamCreateName = "mgpuStreamCreate";
 static constexpr const char *kGpuStreamSynchronizeName =
     "mgpuStreamSynchronize";
 static constexpr const char *kGpuMemHostRegisterName = "mgpuMemHostRegister";
 static constexpr const char *kGpuBinaryStorageSuffix = "_gpubin_cst";
 
 namespace {
 
 /// A pass to convert gpu.launch_func operations into a sequence of GPU

 // uses void pointers. This is fine as they have the same linkage in C.
 void GpuLaunchFuncToGpuRuntimeCallsPass::declareGpuRuntimeFunctions(
     Location loc) {
   ModuleOp module = getOperation();
   OpBuilder builder(module.getBody()->getTerminator());
   if (!module.lookupSymbol(kGpuModuleLoadName)) {
     builder.create<LLVM::LLVMFuncOp>(
         loc, kGpuModuleLoadName,
-        LLVM::LLVMType::getFunctionTy(
-            getGpuRuntimeResultType(),
-            {
-                getPointerPointerType(), /* CUmodule *module */
-                getPointerType()         /* void *cubin */
-            },
+        LLVM::LLVMType::getFunctionTy(getPointerType(),
+                                      getPointerType(), /* void *cubin */
                                       /*isVarArg=*/false));
   }
   if (!module.lookupSymbol(kGpuModuleGetFunctionName)) {
     // The helper uses void* instead of CUDA's opaque CUmodule and
     // CUfunction, or ROCm (HIP)'s opaque hipModule_t and hipFunction_t.
     builder.create<LLVM::LLVMFuncOp>(
         loc, kGpuModuleGetFunctionName,
-        LLVM::LLVMType::getFunctionTy(
-            getGpuRuntimeResultType(),
+        LLVM::LLVMType::getFunctionTy(getPointerType(),
                                       {
-                getPointerPointerType(), /* void **function */
-                getPointerType(),        /* void *module */
-                getPointerType()         /* char *name */
+                                          getPointerType(), /* void *module */
+                                          getPointerType()  /* char *name   */
                                       },
                                       /*isVarArg=*/false));
   }
   if (!module.lookupSymbol(kGpuLaunchKernelName)) {
     // Other than the CUDA or ROCm (HIP) api, the wrappers use uintptr_t to
     // match the LLVM type if MLIR's index type, which the GPU dialect uses.
     // Furthermore, they use void* instead of CUDA's opaque CUfunction and
     // CUstream, or ROCm (HIP)'s opaque hipFunction_t and hipStream_t.
     builder.create<LLVM::LLVMFuncOp>(
         loc, kGpuLaunchKernelName,
         LLVM::LLVMType::getFunctionTy(
-            getGpuRuntimeResultType(),
+            getVoidType(),
             {
                 getPointerType(),        /* void* f */
                 getIntPtrType(),         /* intptr_t gridXDim */
                 getIntPtrType(),         /* intptr_t gridyDim */
                 getIntPtrType(),         /* intptr_t gridZDim */
                 getIntPtrType(),         /* intptr_t blockXDim */
                 getIntPtrType(),         /* intptr_t blockYDim */
                 getIntPtrType(),         /* intptr_t blockZDim */
                 getInt32Type(),          /* unsigned int sharedMemBytes */
                 getPointerType(),        /* void *hstream */
                 getPointerPointerType(), /* void **kernelParams */
                 getPointerPointerType()  /* void **extra */
             },
             /*isVarArg=*/false));
   }
-  if (!module.lookupSymbol(kGpuGetStreamHelperName)) {
+  if (!module.lookupSymbol(kGpuStreamCreateName)) {
     // Helper function to get the current GPU compute stream. Uses void*
     // instead of CUDA's opaque CUstream, or ROCm (HIP)'s opaque hipStream_t.
     builder.create<LLVM::LLVMFuncOp>(
-        loc, kGpuGetStreamHelperName,
+        loc, kGpuStreamCreateName,
         LLVM::LLVMType::getFunctionTy(getPointerType(), /*isVarArg=*/false));
   }
   if (!module.lookupSymbol(kGpuStreamSynchronizeName)) {
     builder.create<LLVM::LLVMFuncOp>(
         loc, kGpuStreamSynchronizeName,
-        LLVM::LLVMType::getFunctionTy(getGpuRuntimeResultType(),
-                                      getPointerType() /* CUstream stream */,
+        LLVM::LLVMType::getFunctionTy(getVoidType(),
+                                      getPointerType() /* void *stream */,
                                       /*isVarArg=*/false));
   }
   if (!module.lookupSymbol(kGpuMemHostRegisterName)) {
     builder.create<LLVM::LLVMFuncOp>(
         loc, kGpuMemHostRegisterName,
         LLVM::LLVMType::getFunctionTy(getVoidType(),
                                       {
                                           getPointerType(), /* void *ptr */
return LLVM::createGlobalString(
       LLVM::Linkage::Internal, llvmDialect);
 }
 
 // Emits LLVM IR to launch a kernel function. Expects the module that contains
 // the compiled kernel function as a cubin in the 'nvvm.cubin' attribute, or a
 // hsaco in the 'rocdl.hsaco' attribute of the kernel function in the IR.
 //
 // %0 = call %binarygetter
-// %1 = alloca sizeof(void*)
-// call %moduleLoad(%2, %1)
-// %2 = alloca sizeof(void*)
-// %3 = load %1
-// %4 = <see generateKernelNameConstant>
-// call %moduleGetFunction(%2, %3, %4)
-// %5 = call %getStreamHelper()
-// %6 = load %2
-// %7 = <see setupParamsArray>
-// call %launchKernel(%6, <launchOp operands 0..5>, 0, %5, %7, nullptr)
-// call %streamSynchronize(%5)
+// %1 = call %moduleLoad(%0)
+// %2 = <see generateKernelNameConstant>
+// %3 = call %moduleGetFunction(%1, %2)
+// %4 = call %streamCreate()
+// %5 = <see setupParamsArray>
+// call %launchKernel(%3, <launchOp operands 0..5>, 0, %4, %5, nullptr)
+// call %streamSynchronize(%4)
 void GpuLaunchFuncToGpuRuntimeCallsPass::translateGpuLaunchCalls(
     mlir::gpu::LaunchFuncOp launchOp) {
   OpBuilder builder(launchOp);
   Location loc = launchOp.getLoc();
   declareGpuRuntimeFunctions(loc);
 
   auto zero = builder.create<LLVM::ConstantOp>(loc, getInt32Type(),
                                                builder.getI32IntegerAttr(0));
void GpuLaunchFuncToGpuRuntimeCallsPass::translateGpuLaunchCalls(
   SmallString<128> nameBuffer(kernelModule.getName());
   nameBuffer.append(kGpuBinaryStorageSuffix);
   Value data = LLVM::createGlobalString(
       loc, builder, nameBuffer.str(), binaryAttr.getValue(),
       LLVM::Linkage::Internal, getLLVMDialect());
 
   // Emit the load module call to load the module data. Error checking is done
   // in the called helper function.
-  auto gpuModule = allocatePointer(builder, loc);
   auto gpuModuleLoad =
       getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuModuleLoadName);
-  builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getGpuRuntimeResultType()},
-                               builder.getSymbolRefAttr(gpuModuleLoad),
-                               ArrayRef<Value>{gpuModule, data});
+  auto module = builder.create<LLVM::CallOp>(
+      loc, ArrayRef<Type>{getVoidType()},
+      builder.getSymbolRefAttr(gpuModuleLoad), ArrayRef<Value>{data});
   // Get the function from the module. The name corresponds to the name of
   // the kernel function.
-  auto gpuOwningModuleRef =
-      builder.create<LLVM::LoadOp>(loc, getPointerType(), gpuModule);
   auto kernelName = generateKernelNameConstant(
       launchOp.getKernelModuleName(), launchOp.getKernelName(), loc, builder);
-  auto gpuFunction = allocatePointer(builder, loc);
   auto gpuModuleGetFunction =
       getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuModuleGetFunctionName);
-  builder.create<LLVM::CallOp>(
-      loc, ArrayRef<Type>{getGpuRuntimeResultType()},
+  auto function = builder.create<LLVM::CallOp>(
+      loc, ArrayRef<Type>{getVoidType()},
       builder.getSymbolRefAttr(gpuModuleGetFunction),
-      ArrayRef<Value>{gpuFunction, gpuOwningModuleRef, kernelName});
+      ArrayRef<Value>{module.getResult(0), kernelName});
   // Grab the global stream needed for execution.
-  auto gpuGetStreamHelper =
-      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuGetStreamHelperName);
-  auto gpuStream = builder.create<LLVM::CallOp>(
+  auto gpuStreamCreate =
+      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuStreamCreateName);
+  auto stream = builder.create<LLVM::CallOp>(
       loc, ArrayRef<Type>{getPointerType()},
-      builder.getSymbolRefAttr(gpuGetStreamHelper), ArrayRef<Value>{});
+      builder.getSymbolRefAttr(gpuStreamCreate), ArrayRef<Value>{});
   // Invoke the function with required arguments.
   auto gpuLaunchKernel =
       getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuLaunchKernelName);
-  auto gpuFunctionRef =
-      builder.create<LLVM::LoadOp>(loc, getPointerType(), gpuFunction);
   auto paramsArray = setupParamsArray(launchOp, builder);
   if (!paramsArray) {
     launchOp.emitOpError() << "cannot pass given parameters to the kernel";
     return signalPassFailure();
   }
   auto nullpointer =
       builder.create<LLVM::IntToPtrOp>(loc, getPointerPointerType(), zero);
   builder.create<LLVM::CallOp>(
       loc, ArrayRef<Type>{getGpuRuntimeResultType()},
       builder.getSymbolRefAttr(gpuLaunchKernel),
-      ArrayRef<Value>{gpuFunctionRef, launchOp.getOperand(0),
+      ArrayRef<Value>{function.getResult(0), launchOp.getOperand(0),
                       launchOp.getOperand(1), launchOp.getOperand(2),
                       launchOp.getOperand(3), launchOp.getOperand(4),
                       launchOp.getOperand(5), zero, /* sharedMemBytes */
-                      gpuStream.getResult(0),       /* stream */
+                      stream.getResult(0),          /* stream */
                       paramsArray,                  /* kernel params */
                       nullpointer /* extra */});
   // Sync on the stream to make it synchronous.
   auto gpuStreamSync =
       getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuStreamSynchronizeName);
   builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getGpuRuntimeResultType()},
                                builder.getSymbolRefAttr(gpuStreamSync),
-                               ArrayRef<Value>(gpuStream.getResult(0)));
+                               ArrayRef<Value>(stream.getResult(0)));
   launchOp.erase();
 }
 
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
     StringRef gpuBinaryAnnotation) {
   if (gpuBinaryAnnotation.empty())
     return std::make_unique<GpuLaunchFuncToGpuRuntimeCallsPass>();
   return std::make_unique<GpuLaunchFuncToGpuRuntimeCallsPass>(
       gpuBinaryAnnotation);
 }

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

llvm.func @foo() {
     // CHECK: %[[addressof:.*]] = llvm.mlir.addressof @[[global]]
     // CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : index)
     // CHECK: %[[binary_ptr:.*]] = llvm.getelementptr %[[addressof]][%[[c0]], %[[c0]]]
     // CHECK-SAME: -> !llvm<"i8*">
     // CHECK: %[[module_ptr:.*]] = llvm.alloca {{.*}} x !llvm<"i8*"> : (!llvm.i32) -> !llvm<"i8**">
     // CHECK: llvm.call @mgpuModuleLoad(%[[module_ptr]], %[[binary_ptr]]) : (!llvm<"i8**">, !llvm<"i8*">) -> !llvm.i32
     // CHECK: %[[func_ptr:.*]] = llvm.alloca {{.*}} x !llvm<"i8*"> : (!llvm.i32) -> !llvm<"i8**">
     // CHECK: llvm.call @mgpuModuleGetFunction(%[[func_ptr]], {{.*}}, {{.*}}) : (!llvm<"i8**">, !llvm<"i8*">, !llvm<"i8*">) -> !llvm.i32
-    // CHECK: llvm.call @mgpuGetStreamHelper
+    // CHECK: llvm.call @mgpuStreamCreate
     // CHECK: llvm.call @mgpuLaunchKernel
     // CHECK: llvm.call @mgpuStreamSynchronize
     "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernel_module::@kernel }
         : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.float, !llvm<"float*">) -> ()
 
     llvm.return
   }
 
 }

mlir/test/mlir-cuda-runner/all-reduce-and.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
-func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
+func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-max.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [16, 11]
 
   return
 }
 
-func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
+func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-min.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
-func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
+func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-op.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK-COUNT-8: [{{(5356, ){12}5356}}]
 func @main() {
   %arg = alloc() : memref<2x4x13xf32>
   %dst = memref_cast %arg : memref<2x4x13xf32> to memref<?x?x?xf32>
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %sx = dim %dst, %c2 : memref<?x?x?xf32>
   %sy = dim %dst, %c1 : memref<?x?x?xf32>
   %sz = dim %dst, %c0 : memref<?x?x?xf32>
   %cast_dst = memref_cast %dst : memref<?x?x?xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%cast_dst) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%cast_dst) : (memref<*xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %sy, %block_z = %sz) {
     %t0 = muli %tz, %block_y : index
     %t1 = addi %ty, %t0 : index
     %t2 = muli %t1, %block_x : index
     %idx = addi %tx, %t2 : index
     %t3 = index_cast %idx : index to i32
     %val = sitofp %t3 : i32 to f32
     %sum = "gpu.all_reduce"(%val) ({}) { op = "add" } : (f32) -> (f32)
     store %sum, %dst[%tz, %ty, %tx] : memref<?x?x?xf32>
     gpu.terminator
   }
   call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/all-reduce-or.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 15]
 
   return
 }
 
-func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
+func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/all-reduce-region.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [{{(35, ){34}35}}]
 func @main() {
   %arg = alloc() : memref<35xf32>
   %dst = memref_cast %arg : memref<35xf32> to memref<?xf32>
   %one = constant 1 : index
   %c0 = constant 0 : index
   %sx = dim %dst, %c0 : memref<?xf32>
   %cast_dst = memref_cast %dst : memref<?xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%cast_dst) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%cast_dst) : (memref<*xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %val = index_cast %tx : index to i32
     %xor = "gpu.all_reduce"(%val) ({
     ^bb(%lhs : i32, %rhs : i32):
       %xor = xor %lhs, %rhs : i32
       "gpu.yield"(%xor) : (i32) -> ()
     }) : (i32) -> (i32)
     %res = sitofp %xor : i32 to f32
     store %res, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
   call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(memref<*xf32>)

mlir/test/mlir-cuda-runner/all-reduce-xor.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
   %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_sum) : (memref<*xi32>) -> ()
 
   store %cst0, %data[%c0, %c0] : memref<2x6xi32>
   store %cst1, %data[%c0, %c1] : memref<2x6xi32>
   store %cst2, %data[%c0, %c2] : memref<2x6xi32>
   store %cst4, %data[%c0, %c3] : memref<2x6xi32>
   store %cst8, %data[%c0, %c4] : memref<2x6xi32>
   store %cst16, %data[%c0, %c5] : memref<2x6xi32>
 
func @main() {
   }
 
   call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 1]
 
   return
 }
 
-func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
+func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

mlir/test/mlir-cuda-runner/gpu-to-cubin.mlir

 }
 
 // CHECK: [1, 1, 1, 1, 1]
 func @main() {
   %arg0 = alloc() : memref<5xf32>
   %21 = constant 5 : i32
   %22 = memref_cast %arg0 : memref<5xf32> to memref<?xf32>
   %23 = memref_cast %22 : memref<?xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%23) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%23) : (memref<*xf32>) -> ()
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   %24 = constant 1.0 : f32
   call @other_func(%24, %22) : (f32, memref<?xf32>) -> ()
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/multiple-all-reduce.mlir

func @main() {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
 
   %cast_data = memref_cast %data : memref<2x6xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%cast_data) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%cast_data) : (memref<*xf32>) -> ()
   %cast_sum = memref_cast %sum : memref<2xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%cast_sum) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%cast_sum) : (memref<*xf32>) -> ()
   %cast_mul = memref_cast %mul : memref<2xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%cast_mul) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%cast_mul) : (memref<*xf32>) -> ()
 
   store %cst0, %data[%c0, %c0] : memref<2x6xf32>
   store %cst1, %data[%c0, %c1] : memref<2x6xf32>
   store %cst2, %data[%c0, %c2] : memref<2x6xf32>
   store %cst4, %data[%c0, %c3] : memref<2x6xf32>
   store %cst8, %data[%c0, %c4] : memref<2x6xf32>
   store %cst16, %data[%c0, %c5] : memref<2x6xf32>
 
func @main() {
   // CHECK: [31, 39]
 
   call @print_memref_f32(%cast_mul) : (memref<*xf32>) -> ()
   // CHECK: [0, 27720]
 
   return
 }
 
-func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(memref<*xf32>)

mlir/test/mlir-cuda-runner/shuffle.mlir

 // RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 // CHECK: [4, 5, 6, 7, 0, 1, 2, 3, 12, -1, -1, -1, 8]
 func @main() {
   %arg = alloc() : memref<13xf32>
   %dst = memref_cast %arg : memref<13xf32> to memref<?xf32>
   %one = constant 1 : index
   %c0 = constant 0 : index
   %sx = dim %dst, %c0 : memref<?xf32>
   %cast_dest = memref_cast %dst : memref<?xf32> to memref<*xf32>
-  call @mcuMemHostRegisterFloat(%cast_dest) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%cast_dest) : (memref<*xf32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %t0 = index_cast %tx : index to i32
     %val = sitofp %t0 : i32 to f32
     %width = index_cast %block_x : index to i32
     %offset = constant 4 : i32
     %shfl, %valid = gpu.shuffle %val, %offset, %width xor : f32
     cond_br %valid, ^bb1(%shfl : f32), ^bb0
   ^bb0:
     %m1 = constant -1.0 : f32
     br ^bb1(%m1 : f32)
   ^bb1(%value : f32):
     store %value, %dst[%tx] : memref<?xf32>
     gpu.terminator
   }
   call @print_memref_f32(%cast_dest) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-cuda-runner/two-modules.mlir

 // RUN: mlir-cuda-runner %s --print-ir-after-all --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s --dump-input=always
 
 // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
 func @main() {
   %arg = alloc() : memref<13xi32>
   %dst = memref_cast %arg : memref<13xi32> to memref<?xi32>
   %one = constant 1 : index
   %c0 = constant 0 : index
   %sx = dim %dst, %c0 : memref<?xi32>
   %cast_dst = memref_cast %dst : memref<?xi32> to memref<*xi32>
-  call @mcuMemHostRegisterInt32(%cast_dst) : (memref<*xi32>) -> ()
+  call @mgpuMemHostRegisterInt32(%cast_dst) : (memref<*xi32>) -> ()
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %t0 = index_cast %tx : index to i32
     store %t0, %dst[%tx] : memref<?xi32>
     gpu.terminator
   }
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
     %t0 = index_cast %tx : index to i32
     store %t0, %dst[%tx] : memref<?xi32>
     gpu.terminator
   }
   call @print_memref_i32(%cast_dst) : (memref<*xi32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterInt32(%memref : memref<*xi32>)
+func @mgpuMemHostRegisterInt32(%memref : memref<*xi32>)
 func @print_memref_i32(%memref : memref<*xi32>)

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

 
 #include <cassert>
 #include <numeric>
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include "cuda.h"
 
-namespace {
-int32_t reportErrorIfAny(CUresult result, const char *where) {
-  if (result != CUDA_SUCCESS) {
-    llvm::errs() << "CUDA failed with " << result << " in " << where << "\n";
-  }
-  return result;
-}
-} // anonymous namespace
-
-extern "C" int32_t mgpuModuleLoad(void **module, void *data) {
-  int32_t err = reportErrorIfAny(
-      cuModuleLoadData(reinterpret_cast<CUmodule *>(module), data),
-      "ModuleLoad");
-  return err;
-}
-
-extern "C" int32_t mgpuModuleGetFunction(void **function, void *module,
-                                         const char *name) {
-  return reportErrorIfAny(
-      cuModuleGetFunction(reinterpret_cast<CUfunction *>(function),
-                          reinterpret_cast<CUmodule>(module), name),
-      "GetFunction");
+#define CUDA_REPORT_IF_ERROR(expr)                                             \
+  [](CUresult result) {                                                        \
+    if (!result)                                                               \
+      return;                                                                  \
+    const char *name = nullptr;                                                \
+    cuGetErrorName(result, &name);                                             \
+    if (!name)                                                                 \
+      name = "<unknown>";                                                      \
+    llvm::errs() << "'" << #expr << "' failed with '" << name << "'\n";        \
+  }(expr)
+
+extern "C" CUmodule mgpuModuleLoad(void *data) {
+  CUmodule module = nullptr;
+  CUDA_REPORT_IF_ERROR(cuModuleLoadData(&module, data));
+  return module;
+}
+
+extern "C" CUfunction mgpuModuleGetFunction(CUmodule module, const char *name) {
+  CUfunction function = nullptr;
+  CUDA_REPORT_IF_ERROR(cuModuleGetFunction(&function, module, name));
+  return function;
 }
 
 // The wrapper uses intptr_t instead of CUDA's unsigned int to match
 // the type of MLIR's index type. This avoids the need for casts in the
 // generated MLIR code.
-extern "C" int32_t mgpuLaunchKernel(void *function, intptr_t gridX,
-                                    intptr_t gridY, intptr_t gridZ,
-                                    intptr_t blockX, intptr_t blockY,
-                                    intptr_t blockZ, int32_t smem, void *stream,
-                                    void **params, void **extra) {
-  return reportErrorIfAny(
-      cuLaunchKernel(reinterpret_cast<CUfunction>(function), gridX, gridY,
-                     gridZ, blockX, blockY, blockZ, smem,
-                     reinterpret_cast<CUstream>(stream), params, extra),
-      "LaunchKernel");
+extern "C" void mgpuLaunchKernel(CUfunction function, intptr_t gridX,
+                                 intptr_t blockZ, int32_t smem, CUstream stream,
+  CUDA_REPORT_IF_ERROR(cuLaunchKernel(function, gridX, gridY, gridZ, blockX,
+                                      blockY, blockZ, smem, stream, params,
+                                      extra));
 }
 
-extern "C" void *mgpuGetStreamHelper() {
-  CUstream stream;
-  reportErrorIfAny(cuStreamCreate(&stream, CU_STREAM_DEFAULT), "StreamCreate");
+extern "C" CUstream mgpuStreamCreate() {
+  CUstream stream = nullptr;
+  CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING));
   return stream;
 }
 
-extern "C" int32_t mgpuStreamSynchronize(void *stream) {
-  return reportErrorIfAny(
-      cuStreamSynchronize(reinterpret_cast<CUstream>(stream)), "StreamSync");
+extern "C" void mgpuStreamSynchronize(CUstream stream) {
+  CUDA_REPORT_IF_ERROR(cuStreamSynchronize(stream));
 }
 
 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the CUDA runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
-  reportErrorIfAny(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0),
-                   "MemHostRegister");
+  CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0));
 }
 
 // Allows to register a MemRef with the CUDA runtime. Initializes array with
 // value. Helpful until we have transfer functions implemented.
 template <typename T>
-void mcuMemHostRegisterMemRef(const DynamicMemRefType<T> &mem_ref, T value) {
+void mgpuMemHostRegisterMemRef(const DynamicMemRefType<T> &mem_ref, T value) {
   llvm::SmallVector<int64_t, 4> denseStrides(mem_ref.rank);
   llvm::ArrayRef<int64_t> sizes(mem_ref.sizes, mem_ref.rank);
   llvm::ArrayRef<int64_t> strides(mem_ref.strides, mem_ref.rank);
 
   std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),
                    std::multiplies<int64_t>());
   auto count = denseStrides.front();
 
   // Only densely packed tensors are currently supported.
   std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
               denseStrides.end());
   denseStrides.back() = 1;
   assert(strides == llvm::makeArrayRef(denseStrides));
 
   auto *pointer = mem_ref.data + mem_ref.offset;
   std::fill_n(pointer, count, value);
   mgpuMemHostRegister(pointer, count * sizeof(T));
 }
 
-extern "C" void mcuMemHostRegisterFloat(int64_t rank, void *ptr) {
+extern "C" void mgpuMemHostRegisterFloat(int64_t rank, void *ptr) {
   UnrankedMemRefType<float> mem_ref = {rank, ptr};
-  mcuMemHostRegisterMemRef(DynamicMemRefType<float>(mem_ref), 1.23f);
+  mgpuMemHostRegisterMemRef(DynamicMemRefType<float>(mem_ref), 1.23f);
 }
 
-extern "C" void mcuMemHostRegisterInt32(int64_t rank, void *ptr) {
+extern "C" void mgpuMemHostRegisterInt32(int64_t rank, void *ptr) {
   UnrankedMemRefType<int32_t> mem_ref = {rank, ptr};
-  mcuMemHostRegisterMemRef(DynamicMemRefType<int32_t>(mem_ref), 123);
+  mgpuMemHostRegisterMemRef(DynamicMemRefType<int32_t>(mem_ref), 123);
 }

mlir/tools/mlir-rocm-runner/rocm-runtime-wrappers.cpp

 
 #include <cassert>
 #include <numeric>
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include "hip/hip_runtime.h"
 
-namespace {
-int32_t reportErrorIfAny(hipError_t result, const char *where) {
-  if (result != hipSuccess) {
-    llvm::errs() << "HIP failed with " << result << " in " << where << "\n";
-  }
-  return result;
-}
-} // anonymous namespace
-
-extern "C" int32_t mgpuModuleLoad(void **module, void *data) {
-  int32_t err = reportErrorIfAny(
-      hipModuleLoadData(reinterpret_cast<hipModule_t *>(module), data),
-      "ModuleLoad");
-  return err;
+#define HIP_REPORT_IF_ERROR(expr)                                              \
+  [](hipError_t result) {                                                      \
+    if (!result)                                                               \
+      return;                                                                  \
+    const char *name = nullptr;                                                \
+    hipGetErrorName(result, &name);                                            \
+    if (!name)                                                                 \
+      name = "<unknown>";                                                      \
+    llvm::errs() << "'" << #expr << "' failed with '" << name << "'\n";        \
+  }(expr)
+
+extern "C" hipModule_t mgpuModuleLoad(void *data) {
+  hipModule_t module = nullptr;
+  HIP_REPORT_IF_ERROR(hipModuleLoadData(&module, data));
+  return module;
 }
 
-extern "C" int32_t mgpuModuleGetFunction(void **function, void *module,
+extern "C" hipFunction_t mgpuModuleGetFunction(hipModule_t module,
                                                const char *name) {
-  return reportErrorIfAny(
-      hipModuleGetFunction(reinterpret_cast<hipFunction_t *>(function),
-                           reinterpret_cast<hipModule_t>(module), name),
-      "GetFunction");
+  hipFunction_t function = nullptr;
+  HIP_REPORT_IF_ERROR(hipModuleGetFunction(&function, module, name));
+  return function;
 }
 
 // The wrapper uses intptr_t instead of ROCM's unsigned int to match
 // the type of MLIR's index type. This avoids the need for casts in the
 // generated MLIR code.
-extern "C" int32_t mgpuLaunchKernel(void *function, intptr_t gridX,
-                                    intptr_t gridY, intptr_t gridZ,
-                                    intptr_t blockX, intptr_t blockY,
-                                    intptr_t blockZ, int32_t smem, void *stream,
-                                    void **params, void **extra) {
-  return reportErrorIfAny(
-      hipModuleLaunchKernel(reinterpret_cast<hipFunction_t>(function), gridX,
-                            gridY, gridZ, blockX, blockY, blockZ, smem,
-                            reinterpret_cast<hipStream_t>(stream), params,
-                            extra),
-      "LaunchKernel");
+extern "C" void mgpuLaunchKernel(hipFunction_t function, intptr_t gridX,
+                                 intptr_t blockZ, int32_t smem,
+                                 hipStream_t stream, void **params,
+                                 void **extra) {
+  HIP_REPORT_IF_ERROR(hipModuleLaunchKernel(function, gridX, gridY, gridZ,
+                                            blockX, blockY, blockZ, smem,
+                                            stream, params, extra));
 }
 
-extern "C" void *mgpuGetStreamHelper() {
-  hipStream_t stream;
-  reportErrorIfAny(hipStreamCreate(&stream), "StreamCreate");
+extern "C" void *mgpuStreamCreate() {
+  hipStream_t stream = nullptr;
+  HIP_REPORT_IF_ERROR(hipStreamCreate(&stream));
   return stream;
 }
 
-extern "C" int32_t mgpuStreamSynchronize(void *stream) {
-  return reportErrorIfAny(
-      hipStreamSynchronize(reinterpret_cast<hipStream_t>(stream)),
-      "StreamSync");
+extern "C" int32_t mgpuStreamSynchronize(hipStream_t stream) {
+  return HIP_REPORT_IF_ERROR(hipStreamSynchronize(stream));
 }
 
 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the ROCM runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
-  reportErrorIfAny(hipHostRegister(ptr, sizeBytes, /*flags=*/0),
+  HIP_REPORT_IF_ERROR(hipHostRegister(ptr, sizeBytes, /*flags=*/0),
                       "MemHostRegister");
 }
 
 // Allows to register a MemRef with the ROCM runtime. Initializes array with
 // value. Helpful until we have transfer functions implemented.
 template <typename T>
 void mgpuMemHostRegisterMemRef(T *pointer, llvm::ArrayRef<int64_t> sizes,
                                llvm::ArrayRef<int64_t> strides, T value) {
   assert(sizes.size() == strides.size());
extern "C" void mgpuMemHostRegisterInt32(int64_t rank, void *ptr) {
   auto *desc = static_cast<StridedMemRefType<int32_t, 1> *>(ptr);
   auto sizes = llvm::ArrayRef<int64_t>(desc->sizes, rank);
   auto strides = llvm::ArrayRef<int64_t>(desc->sizes + rank, rank);
   mgpuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 123);
 }
 
 template <typename T>
 void mgpuMemGetDevicePointer(T *hostPtr, T **devicePtr) {
-  reportErrorIfAny(hipSetDevice(0), "hipSetDevice");
-  reportErrorIfAny(
+  HIP_REPORT_IF_ERROR(hipSetDevice(0), "hipSetDevice");
+  HIP_REPORT_IF_ERROR(
       hipHostGetDevicePointer((void **)devicePtr, hostPtr, /*flags=*/0),
       "hipHostGetDevicePointer");
 }
 
 extern "C" StridedMemRefType<float, 1>
 mgpuMemGetDeviceMemRef1dFloat(float *allocated, float *aligned, int64_t offset,
                               int64_t size, int64_t stride) {
   float *devicePtr = nullptr;