Unrank mcuMemHostRegister tensor argument.

[csigg]

The presubmit failure is from clang-tidy:

/mnt/disks/ssd0/agent/workspace/BETA_amd64_debian_testing_clang8/llvm-project/mlir/tools/mlir-cuda-runner/cuda-runtime-wrappers.cpp:22:10: error: 'cuda.h' file not found [clang-diagnostic-error]
#include "cuda.h"
         ^

Is there anything I can/need to do about it?

[herhut]

Nice cleanup, thanks!

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

106

This is a bit fishy, as you rely on the inner layout of the descriptor declared in a different file. It would be nicer to have an interface in that header that allows to do this.

[csigg]

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

106

Will do that in a separate change.

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

func @main() {
   %c0 = constant 0 : index
   %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<?x?xi32>
-  call @mcuMemHostRegisterMemRef2dInt32(%cast_data) : (memref<?x?xi32>) -> ()
-  %cast_sum = memref_cast %sum : memref<2xi32> to memref<?xi32>
-  call @mcuMemHostRegisterMemRef1dInt32(%cast_sum) : (memref<?xi32>) -> ()
+  %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%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>
 
   store %cst2, %data[%c1, %c0] : memref<2x6xi32>
   store %cst3, %data[%c1, %c1] : memref<2x6xi32>
   store %cst6, %data[%c1, %c2] : memref<2x6xi32>
   store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // AND
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = load %data[%bx, %tx] : memref<2x6xi32>
     %reduced = "gpu.all_reduce"(%val) ({}) { op = "and" } : (i32) -> (i32)
     store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
-  %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
+  call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
-func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
-func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
+func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

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

func @main() {
   %c0 = constant 0 : index
   %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<?x?xi32>
-  call @mcuMemHostRegisterMemRef2dInt32(%cast_data) : (memref<?x?xi32>) -> ()
-  %cast_sum = memref_cast %sum : memref<2xi32> to memref<?xi32>
-  call @mcuMemHostRegisterMemRef1dInt32(%cast_sum) : (memref<?xi32>) -> ()
+  %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%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>
 
   store %cst2, %data[%c1, %c0] : memref<2x6xi32>
   store %cst3, %data[%c1, %c1] : memref<2x6xi32>
   store %cst6, %data[%c1, %c2] : memref<2x6xi32>
   store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // MAX
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = load %data[%bx, %tx] : memref<2x6xi32>
     %reduced = "gpu.all_reduce"(%val) ({}) { op = "max" } : (i32) -> (i32)
     store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
-  %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
+  call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [16, 11]
 
   return
 }
 
-func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
-func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
+func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

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

func @main() {
   %c0 = constant 0 : index
   %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<?x?xi32>
-  call @mcuMemHostRegisterMemRef2dInt32(%cast_data) : (memref<?x?xi32>) -> ()
-  %cast_sum = memref_cast %sum : memref<2xi32> to memref<?xi32>
-  call @mcuMemHostRegisterMemRef1dInt32(%cast_sum) : (memref<?xi32>) -> ()
+  %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%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>
 
   store %cst2, %data[%c1, %c0] : memref<2x6xi32>
   store %cst3, %data[%c1, %c1] : memref<2x6xi32>
   store %cst6, %data[%c1, %c2] : memref<2x6xi32>
   store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // MIN
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = load %data[%bx, %tx] : memref<2x6xi32>
     %reduced = "gpu.all_reduce"(%val) ({}) { op = "min" } : (i32) -> (i32)
     store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
-  %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
+  call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
-func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
-func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
+func @mcuMemHostRegisterInt32(%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>
   %one = constant 1 : index
   %sx = dim %dst, 2 : memref<?x?x?xf32>
   %sy = dim %dst, 1 : memref<?x?x?xf32>
   %sz = dim %dst, 0 : memref<?x?x?xf32>
-  call @mcuMemHostRegisterMemRef3dFloat(%dst) : (memref<?x?x?xf32>) -> ()
+  %cast_dst = memref_cast %dst : memref<?x?x?xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%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 = %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
   }
-  %U = memref_cast %dst : memref<?x?x?xf32> to memref<*xf32>
-  call @print_memref_f32(%U) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterMemRef3dFloat(%ptr : memref<?x?x?xf32>)
+func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

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

func @main() {
   %c0 = constant 0 : index
   %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<?x?xi32>
-  call @mcuMemHostRegisterMemRef2dInt32(%cast_data) : (memref<?x?xi32>) -> ()
-  %cast_sum = memref_cast %sum : memref<2xi32> to memref<?xi32>
-  call @mcuMemHostRegisterMemRef1dInt32(%cast_sum) : (memref<?xi32>) -> ()
+  %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%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>
 
   store %cst2, %data[%c1, %c0] : memref<2x6xi32>
   store %cst3, %data[%c1, %c1] : memref<2x6xi32>
   store %cst6, %data[%c1, %c2] : memref<2x6xi32>
   store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // OR
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = load %data[%bx, %tx] : memref<2x6xi32>
     %reduced = "gpu.all_reduce"(%val) ({}) { op = "or" } : (i32) -> (i32)
     store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
-  %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
+  call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 15]
 
   return
 }
 
-func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
-func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
+func @mcuMemHostRegisterInt32(%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
   %sx = dim %dst, 0 : memref<?xf32>
-  call @mcuMemHostRegisterMemRef1dFloat(%dst) : (memref<?xf32>) -> ()
+  %cast_dst = memref_cast %dst : memref<?xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%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
   }
-  %U = memref_cast %dst : memref<?xf32> to memref<*xf32>
-  call @print_memref_f32(%U) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%cast_dst) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
+func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(memref<*xf32>)

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

