[MLIR] NFC: Rename mcuMemHostRegister* to mgpuMemHostRegister* to make it consistent with the other cuda-runner functions and ROCm.

[csigg]

Summary

Rename mcuMemHostRegister* to mgpuMemHostRegister*.

[herhut]

Thanks.

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;
 }

 extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   reportErrorIfAny(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0),
                    "MemHostRegister");
 }
 
 // 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);
 }