[mlir] resolve types from attributes in assemblyFormat

[tali]

Summary

An operation can specify that an operation or result type matches the
type of another operation, result, or attribute via the AllTypesMatch
or TypesMatchWith constraints.

Use these constraints to also automatically resolve types in the
automatically generated assembly parser.
This way, only the attribute needs to be listed in assemblyFormat,
e.g. for constant operations.

[rriddle]

Sorry for the delay. Thanks for tackling this! Can you add a few tests?
https://github.com/llvm/llvm-project/blob/master/mlir/test/mlir-tblgen/op-format-spec.td
https://github.com/llvm/llvm-project/blob/master/mlir/test/mlir-tblgen/op-format.mlir
https://github.com/llvm/llvm-project/blob/6b3168f8cdb46656330929877b0b4daab35d30de/mlir/test/lib/Dialect/Test/TestOps.td#L1165

mlir/tools/mlir-tblgen/OpFormatGen.cpp

318

Can you use an llvm::PointerUnion for the variable and attribute to avoid another field?

1345

Same here, could you use PointerUnion instead?

1745

Is there a reason why we couldn't use the transformer with the attribute type?

[tali]

Use llvm::PointerUnion and add some tests

[tali]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1745

Haven't found a need to do so, but added it in next revision.

[rriddle]

Looking really good.

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

1259

Can you add an example for AllTypesMatch?

mlir/tools/mlir-tblgen/OpFormatGen.cpp

299

Please run clang-format.

1702

Can you add support for attributes here as well?

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

739

We will also need to add verification of the attribute here if the transformer is non-trivial.

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

739

Although I suspect 99% of the cases are already handled because we check the storage type when parsing, happy to leave this as a followup for now.

[rriddle]

LGTM after adding handling for AllTypesMatch and a test for it.

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

1259

With attributes I mean.

[tali]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

739

Which checks would you find useful here?

[tali]

Add attribute support for AllTypesMatch

[tali]

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

106

I would really like to be able to format this example as test.format_all_types_match_attr 1, %i64 : i64, with the type at the end.
We could add support for type($attribute), and suppress printing of the attribute type if it is specified explicitly or when it can be inferred.

[rriddle]

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

106

The type of the attribute is required to be resolved when parsing it(i.e. at creation time). You would need some fairly fundamental changes to attributes to support that.

mlir/tools/mlir-tblgen/OpFormatGen.cpp

739

I mean that if the transformer is dependent on specific invariants of the attribute constraint, the attribute has to be verified first otherwise the transformer could assert/crash. This is the same reasoning as to why the operand/result variables get verified here.

[tali]

rebase against HEAD

[tali]

Rebase to HEAD

[tali]

Rename NamedArgOrAttr to ConstArgument

[mehdi_amini]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1766

Can you name the types instead of auto? (when the type isn't immediately obvious and/or it does not improve readability significantly)

[tali]

Used more concrete types instead of auto and simplified findSeenArg().

[tali]

Rebase on top of HEAD.

mlir/docs/OpDefinitions.md

 inheriting from the `BuildableType` class.
 
 * Trait Equality Constraints
 
 There are many operations that have known type equality constraints registered
 as traits on the operation; for example the true, false, and result values of a
 `select` operation often have the same type. The assembly format may inspect
 these equal constraints to discern the types of missing variables. The currently
-supported traits are: `AllTypesMatch`, `SameTypeOperands`, and
-`SameOperandsAndResultType`.
+supported traits are: `AllTypesMatch`, `TypesMatchWith`, `SameTypeOperands`,
+and `SameOperandsAndResultType`.
 
 ### `hasCanonicalizer`
 
 This boolean field indicate whether canonicalization patterns have been defined
 for this operation. If it is `1`, then `::getCanonicalizationPatterns()` should
 be defined.
 
 ### `hasFolder`

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

 }]>;
 
 def FormatOptionalOperandResultBOp : FormatOptionalOperandResultOpBase<"b", [{
   (`(` $optional^ `:` type($optional) `)`)? `:` type($optional_res)
   (`[` $variadic^ `]`)? attr-dict
 }]>;
 
 //===----------------------------------------------------------------------===//
