create InputAddressIsCombinationOf

Author: modiking

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

@@ -3015,11 +3015,10 @@ def NVVM_GriddepcontrolLaunchDependentsOp
 // NVVM Mapa Op
 //===----------------------------------------------------------------------===//
 
-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>]> {
+def NVVM_MapaOp: NVVM_Op<"mapa",
+  [InputAddressIsCombinationOf<["a", "res"],
+    [[LLVM_PointerShared, LLVM_PointerSharedCluster], [LLVM_PointerGeneric, LLVM_PointerGeneric]],
+    "Valid address-space check(or mapping) for mapa Op">, 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

@@ -619,6 +619,35 @@ class InputMatchesTypes<list<string> inputArgs, list<Type> allowedTypes> :
     list<Type> allowedTypeList = allowedTypes;
 }
 
+// Checks that inputArgs match one of the allowed type combinations.
+// Each combination in allowedCombinations must have the same number of types
+// as there are inputArgs.
+class InputAddressIsCombinationOf<list<string> inputArgs, 
+                                  list<list<Type>> allowedCombinations,
+                                  string description = ""> :
+    PredOpTrait<!if(!empty(description), 
+                    "operands {" # !interleave(inputArgs, ", ") # "} match one of the allowed type combinations",
+                    description),
+                Or<!foreach(combination, allowedCombinations,
+                           !foldl(TruePred, !range(!size(inputArgs)), acc, i,
+                                  And<[acc,
+                                      SubstLeaves<"$_self", "$" # inputArgs[i] # ".getType()",
+                                                 combination[i].predicate>
+                                  ]>))>> {
+    assert !gt(!size(allowedCombinations), 0),
+           "allowedCombinations must not be empty";
+    
+    // Validate that each combination has the same number of types as inputArgs
+    defvar inputArgSize = !size(inputArgs);
+    defvar validSizes = !foldl(1, allowedCombinations, acc, combination,
+                              !and(acc, !eq(inputArgSize, !size(combination))));
+    assert validSizes,
+           "each combination in allowedCombinations must have the same length as inputArgs";
+    
+    list<string> inputArgList = inputArgs;
+    list<list<Type>> allowedCombinationList = allowedCombinations;
+}
+
 // Type Constraint operand `idx`'s Element type is `type`.
 class TCopVTEtIs<int idx, Type type> : And<[
    CPred<"$_op.getNumOperands() > " # idx>,