[Mlir-commits] [mlir] [MLIR][NVVM] Add support for mapa MLIR Ops (PR #124514)

[Srinivasa Ravi]

https://github.com/Wolfram70 updated https://github.com/llvm/llvm-project/pull/124514

>From 5b97941d1c99532402fdbaa2819fd861b805e241 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Thu, 23 Jan 2025 17:02:11 +0530
Subject: [PATCH] [MLIR][NVVM] Add support for mapa MLIR Ops

Adds `mapa` and `mapa.shared.cluster` MLIR Ops to generate mapa
instructions.

`mapa` - Map the address of the shared variable in the target CTA.

- `mapa` - source is a register containing generic address pointing to
  shared memory.
- `mapa.shared.cluster` - source is a shared memory variable or a
  register containing a valid shared memory address.
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 25 +++++++++++++++++++++
 mlir/test/Dialect/LLVMIR/invalid.mlir       |  9 ++++++++
 mlir/test/Dialect/LLVMIR/nvvm.mlir          |  9 +++++++-
 mlir/test/Target/LLVMIR/nvvmir.mlir         | 10 +++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 8c8e44a054a627..d5bafb52b38830 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2512,6 +2512,31 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM Mapa Op
+//===----------------------------------------------------------------------===//
+
+def NVVM_MapaOp: NVVM_Op<"mapa",
+    [TypesMatchWith<"`res` and `a` should have the same type",
+                    "a", "res", "$_self">]> {
+  let results = (outs AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$res);
+  let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$a, I32:$b);
+
+  string llvmBuilder = [{
+    int addrSpace = llvm::cast<LLVMPointerType>(op.getA().getType()).getAddressSpace();
+    
+    bool isSharedMemory = addrSpace == NVVM::NVVMMemorySpace::kSharedMemorySpace;
+
+    if(isSharedMemory)
+      $res = createIntrinsicCall(builder,
+        llvm::Intrinsic::nvvm_mapa_shared_cluster, {$a, $b});
+    else
+      $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_mapa, {$a, $b});
+  }];
+  
+  let assemblyFormat = "$a`,` $b attr-dict `:` type($a) `->` type($res)";
+}
+
 def NVVM_Exit : NVVM_Op<"exit"> {
   let summary = "Exit Op";
   let description = [{
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index 25806d9d0edd72..dc7624b0f87bd5 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -1189,6 +1189,15 @@ func.func @cp_async(%arg0: !llvm.ptr<3>, %arg1: !llvm.ptr<1>) {
 
 // -----
 
+func.func @mapa(%a: !llvm.ptr, %b : i32) {
+  // expected-error @below {{`res` and `a` should have the same type}}
+  %0 = nvvm.mapa %a, %b: !llvm.ptr -> !llvm.ptr<3>
+  return
+}
+// -----
+
+// -----
+
 func.func @gep_struct_variable(%arg0: !llvm.ptr, %arg1: i32, %arg2: i32) {
   // expected-error @below {{op expected index 1 indexing a struct to be constant}}
   llvm.getelementptr %arg0[%arg1, %arg1] : (!llvm.ptr, i32, i32) -> !llvm.ptr, !llvm.struct<(i32)>
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 4c3b6648a41c00..20bf74d19aea41 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -509,7 +509,14 @@ func.func @wgmma_wait_group_sync_aligned() {
   return
 }
 
-// -----
+// CHECK-LABEL: @mapa
+func.func @mapa(%a: !llvm.ptr, %a_shared: !llvm.ptr<3>, %b : i32) {
+  // CHECK:   nvvm.mapa %{{.*}}
+  %0 = nvvm.mapa %a, %b: !llvm.ptr -> !llvm.ptr
+  // CHECK:   nvvm.mapa %{{.*}}
+  %1 = nvvm.mapa %a_shared, %b: !llvm.ptr<3> -> !llvm.ptr<3>
+  return
+}
 
 // Just check these don't emit errors.
 gpu.module @module_1 [#nvvm.target<chip = "sm_90", features = "+ptx70", link = ["my_device_lib.bc"], flags = {fast, ftz}>] {
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 7dad9a403def0e..19332af78ec11f 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -757,3 +757,13 @@ llvm.func @nvvm_wgmma_wait_group_aligned() {
   nvvm.wgmma.wait.group.sync.aligned 20
   llvm.return
 }
+
+// -----
+// CHECK-LABEL: @nvvm_mapa
+llvm.func @nvvm_mapa(%a: !llvm.ptr, %a_shared: !llvm.ptr<3>, %b : i32) {
+  // CHECK-LLVM: call ptr @llvm.nvvm.mapa(ptr %{{.*}}, i32 %{{.*}})
+  %0 = nvvm.mapa %a, %b: !llvm.ptr -> !llvm.ptr
+  // CHECK-LLVM: call ptr addrspace(3) @llvm.nvvm.mapa.shared.cluster(ptr addrspace(3) %{{.*}}, i32 %{{.*}})
+  %1 = nvvm.mapa %a_shared, %b: !llvm.ptr<3> -> !llvm.ptr<3>
+  llvm.return
+}