func @main() {
   %c0 = constant 0 : index
   %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<?x?xi32>
-  call @mcuMemHostRegisterMemRef2dInt32(%cast_data) : (memref<?x?xi32>) -> ()
-  %cast_sum = memref_cast %sum : memref<2xi32> to memref<?xi32>
-  call @mcuMemHostRegisterMemRef1dInt32(%cast_sum) : (memref<?xi32>) -> ()
+  %cast_data = memref_cast %data : memref<2x6xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%cast_data) : (memref<*xi32>) -> ()
+  %cast_sum = memref_cast %sum : memref<2xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%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>
 
   store %cst2, %data[%c1, %c0] : memref<2x6xi32>
   store %cst3, %data[%c1, %c1] : memref<2x6xi32>
   store %cst6, %data[%c1, %c2] : memref<2x6xi32>
   store %cst7, %data[%c1, %c3] : memref<2x6xi32>
   store %cst10, %data[%c1, %c4] : memref<2x6xi32>
   store %cst11, %data[%c1, %c5] : memref<2x6xi32>
 
   // XOR
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c6, %block_y = %c1, %block_z = %c1) {
     %val = load %data[%bx, %tx] : memref<2x6xi32>
     %reduced = "gpu.all_reduce"(%val) ({}) { op = "xor" } : (i32) -> (i32)
     store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
-  %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
-  call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
+  call @print_memref_i32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [31, 1]
 
   return
 }
 
-func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
-func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
+func @mcuMemHostRegisterInt32(%ptr : memref<*xi32>)
 func @print_memref_i32(memref<*xi32>)
 

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

func @other_func(%arg0 : f32, %arg1 : memref<?xf32>) {
   return
 }
 
 // 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>
-  call @mcuMemHostRegisterMemRef1dFloat(%22) : (memref<?xf32>) -> ()
   %23 = memref_cast %22 : memref<?xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%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 @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
+func @mcuMemHostRegisterFloat(%ptr : memref<*xf32>)
 func @print_memref_f32(%ptr : memref<*xf32>)

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

func @main() {
   %c0 = constant 0 : index
   %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<?x?xf32>
-  call @mcuMemHostRegisterMemRef2dFloat(%cast_data) : (memref<?x?xf32>) -> ()
-  %cast_sum = memref_cast %sum : memref<2xf32> to memref<?xf32>
-  call @mcuMemHostRegisterMemRef1dFloat(%cast_sum) : (memref<?xf32>) -> ()
-  %cast_mul = memref_cast %mul : memref<2xf32> to memref<?xf32>
-  call @mcuMemHostRegisterMemRef1dFloat(%cast_mul) : (memref<?xf32>) -> ()
+  %cast_data = memref_cast %data : memref<2x6xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%cast_data) : (memref<*xf32>) -> ()
+  %cast_sum = memref_cast %sum : memref<2xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%cast_sum) : (memref<*xf32>) -> ()
+  %cast_mul = memref_cast %mul : memref<2xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%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>
 
gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c2, %grid_y = %c1, %grid_z = %c1)
     %val = load %data[%bx, %tx] : memref<2x6xf32>
     %reduced0 = "gpu.all_reduce"(%val) ({}) { op = "add" } : (f32) -> (f32)
     store %reduced0, %sum[%bx] : memref<2xf32>
     %reduced1 = "gpu.all_reduce"(%val) ({}) { op = "mul" } : (f32) -> (f32)
     store %reduced1, %mul[%bx] : memref<2xf32>
     gpu.terminator
   }
 
-  %ptr_sum = memref_cast %sum : memref<2xf32> to memref<*xf32>
-  call @print_memref_f32(%ptr_sum) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%cast_sum) : (memref<*xf32>) -> ()
   // CHECK: [31, 39]
 
-  %ptr_mul = memref_cast %mul : memref<2xf32> to memref<*xf32>
-  call @print_memref_f32(%ptr_mul) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%cast_mul) : (memref<*xf32>) -> ()
   // CHECK: [0, 27720]
 
   return
 }
 
