[Mlir-commits] [mlir] [MLIR][NVVM] Update MLIR mapa to reflect new address space (PR #146031)

[Guray Ozen]

================
@@ -3009,10 +3009,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_MapaASCheck<string inputArg, string resultArg> : 
+    PredOpTrait<"Valid address-space check(or mapping) for mapa Op",
+                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>
+                  ]>
+                ]>>;
----------------
grypp wrote:

One can use it like following
```
def NVVM_MapaOp: NVVM_Op<"mapa",
    [NVVM_AddressCheckCheck<"res", [LLVM_PointerShared, LLVM_PointerSharedCluster]>, NVVMRequiresSM<90>]> {
```

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