[MLIR] Provide bare pointer memref lowering option on gpu-to-nvvm pass

[bondhugula]

Summary

Provide the bare pointer memref lowering option on gpu-to-nvvm pass.
This is needed whenever we lower memrefs on the host function side and
the kernel calls on the host-side (gpu-to-llvm) with the bare ptr
convention. The GPU module side of the lowering should also "align" and
use the bare pointer convention.

[bondhugula]

Set the option more naturally.

[krzysz00]

LGTM

[krzysz00]

You might want to add a test though

[bondhugula]

Add integration test cases for bare pointer memref lowering for kernels as well
as host-side IR. Tests the options for both gpu-to-nvvm and gpu-to-llvm.

[bondhugula]

@mehdi_amini Here's the test case for using bare pointers on the host side as well.

[bondhugula]

Adjust comments and dealloc.

[mehdi_amini]

Can you add a unit-test please? Integration test are never replacement for unit-tests.

[bondhugula]

Add unit test case for memref lowering for gpu-to-nvvm.

[bondhugula]

In D152480#4430524, @mehdi_amini wrote:

Can you add a unit-test please? Integration test are never replacement for unit-tests.

Done - added unit test case. (The reason I thought it wasn't necessary was that memref bare ptr lowering already had unit test cases in func-to-llvm. The testing here is repeating/exercising the same lowering. But it's good to have it anyway to show all the users of the feature in the tree and to ensure the option is indeed working.)

mlir/include/mlir/Conversion/Passes.td

let dependentDialects = [
     "NVVM::NVVMDialect",
   ];
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
     Option<"hasRedux", "has-redux", "bool", /*default=*/"false",
            "Target gpu supports redux">,
+    Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
+           /*default=*/"false",
+           "Replace memref arguments in GPU functions with bare pointers. "
+           "All memrefs must have static shape.">,
     Option<"useOpaquePointers", "use-opaque-pointers", "bool",
                    /*default=*/"true", "Generate LLVM IR using opaque pointers "
                    "instead of typed pointers">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToROCDL

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

