[MLIR][NVVM] Update MLIR mapa to reflect new address space

[modiking]

The mapa.shared.cluster variant that takes in addrspace 3 now should output addrspace 7. Update this the td file to reflect this.

Testing:
Github CI

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: None (modiking)

Changes

The mapa.shared.cluster variant that takes in addrspace 3 now should output addrspace 7. Update this the td file to reflect this.

Testing:
Github CI

---

Full diff: https://github.com/llvm/llvm-project/pull/146031.diff

4 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+22-3)
• (modified) mlir/test/Dialect/LLVMIR/invalid.mlir (+2-2)
• (modified) mlir/test/Dialect/LLVMIR/nvvm.mlir (+1-1)
• (modified) mlir/test/Target/LLVMIR/nvvmir.mlir (+2-2)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 0a6e66919f021..c9d7e2833fcd3 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2600,10 +2600,29 @@ def NVVM_GriddepcontrolLaunchDependentsOp
 // NVVM Mapa Op
 //===----------------------------------------------------------------------===//
 
+// Helper predicates for address space checking
+def IsGenericAddressSpace : CPred<"llvm::cast<LLVM::LLVMPointerType>($_self).getAddressSpace() == 0">;
+def IsSharedAddressSpace : CPred<"llvm::cast<LLVM::LLVMPointerType>($_self).getAddressSpace() == 3">;
+def IsSharedClusterAddressSpace : CPred<"llvm::cast<LLVM::LLVMPointerType>($_self).getAddressSpace() == 7">;
+
+class NVVM_AddressSpaceMapping<string inputArg, string resultArg> : 
+    PredOpTrait<"valid address space mapping for NVVM mapa operation",
+                Or<[
+                  // Generic -> Generic
+                  And<[
+                    SubstLeaves<"$_self", "$" # inputArg # ".getType()", IsGenericAddressSpace>,
+                    SubstLeaves<"$_self", "$" # resultArg # ".getType()", IsGenericAddressSpace>
+                  ]>,
+                  // Shared -> SharedCluster
+                  And<[
+                    SubstLeaves<"$_self", "$" # inputArg # ".getType()", IsSharedAddressSpace>,
+                    SubstLeaves<"$_self", "$" # resultArg # ".getType()", IsSharedClusterAddressSpace>
+                  ]>
+                ]>>;
+
 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);
+    [NVVM_AddressSpaceMapping<"a", "res">, NVVMRequiresSM<90>]> {
+  let results = (outs AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerSharedCluster]>:$res);
   let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$a, I32:$b);
 
   string llvmBuilder = [{
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index 0cd6b1f20a1bf..f239816bcbad9 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -1201,8 +1201,8 @@ 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>
+  // expected-error @below {{'nvvm.mapa' op failed to verify that valid address space mapping for NVVM mapa operation}}
+  %0 = nvvm.mapa %a, %b: !llvm.ptr -> !llvm.ptr<7>
   return
 }
 
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index d3915492c38a0..46e3f1e5200fd 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -552,7 +552,7 @@ 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>
+  %1 = nvvm.mapa %a_shared, %b: !llvm.ptr<3> -> !llvm.ptr<7>
   return
 }
 
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 3a0713f2feee8..0cfcd8a4ef18d 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -769,8 +769,8 @@ llvm.func @nvvm_griddepcontrol_launch_dependents() {
 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>
+  // CHECK-LLVM: call ptr addrspace(7) @llvm.nvvm.mapa.shared.cluster(ptr addrspace(3) %{{.*}}, i32 %{{.*}})
+  %1 = nvvm.mapa %a_shared, %b: !llvm.ptr<3> -> !llvm.ptr<7>
   llvm.return
 }
 

[durga4github]

Let us have Mapa in this class-name so that it is clear that this does the checks for the Mapa Op.
Some options below are:
NVVM_MapaASCheck<>
NVVM_MapaASMappingCheck<>
etc..

[durga4github]

nit: Valid address-space check(or mapping) for mapa Op
(looking at how the error message appears in the negative test below)