make generic helper and update usage to it

Author: modiking

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -3009,28 +3009,11 @@ 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>
-                  ]>
-                ]>>;
-
-def NVVM_MapaOp: NVVM_Op<"mapa",
-    [NVVM_MapaASCheck<"a", "res">, NVVMRequiresSM<90>]> {
+def NVVM_MapaASCheck : PredOpTrait<"Valid address-space check(or mapping) for mapa Op",
+    Or<[InputMatchesTypes<["a", "res"], [LLVM_PointerShared, LLVM_PointerSharedCluster]>.predicate,
+        InputMatchesTypes<["a", "res"], [LLVM_PointerGeneric, LLVM_PointerGeneric]>.predicate]>>;
+
+def NVVM_MapaOp: NVVM_Op<"mapa", [NVVM_MapaASCheck, NVVMRequiresSM<90>]> {
   let results = (outs AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerSharedCluster]>:$res);
   let arguments = (ins AnyTypeOf<[LLVM_PointerGeneric, LLVM_PointerShared]>:$a, I32:$b);
 

mlir/include/mlir/IR/OpBase.td

@@ -603,6 +603,22 @@ class RangedTypesMatchWith<string summary, string lhsArg, string rhsArg,
                            string transform>
   : TypesMatchWith<summary, lhsArg, rhsArg, transform, "llvm::equal">;
 
+// Checks that each inputArg has the same type as the corresponding entry
+// in allowedTypes
+class InputMatchesTypes<list<string> inputArgs, list<Type> allowedTypes> :
+    PredOpTrait<"operands {" # !interleave(inputArgs, ", ") # "} match expected types",
+                !foldl(TruePred, !range(!size(inputArgs)), acc, i,
+                       And<[acc,
+                           SubstLeaves<"$_self", "$" # inputArgs[i] # ".getType()",
+                                      allowedTypes[i].predicate>
+                       ]>)> {
+    assert !eq(!size(inputArgs), !size(allowedTypes)),
+           "inputArgs and allowedTypes lists must have the same length";
+
+    list<string> inputArgList = inputArgs;
+    list<Type> allowedTypeList = allowedTypes;
+}
+
 // Type Constraint operand `idx`'s Element type is `type`.
 class TCopVTEtIs<int idx, Type type> : And<[
    CPred<"$_op.getNumOperands() > " # idx>,