Fix all-reduce int tests by host-registering memrefs. Reduce amount of boiler plate to register host memory.

[csigg]

Summary

Fix all-reduce int tests by host-registering memrefs.

[mehdi_amini]

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

97

FYI: clang-tidy flags this variable name

109

Why does the registration requires to overwrite the array?

[csigg]

Fix variable name (dense_strides > denseStrides)

[csigg]

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

97

Thanks. Fixed.

109

It's not required, and it would probably be good to remove (or have an op to allocate host-registered memory for tests). I don't want to change the functionality as part of this CL though.

[herhut]

Thanks for fixing this. The clang-tidy error seems benign.

[herhut]

@clementval FYI, these tests were failing.

[clementval]

In D76563#1936415, @herhut wrote:

@clementval FYI, these tests were failing.

Ok weird. There were passing on my system. Anyway thanks for fixing.

mlir/test/mlir-cuda-runner/all-reduce-and.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
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
-  %sum_and = alloc() : memref<2xi32>
-  %sum_or = alloc() : memref<2xi32>
-  %sum_min = alloc() : memref<2xi32>
+  %sum = alloc() : memref<2xi32>
   %cst0 = constant 0 : i32
   %cst1 = constant 1 : i32
   %cst2 = constant 2 : i32
   %cst4 = constant 4 : i32
   %cst8 = constant 8 : i32
   %cst16 = constant 16 : i32
 
   %cst3 = constant 3 : i32
   %cst6 = constant 6 : i32
   %cst7 = constant 7 : i32
   %cst10 = constant 10 : i32
   %cst11 = constant 11 : i32
 
   %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>) -> ()
+
   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_and = "gpu.all_reduce"(%val) ({}) { op = "and" } : (i32) -> (i32)
-    store %reduced_and, %sum_and[%bx] : memref<2xi32>
+    %reduced = "gpu.all_reduce"(%val) ({}) { op = "and" } : (i32) -> (i32)
+    store %reduced, %sum[%bx] : memref<2xi32>
     gpu.terminator
   }
 
-  %ptr_and = memref_cast %sum_and : memref<2xi32> to memref<*xi32>
-  call @print_memref_i32(%ptr_and) : (memref<*xi32>) -> ()
+  %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
+  call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
+func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
+func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?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>) -> ()
+
   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>
func @main() {
 
   %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
   call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
   // CHECK: [16, 11]
 
   return
 }
 
+func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
+func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?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>) -> ()
+
   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>
func @main() {
 
   %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
   call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
+func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
+func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
 func @print_memref_i32(memref<*xi32>)
 

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>) -> ()
+
   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>
func @main() {
 
   %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
   call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
   // CHECK: [31, 15]
 
   return
 }
 
+func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
+func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
 func @print_memref_i32(memref<*xi32>)
 

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>) -> ()
+
   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>
func @main() {
 
   %ptr = memref_cast %sum : memref<2xi32> to memref<*xi32>
   call @print_memref_i32(%ptr) : (memref<*xi32>) -> ()
   // CHECK: [31, 1]
 
   return
 }
 
+func @mcuMemHostRegisterMemRef1dInt32(%ptr : memref<?xi32>)
+func @mcuMemHostRegisterMemRef2dInt32(%ptr : memref<?x?xi32>)
 func @print_memref_i32(memref<*xi32>)
 

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>) -> ()
+
   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>
 
   store %cst2, %data[%c1, %c0] : memref<2x6xf32>
func @main() {
 
   %ptr_mul = memref_cast %mul : memref<2xf32> to memref<*xf32>
   call @print_memref_f32(%ptr_mul) : (memref<*xf32>) -> ()
   // CHECK: [0, 27720]
 
   return
 }
 
+func @mcuMemHostRegisterMemRef1dFloat(%ptr : memref<?xf32>)
+func @mcuMemHostRegisterMemRef2dFloat(%ptr : memref<?x?xf32>)
 func @print_memref_f32(memref<*xf32>)

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 "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";
template <typename T, int N> struct MemRefType {
   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, int N>
-void mcuMemHostRegisterMemRef(const MemRefType<T, N> *arg, T value) {
-  auto count = std::accumulate(arg->sizes, arg->sizes + N, 1,
-                               std::multiplies<int64_t>());
-  std::fill_n(arg->data, count, value);
-  mcuMemHostRegister(arg->data, count * sizeof(T));
+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> dense_strides(strides.size());
+
+  std::partial_sum(sizes.rbegin(), sizes.rend(), dense_strides.rbegin(),
+  auto count = dense_strides.front();
+
+  // Only densely backed tensors are currently supported.
+  std::rotate(dense_strides.begin(), dense_strides.begin() + 1,
+              dense_strides.end());
+  dense_strides.back() = 1;
+  assert(strides == llvm::makeArrayRef(dense_strides));
+
+  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) {
-  MemRefType<float, 1> descriptor;
-  descriptor.basePtr = allocated;
-  descriptor.data = aligned;
-  descriptor.offset = offset;
-  descriptor.sizes[0] = size;
-  descriptor.strides[0] = stride;
-  mcuMemHostRegisterMemRef(&descriptor, 1.23f);
+  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) {
-  MemRefType<float, 3> descriptor;
-  descriptor.basePtr = allocated;
-  descriptor.data = aligned;
-  descriptor.offset = offset;
-  descriptor.sizes[0] = size0;
-  descriptor.strides[0] = stride0;
-  descriptor.sizes[1] = size1;
-  descriptor.strides[1] = stride1;
-  descriptor.sizes[2] = size2;
-  descriptor.strides[2] = stride2;
-  mcuMemHostRegisterMemRef(&descriptor, 1.23f);
+  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);
 }