+// AllTypesMatch type inference
+//===----------------------------------------------------------------------===//
+
+def FormatAllTypesMatchVarOp : TEST_Op<"format_all_types_match_var", [
+    AllTypesMatch<["value1", "value2", "result"]>
+  ]> {
+  let arguments = (ins AnyType:$value1, AnyType:$value2);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "attr-dict $value1 `,` $value2 `:` type($value1)";
+}
+
+def FormatAllTypesMatchAttrOp : TEST_Op<"format_all_types_match_attr", [
+    AllTypesMatch<["value1", "value2", "result"]>
+  ]> {
+  let arguments = (ins AnyAttr:$value1, AnyType:$value2);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "attr-dict $value1 `,` $value2";
+}
+
+//===----------------------------------------------------------------------===//
+// TypesMatchWith type inference
+//===----------------------------------------------------------------------===//
+
+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 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";
+}
+
+//===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//
 
 def SideEffectOp : TEST_Op<"side_effect_op",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let results = (outs AnyType:$result);
 }
 
 #endif // TEST_OPS

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

 // CHECK: test.format_optional_operand_result_b_op(%[[I64]] : i64) : i64
 test.format_optional_operand_result_b_op(%i64 : i64) : i64
 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op( : ) : i64
 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op : i64
+
+//===----------------------------------------------------------------------===//
+// AllTypesMatch type inference
+//===----------------------------------------------------------------------===//
+
+// CHECK: test.format_all_types_match_var %[[I64]], %[[I64]] : i64
+%ignored_res1 = test.format_all_types_match_var %i64, %i64 : i64
+
+// CHECK: test.format_all_types_match_attr 1 : i64, %[[I64]]
+%ignored_res2 = test.format_all_types_match_attr 1 : i64, %i64
+
+//===----------------------------------------------------------------------===//
+// 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_attr 1 : i64
+%ignored_res4 = test.format_types_match_attr 1 : i64

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // OperationFormat
 //===----------------------------------------------------------------------===//
 
 namespace {
+
+using NamedArgOrAttr =
+    llvm::PointerUnion<const NamedTypeConstraint *, const NamedAttribute *>;
+
 struct OperationFormat {
   /// This class represents a specific resolver for an operand or result type.
   class TypeResolution {
   public:
     TypeResolution() = default;
 
     /// Get the index into the buildable types for this type, or None.
     Optional<int> getBuilderIdx() const { return builderIdx; }
     void setBuilderIdx(int idx) { builderIdx = idx; }
 
-    /// Get the variable this type is resolved to, or None.
-    const NamedTypeConstraint *getVariable() const { return variable; }
+    /// Get the variable this type is resolved to, or nullptr.
+    const NamedTypeConstraint *getVariable() const {
+      return resolver.dyn_cast<const NamedTypeConstraint *>();
+    }
+    /// Get the attribute this type is resolved to, or nullptr.
+    const NamedAttribute *getAttribute() const {
+      return resolver.dyn_cast<const NamedAttribute *>();
+    }
+    /// Get the transformer for the type of the variable, or None.
     Optional<StringRef> getVarTransformer() const {
       return variableTransformer;
     }
-    void setVariable(const NamedTypeConstraint *var,
+    void setResolver(NamedArgOrAttr varOrAttr,
                      Optional<StringRef> transformer) {
-      variable = var;
+      resolver = varOrAttr;
       variableTransformer = transformer;
+      assert(getVariable() || getAttribute());
     }
 
   private:
     /// If the type is resolved with a buildable type, this is the index into
     /// 'buildableTypes' in the parent format.
     Optional<int> builderIdx;
     /// If the type is resolved based upon another operand or result, this is
-    /// the variable that this type is resolved to.
-    const NamedTypeConstraint *variable;
+    /// the variable or the attribute that this type is resolved to.
+    NamedArgOrAttr resolver;
     /// If the type is resolved based upon another operand or result, this is
     /// a transformer to apply to the variable when resolving.
     Optional<StringRef> variableTransformer;
   };
 
   OperationFormat(const Operator &op)
       : allOperands(false), allOperandTypes(false), allResultTypes(false) {
     operandTypes.resize(op.getNumOperands(), TypeResolution());
     resultTypes.resize(op.getNumResults(), TypeResolution());
   }
 
   /// Generate the operation parser from this format.
   void genParser(Operator &op, OpClass &opClass);
   /// Generate the c++ to resolve the types of operands and results during
void OperationFormat::genParserTypeResolution(Operator &op,
   // types of those variables are resolved.
   SmallPtrSet<const NamedTypeConstraint *, 8> verifiedVariables;
   FmtContext verifierFCtx;
   for (TypeResolution &resolver :
        llvm::concat<TypeResolution>(resultTypes, operandTypes)) {
     Optional<StringRef> transformer = resolver.getVarTransformer();
     if (!transformer)
       continue;
     // Ensure that we don't verify the same variables twice.
     const NamedTypeConstraint *variable = resolver.getVariable();
-    if (!verifiedVariables.insert(variable).second)
+    if (!variable || !verifiedVariables.insert(variable).second)
       continue;
 
     auto constraint = variable->constraint;
     body << "  for (Type type : " << variable->name << "Types) {\n"
          << "    (void)type;\n"
          << "    if (!("
          << tgfmt(constraint.getConditionTemplate(),
                   &verifierFCtx.withSelf("type"))
void OperationFormat::genParserTypeResolution(Operator &op,
   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]"));
       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";
     }
   };
 
   // Resolve each of the result types.
   if (allResultTypes) {
     body << "  result.addTypes(allResultTypes);\n";

 }
 
 //===----------------------------------------------------------------------===//
 // FormatParser
 //===----------------------------------------------------------------------===//
 
 /// Function to find an element within the given range that has the same name as
 /// 'name'.
-template <typename RangeT> static auto findArg(RangeT &&range, StringRef name) {
+template <typename RangeT>
+static auto findArg(RangeT &&range, StringRef name) {
   auto it = llvm::find_if(range, [=](auto &arg) { return arg.name == name; });
   return it != range.end() ? &*it : nullptr;
 }
 
 namespace {
 /// This class implements a parser for an instance of an operation assembly
 /// format.
 class FormatParser {
 public:
   FormatParser(llvm::SourceMgr &mgr, OperationFormat &format, Operator &op)
       : lexer(mgr, op), curToken(lexer.lexToken()), fmt(format), op(op),
         seenOperandTypes(op.getNumOperands()),
         seenResultTypes(op.getNumResults()) {}
 
   /// Parse the operation assembly format.
   LogicalResult parse();
 
 private:
   /// This struct represents a type resolution instance. It includes a specific
   /// type as well as an optional transformer to apply to that type in order to
   /// properly resolve the type of a variable.
   struct TypeResolutionInstance {
-    const NamedTypeConstraint *type;
+    NamedArgOrAttr resolver;
     Optional<StringRef> transformer;
   };
 
   /// An iterator over the elements of a format group.
   using ElementsIterT = llvm::pointee_iterator<
       std::vector<std::unique_ptr<Element>>::const_iterator>;
 
   /// Verify the state of operation attributes within the format.
   LogicalResult verifyAttributes(llvm::SMLoc loc);
void handleAllTypesMatchConstraint(
       ArrayRef<StringRef> values,
       llvm::StringMap<TypeResolutionInstance> &variableTyResolver);
   /// Check for inferable type resolution given all operands, and or results,
   /// have the same type. If 'includeResults' is true, the results also have the
   /// same type as all of the operands.
   void handleSameTypesConstraint(
       llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
       bool includeResults);
+  /// Check for inferable type resolution based on another operand, result, or
+  /// attribute.
+  void handleTypesMatchConstraint(
+      llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
+      llvm::Record def);
 
   /// Returns an argument with the given name that has been seen within the
   /// format.
   const NamedTypeConstraint *findSeenArg(StringRef name);
+  /// Returns an attribute with the given name that has been seen within the
+  /// format.
+  const NamedAttribute *findSeenAttr(StringRef name);
+  /// Returns an argument or attribute with the given name that has been seen
+  /// within the format.
+  NamedArgOrAttr findSeenArgOrAttr(StringRef name);
 
   /// Parse a specific element.
   LogicalResult parseElement(std::unique_ptr<Element> &element,
                              bool isTopLevel);
   LogicalResult parseVariable(std::unique_ptr<Element> &element,
                               bool isTopLevel);
   LogicalResult parseDirective(std::unique_ptr<Element> &element,
                                bool isTopLevel);
for (const OpTrait &trait : op.getTraits()) {
     if (def.isSubClassOf("AllTypesMatch")) {
       handleAllTypesMatchConstraint(def.getValueAsListOfStrings("values"),
                                     variableTyResolver);
     } else if (def.getName() == "SameTypeOperands") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/false);
     } else if (def.getName() == "SameOperandsAndResultType") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/true);
     } else if (def.isSubClassOf("TypesMatchWith")) {
-      if (const auto *lhsArg = findSeenArg(def.getValueAsString("lhs")))
-        variableTyResolver[def.getValueAsString("rhs")] = {
-            lhsArg, def.getValueAsString("transformer")};
+      handleTypesMatchConstraint(variableTyResolver, def);
     }
   }
 
   // Verify the state of the various operation components.
   if (failed(verifyAttributes(loc)) ||
       failed(verifyResults(loc, variableTyResolver)) ||
       failed(verifyOperands(loc, variableTyResolver)) ||
       failed(verifySuccessors(loc)))
for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
 
     // Check that the operand type is in the format, or that it can be inferred.
     if (fmt.allOperandTypes || seenOperandTypes.test(i))
       continue;
 
     // Check to see if we can infer this type from another variable.
     auto varResolverIt = variableTyResolver.find(op.getOperand(i).name);
     if (varResolverIt != variableTyResolver.end()) {
-      fmt.operandTypes[i].setVariable(varResolverIt->second.type,
-                                      varResolverIt->second.transformer);
+      auto resolver = varResolverIt->second;
+      fmt.operandTypes[i].setResolver(resolver.resolver, resolver.transformer);
       continue;
     }
 
     // Similarly to results, allow a custom builder for resolving the type if
     // we aren't using the 'operands' directive.
     Optional<StringRef> builder = operand.constraint.getBuilderCall();
     if (!builder || (hasAllOperands && operand.isVariableLength())) {
       return emitErrorAndNote(
LogicalResult FormatParser::verifyResults(
   auto &buildableTypes = fmt.buildableTypes;
   for (unsigned i = 0, e = op.getNumResults(); i != e; ++i) {
     if (seenResultTypes.test(i))
       continue;
 
     // Check to see if we can infer this type from another variable.
     auto varResolverIt = variableTyResolver.find(op.getResultName(i));
     if (varResolverIt != variableTyResolver.end()) {
-      fmt.resultTypes[i].setVariable(varResolverIt->second.type,
-                                     varResolverIt->second.transformer);
+      auto resolver = varResolverIt->second;
+      fmt.resultTypes[i].setResolver(resolver.resolver, resolver.transformer);
       continue;
     }
 
     // If the result is not variable length, allow for the case where the type
     // has a builder that we can use.
     NamedTypeConstraint &result = op.getResult(i);
     Optional<StringRef> builder = result.constraint.getBuilderCall();
     if (!builder || result.isVariableLength()) {
LogicalResult FormatParser::verifySuccessors(llvm::SMLoc loc) {
   return success();
 }
 
 void FormatParser::handleAllTypesMatchConstraint(
     ArrayRef<StringRef> values,
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver) {
   for (unsigned i = 0, e = values.size(); i != e; ++i) {
     // Check to see if this value matches a resolved operand or result type.
-    const NamedTypeConstraint *arg = findSeenArg(values[i]);
+    NamedArgOrAttr arg = findSeenArgOrAttr(values[i]);
     if (!arg)
       continue;
 
     // Mark this value as the type resolver for the other variables.
     for (unsigned j = 0; j != i; ++j)
       variableTyResolver[values[j]] = {arg, llvm::None};
     for (unsigned j = i + 1; j != e; ++j)
       variableTyResolver[values[j]] = {arg, llvm::None};
if (!seenOperandTypes.test(i) && !op.getOperand(i).name.empty())
       variableTyResolver[op.getOperand(i).name] = {resolver, llvm::None};
   if (includeResults) {
     for (unsigned i = 0, e = op.getNumResults(); i != e; ++i)
       if (!seenResultTypes.test(i) && !op.getResultName(i).empty())
         variableTyResolver[op.getResultName(i)] = {resolver, llvm::None};
   }
 }
 
+void FormatParser::handleTypesMatchConstraint(
+    llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
+    llvm::Record def) {
+  auto lhsName = def.getValueAsString("lhs");
+  auto rhsName = def.getValueAsString("rhs");
+  auto transformer = def.getValueAsString("transformer");
+  if (auto var = findSeenArgOrAttr(lhsName))
+    variableTyResolver[rhsName] = {var, transformer};
+}
+
 const NamedTypeConstraint *FormatParser::findSeenArg(StringRef name) {
   if (auto *arg = findArg(op.getOperands(), name))
     return seenOperandTypes.test(arg - op.operand_begin()) ? arg : nullptr;
   if (auto *arg = findArg(op.getResults(), name))
     return seenResultTypes.test(arg - op.result_begin()) ? arg : nullptr;
   return nullptr;
 }
 
+const NamedAttribute *FormatParser::findSeenAttr(StringRef name) {
+  if (const auto *attr = findArg(op.getAttributes(), name))
+    return seenAttrs.find_as(attr) != seenAttrs.end() ? attr : nullptr;
+  return nullptr;
+}
+
+NamedArgOrAttr FormatParser::findSeenArgOrAttr(StringRef name) {
+  if (auto *arg = findSeenArg(name))
+    return arg;
+  if (auto *attr = findSeenAttr(name))
+    return attr;
+  return nullptr;
+}
+
 LogicalResult FormatParser::parseElement(std::unique_ptr<Element> &element,
                                          bool isTopLevel) {
   // Directives.
   if (curToken.isKeyword())
     return parseDirective(element, isTopLevel);
   // Literals.
   if (curToken.getKind() == Token::literal)
     return parseLiteral(element);