[mlir][ODS] Add new RangedTypesMatchWith operation predicate

[rriddle]

Summary

This is a variant of TypesMatchWith that provides support for variadic arguments. This is necessary because ranges generally can't use the default operator== comparators for checking equality.

[mehdi_amini]

mlir/include/mlir/IR/OpBase.td

2195

Missing closing backquote?

mlir/include/mlir/IR/OpBase.td

 
 class AllShapesMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, Shape<"_self">.result, "shape">;
 
 class AllTypesMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, "$_self.getType()", "type">;
 
 // A type constraint that denotes `transform(lhs.getType()) == rhs.getType()`.
+// An optional comparator function may be provided that changes the above form
+// into: `comparator(transform(lhs.getType()), rhs.getType())
 class TypesMatchWith<string summary, string lhsArg, string rhsArg,
-                     string transform> :
-    PredOpTrait<summary, CPred<
-      !subst("$_self", "$" # lhsArg # ".getType()", transform)
-      # " == $" # rhsArg # ".getType()">> {
+                     string transform, string comparator = "std::equal_to<>()">
+  : PredOpTrait<summary, CPred<
+      comparator # "(" #
+      !subst("$_self", "$" # lhsArg # ".getType()", transform) #
+      ", $" # rhsArg # ".getType())">> {
   string lhs = lhsArg;
   string rhs = rhsArg;
   string transformer = transform;
 }
 
+// Special variant of `TypesMatchWith` that provides a comparator suitable for
+// ranged arguments.
+class RangedTypesMatchWith<string summary, string lhsArg, string rhsArg,
+                           string transform>
+  : TypesMatchWith<summary, lhsArg, rhsArg, transform,
+    "[](auto &&lhs, auto &&rhs) { "
+    "return std::equal(lhs.begin(), lhs.end(), rhs.begin(), rhs.end());"
+    " }">;
+
 // Type Constraint operand `idx`'s Element type is `type`.
 class TCopVTEtIs<int idx, Type type> : And<[
    CPred<"$_op.getNumOperands() > " # idx>,
    SubstLeaves<"$_self", "$_op.getOperand(" # idx # ").getType()",
      IsShapedTypePred>,
    SubstLeaves<"$_self", "getElementTypeOrSelf($_op.getOperand(" # idx # "))",
      type.predicate>]>;
 

mlir/test/lib/Dialect/Test/TestOps.td

 def FormatTypesMatchVarOp : TEST_Op<"format_types_match_var", [
     TypesMatchWith<"result type matches operand", "value", "result", "$_self">
   ]> {
   let arguments = (ins AnyType:$value);
   let results = (outs AnyType:$result);
   let assemblyFormat = "attr-dict $value `:` type($value)";
 }
 
+def FormatTypesMatchVariadicOp : TEST_Op<"format_types_match_variadic", [
+    RangedTypesMatchWith<"result type matches operand", "value", "result",
+                         "llvm::make_range($_self.begin(), $_self.end())">
+  ]> {
+  let arguments = (ins Variadic<AnyType>:$value);
+  let results = (outs Variadic<AnyType>:$result);
+  let assemblyFormat = "attr-dict $value `:` type($value)";
+}
+
 def FormatTypesMatchAttrOp : TEST_Op<"format_types_match_attr", [
     TypesMatchWith<"result type matches constant", "value", "result", "$_self">
   ]> {
   let arguments = (ins AnyAttr:$value);
   let results = (outs AnyType:$result);
   let assemblyFormat = "attr-dict $value";
 }
 

mlir/test/mlir-tblgen/op-format.mlir

 
 //===----------------------------------------------------------------------===//
 // TypesMatchWith type inference
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_types_match_var %[[I64]] : i64
 %ignored_res3 = test.format_types_match_var %i64 : i64
 
+// CHECK: test.format_types_match_variadic %[[I64]], %[[I64]], %[[I64]] : i64, i64, i64
+%ignored_res4:3 = test.format_types_match_variadic %i64, %i64, %i64 : i64, i64, i64
+
 // CHECK: test.format_types_match_attr 1 : i64
-%ignored_res4 = test.format_types_match_attr 1 : i64
+%ignored_res5 = test.format_types_match_attr 1 : i64

mlir/tools/mlir-tblgen/OpFormatGen.cpp

for (auto &it : buildableTypes)
            << tgfmt(it.first, &typeBuilderCtx) << ";\n";
   }
 
   // Emit the code necessary for a type resolver.
   auto emitTypeResolver = [&](TypeResolution &resolver, StringRef curVar) {
     if (Optional<int> val = resolver.getBuilderIdx()) {
       body << "odsBuildableType" << *val;
     } else if (const NamedTypeConstraint *var = resolver.getVariable()) {
-      if (Optional<StringRef> tform = resolver.getVarTransformer())
-        body << tgfmt(*tform, &FmtContext().withSelf(var->name + "Types[0]"));
+      if (Optional<StringRef> tform = resolver.getVarTransformer()) {
+        FmtContext fmtContext;
+        if (var->isVariadic())
+          fmtContext.withSelf(var->name + "Types");
         else
+          fmtContext.withSelf(var->name + "Types[0]");
+        body << tgfmt(*tform, &fmtContext);
+      } else {
         body << var->name << "Types";
+      }
     } else if (const NamedAttribute *attr = resolver.getAttribute()) {
       if (Optional<StringRef> tform = resolver.getVarTransformer())
         body << tgfmt(*tform,
                       &FmtContext().withSelf(attr->name + "Attr.getType()"));
       else
         body << attr->name << "Attr.getType()";
     } else {
       body << curVar << "Types";