-func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
-func @mcuMemHostRegisterMemRef2dFloat(%ptr : memref<?x?xf32>)
+func @mcuMemHostRegisterFloat(%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
   %sx = dim %dst, 0 : memref<?xf32>
-  call @mcuMemHostRegisterMemRef1dFloat(%dst) : (memref<?xf32>) -> ()
+  %cast_dest = memref_cast %dst : memref<?xf32> to memref<*xf32>
+  call @mcuMemHostRegisterFloat(%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
   }
-  %U = memref_cast %dst : memref<?xf32> to memref<*xf32>
-  call @print_memref_f32(%U) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%cast_dest) : (memref<*xf32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
+func @mcuMemHostRegisterFloat(%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
   %sx = dim %dst, 0 : memref<?xi32>
-  call @mcuMemHostRegisterMemRef1dInt32(%dst) : (memref<?xi32>) -> ()
+  %cast_dst = memref_cast %dst : memref<?xi32> to memref<*xi32>
+  call @mcuMemHostRegisterInt32(%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
   }
-  %U = memref_cast %dst : memref<?xi32> to memref<*xi32>
-  call @print_memref_i32(%U) : (memref<*xi32>) -> ()
+  call @print_memref_i32(%cast_dst) : (memref<*xi32>) -> ()
   return
 }
 
-func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
-func @print_memref_i32(%ptr : memref<*xi32>)
+func @mcuMemHostRegisterInt32(%memref : memref<*xi32>)
+func @print_memref_i32(%memref : memref<*xi32>)

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

 // Also adds some debugging helpers that are helpful when writing MLIR code to
 // run on GPUs.
 //
 //===----------------------------------------------------------------------===//
 
 #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) {

 
 // Allows to register byte array with the CUDA runtime. Helpful until we have
 // transfer functions implemented.
 extern "C" void mcuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   reportErrorIfAny(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0),
                    "MemHostRegister");
 }
 
-// A struct that corresponds to how MLIR represents memrefs.
-template <typename T, int N> struct MemRefType {
-  T *basePtr;
-  T *data;
-  int64_t offset;
-  int64_t sizes[N];
-  int64_t strides[N];
-};
-
 // 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(T *pointer, llvm::ArrayRef<int64_t> sizes,
                               llvm::ArrayRef<int64_t> strides, T value) {
   assert(sizes.size() == strides.size());
   llvm::SmallVector<int64_t, 4> denseStrides(strides.size());
 
   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));
 
   std::fill_n(pointer, count, value);
   mcuMemHostRegister(pointer, count * sizeof(T));
 }
 
-extern "C" void mcuMemHostRegisterMemRef1dFloat(float *allocated,
-                                                float *aligned, int64_t offset,
-                                                int64_t size, int64_t stride) {
-  mcuMemHostRegisterMemRef(aligned + offset, {size}, {stride}, 1.23f);
-}
-
-extern "C" void mcuMemHostRegisterMemRef2dFloat(float *allocated,
-                                                float *aligned, int64_t offset,
-                                                int64_t size0, int64_t size1,
-                                                int64_t stride0,
-                                                int64_t stride1) {
-  mcuMemHostRegisterMemRef(aligned + offset, {size0, size1}, {stride0, stride1},
-                           1.23f);
-}
-
-extern "C" void mcuMemHostRegisterMemRef3dFloat(float *allocated,
-                                                float *aligned, int64_t offset,
-                                                int64_t size0, int64_t size1,
-                                                int64_t size2, int64_t stride0,
-                                                int64_t stride1,
-                                                int64_t stride2) {
-  mcuMemHostRegisterMemRef(aligned + offset, {size0, size1, size2},
-                           {stride0, stride1, stride2}, 1.23f);
-}
-
-extern "C" void mcuMemHostRegisterMemRef1dInt32(int32_t *allocated,
-                                                int32_t *aligned,
-                                                int64_t offset, int64_t size,
-                                                int64_t stride) {
-  mcuMemHostRegisterMemRef(aligned + offset, {size}, {stride}, 123);
-}
-
-extern "C" void mcuMemHostRegisterMemRef2dInt32(int32_t *allocated,
-                                                int32_t *aligned,
-                                                int64_t offset, int64_t size0,
-                                                int64_t size1, int64_t stride0,
-                                                int64_t stride1) {
-  mcuMemHostRegisterMemRef(aligned + offset, {size0, size1}, {stride0, stride1},
-                           123);
-}
-
-extern "C" void
-mcuMemHostRegisterMemRef3dInt32(int32_t *allocated, int32_t *aligned,
-                                int64_t offset, int64_t size0, int64_t size1,
-                                int64_t size2, int64_t stride0, int64_t stride1,
-                                int64_t stride2) {
-  mcuMemHostRegisterMemRef(aligned + offset, {size0, size1, size2},
-                           {stride0, stride1, stride2}, 123);
+extern "C" void mcuMemHostRegisterFloat(int64_t rank, void *ptr) {
+  auto *desc = static_cast<StridedMemRefType<float, 1> *>(ptr);
+  auto sizes = llvm::ArrayRef<int64_t>(desc->sizes, rank);
+  auto strides = llvm::ArrayRef<int64_t>(desc->sizes + rank, rank);
+  mcuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 1.23f);
+}
+
+extern "C" void mcuMemHostRegisterInt32(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);
+  mcuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 123);
 }