void runOnOperation() override {
 
     // Customize the bitwidth used for the device side index computations.
     LowerToLLVMOptions options(
         m.getContext(),
         DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
     options.useOpaquePointers = useOpaquePointers;
+    options.useBarePtrCallConv = useBarePtrCallConv;
 
     // Apply in-dialect lowering. In-dialect lowering will replace
     // ops which need to be lowered further, which is not supported by a
     // single conversion pass.
     {
       RewritePatternSet patterns(m.getContext());
       populateGpuRewritePatterns(patterns);
       if (failed(applyPatternsAndFoldGreedily(m, std::move(patterns))))

mlir/test/Conversion/GPUToNVVM/memref.mlir

+// RUN: mlir-opt %s -convert-gpu-to-nvvm | FileCheck %s
+// RUN: mlir-opt %s -convert-gpu-to-nvvm='use-bare-ptr-memref-call-conv=1' \
+// RUN: | FileCheck %s --check-prefix=BARE
+
+gpu.module @memref_conversions {
+  // CHECK: llvm.func @kern
+  // CHECK-SAME: (%{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64)
+  // BARE: llvm.func @kern
+  // BARE-SAME: (%{{.*}}: !llvm.ptr)
+  gpu.func @kern(%arg0: memref<8xf32>) kernel {
+    gpu.return
+  }
+}

mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f32-bare-ptr.mlir

+// Tests memref bare pointer lowering convention both host side and kernel-side;
+// this works for only statically shaped memrefs.
+// Similar to the wmma-matmul-f32 but but with the memref bare pointer lowering convention.
+// This test also uses gpu.memcpy operations (instead of gpu.host_register).
+// RUN: mlir-opt %s \
+// RUN: | mlir-opt -gpu-kernel-outlining \
+// RUN: | mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm{use-bare-ptr-memref-call-conv=1},gpu-to-cubin{chip=sm_70}))' \
+// RUN: | mlir-opt --convert-scf-to-cf -gpu-to-llvm="use-bare-pointers-for-host=1 use-bare-pointers-for-kernels=1" \
+// RUN: | mlir-cpu-runner \
+// RUN:   --shared-libs=%mlir_cuda_runtime \
+// RUN:   --entry-point-result=void \
+// RUN: | FileCheck %s
+
+func.func @main() {
+  %h0 = memref.alloc() : memref<16x16xf16>
+  %h_out = memref.alloc() : memref<16x16xf32>
+
+  %f1 = arith.constant 1.0e+00 : f16
+  %f0 = arith.constant 0.0e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %c32 = arith.constant 32 : index
+  %c1 = arith.constant 1 : index
+
+  // Intialize the Input matrix with ones.
+  scf.for %arg0 = %c0 to %c16 step %c1 {
+    scf.for %arg1 = %c0 to %c16 step %c1 {
+      memref.store %f1, %h0[%arg0, %arg1] : memref<16x16xf16>
+    }
+  }
+  // Intialize the accumulator matrix with zeros.
+  scf.for %arg0 = %c0 to %c16 step %c1 {
+    scf.for %arg1 = %c0 to %c16 step %c1 {
+      memref.store %f0, %h_out[%arg0, %arg1] : memref<16x16xf32>
+    }
+  }
+
+  %token = gpu.wait async
+  %0, %t0 = gpu.alloc async [%token] () : memref<16x16xf16>
+  %out, %t1 = gpu.alloc async [%token]() : memref<16x16xf32>
+  // Copy from host to device.
+  %x = gpu.memcpy async [%token] %0, %h0 : memref<16x16xf16>, memref<16x16xf16>
+  %y = gpu.memcpy async [%token] %out, %h_out : memref<16x16xf32>, memref<16x16xf32>
+
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
+             threads(%tx, %ty, %tz) in (%block_x = %c32, %block_y = %c1, %block_z = %c1) {
+    %A = gpu.subgroup_mma_load_matrix %0[%c0, %c0] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
+    %B = gpu.subgroup_mma_load_matrix %0[%c0, %c0] {leadDimension = 16 : index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "BOp">
+    %C = gpu.subgroup_mma_load_matrix %out[%c0, %c0] {leadDimension = 16 : index} : memref<16x16xf32> -> !gpu.mma_matrix<16x16xf32, "COp">
+
+    %R = gpu.subgroup_mma_compute %A, %B, %C : !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp"> -> !gpu.mma_matrix<16x16xf32, "COp">
+
+    gpu.subgroup_mma_store_matrix %R, %out[%c0, %c0] {leadDimension = 16 : index}: !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32>
+    // Memref print utilities use unranked memrefs. We can't use bare pointer lowering
+    // with them. We simply test for successful execution.
+    gpu.printf "Success\n"
+    // CHECK: Success
+    gpu.terminator
+  }
+  %z = gpu.dealloc async [%token] %0 : memref<16x16xf16>
+  %w = gpu.dealloc async [%token] %out : memref<16x16xf32>
+  gpu.wait [%token]
+  return
+}

mlir/test/Integration/GPU/CUDA/all-reduce-and.mlir

 // RUN: mlir-opt %s \
 // RUN: | mlir-opt -gpu-kernel-outlining \
 // RUN: | mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,gpu-to-cubin))' \
 // RUN: | mlir-opt -gpu-to-llvm \
 // RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%mlir_cuda_runtime \
 // RUN:   --shared-libs=%mlir_runner_utils \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
+// Same as above but with the memref bare pointer lowering convention.
+// RUN: mlir-opt %s \
+// RUN: | mlir-opt -gpu-kernel-outlining \
+// RUN: | mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm{use-bare-ptr-memref-call-conv=1},gpu-to-cubin))' \
+// RUN: | mlir-opt -gpu-to-llvm="use-bare-pointers-for-kernels=1" \
+// RUN: | mlir-cpu-runner \
+// RUN:   --shared-libs=%mlir_cuda_runtime \
+// RUN:   --shared-libs=%mlir_runner_utils \
+// RUN:   --entry-point-result=void \
+// RUN: | FileCheck %s
+
 func.func @main() {
   %data = memref.alloc() : memref<2x6xi32>
   %sum = memref.alloc() : memref<2xi32>
   %cst0 = arith.constant 0 : i32
   %cst1 = arith.constant 1 : i32
   %cst2 = arith.constant 2 : i32
   %cst4 = arith.constant 4 : i32
   %cst8 = arith.constant 8 : i32
func.func @main() {
 
   call @printMemrefI32(%cast_sum) : (memref<*xi32>) -> ()
   // CHECK: [0, 2]
 
   return
 }
 
 func.func private @printMemrefI32(memref<*xi32>)
-