[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

314

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

1339

Same here, could you use PointerUnion instead?

1742

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

1742

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

1250 ↗

(On Diff #259722)

Can you add an example for AllTypesMatch?

mlir/tools/mlir-tblgen/OpFormatGen.cpp

288

Please run clang-format.

1699

Can you add support for attributes here as well?

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

738

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

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

738

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

1250 ↗

(On Diff #259722)

With attributes I mean.

[tali]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

738

Which checks would you find useful here?

[tali]

Add attribute support for AllTypesMatch

[tali]

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

106 ↗

(On Diff #261620)

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 ↗

(On Diff #261620)

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

738

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

1764

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/tools/mlir-tblgen/OpFormatGen.cpp

 //===- OpFormatGen.cpp - MLIR operation asm format generator --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
public:
     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; }
     Optional<StringRef> getVarTransformer() const {
       return variableTransformer;
     }
     void setVariable(const NamedTypeConstraint *var,
                      Optional<StringRef> transformer) {
       variable = var;
       variableTransformer = transformer;
+      attribute = nullptr;
+    }
+    /// Get the attribute this type is resolved to, or None.
+    const NamedAttribute *getAttribute() const { return attribute; }
+    void setAttribute(const NamedAttribute *attr) {
+      variable = nullptr;
+      variableTransformer = llvm::None;
+      attribute = attr;
     }
 
   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;
     /// 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;
+    /// If the type is resolved based upon another attribute, this is
+    /// the attribute that this type is resolved to.
+    const NamedAttribute *attribute;
   };
 
   OperationFormat(const Operator &op)
       : allOperands(false), allOperandTypes(false), allResultTypes(false) {
     operandTypes.resize(op.getNumOperands(), TypeResolution());
     resultTypes.resize(op.getNumResults(), TypeResolution());
   }
 
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)
       continue;
 
     auto constraint = variable->constraint;
     body << "  for (Type type : " << variable->name << "Types) {\n"
          << "    (void)type;\n"
          << "    if (!("
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()) {
+      body << attr->name << "Attr.getType()";
     } else {
       body << curVar << "Types";
     }
   };
 
   // Resolve each of the result types.
   if (allResultTypes) {
     body << "  result.addTypes(allResultTypes);\n";
public:
   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;
+    const NamedAttribute *attr;
     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.
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);
 
   /// 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;
+      if (resolver.type)
+        fmt.operandTypes[i].setVariable(resolver.type, resolver.transformer);
+      else if (resolver.attr)
+        fmt.operandTypes[i].setAttribute(resolver.attr);
       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;
+      if (resolver.type)
+        fmt.resultTypes[i].setVariable(resolver.type, resolver.transformer);
+      else if (resolver.attr)
+        fmt.resultTypes[i].setAttribute(resolver.attr);
       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]);
     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};
+      variableTyResolver[values[j]] = {arg, nullptr, llvm::None};
     for (unsigned j = i + 1; j != e; ++j)
-      variableTyResolver[values[j]] = {arg, llvm::None};
+      variableTyResolver[values[j]] = {arg, nullptr, llvm::None};
   }
 }
 
 void FormatParser::handleSameTypesConstraint(
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
     bool includeResults) {
   const NamedTypeConstraint *resolver = nullptr;
   int resolvedIt = -1;
 
   // Check to see if there is an operand or result to use for the resolution.
   if ((resolvedIt = seenOperandTypes.find_first()) != -1)
     resolver = &op.getOperand(resolvedIt);
   else if (includeResults && (resolvedIt = seenResultTypes.find_first()) != -1)
     resolver = &op.getResult(resolvedIt);
   else
     return;
 
   // Set the resolvers for each operand and result.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i)
     if (!seenOperandTypes.test(i) && !op.getOperand(i).name.empty())
-      variableTyResolver[op.getOperand(i).name] = {resolver, llvm::None};
+      variableTyResolver[op.getOperand(i).name] = {resolver, nullptr, 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};
+        variableTyResolver[op.getResultName(i)] = {resolver, nullptr, llvm::None};
   }
 }
 
+ void FormatParser::handleTypesMatchConstraint(
+     llvm::StringMap<TypeResolutionInstance> &variableTyResolver, llvm::Record def) {
+  auto lhsName = def.getValueAsString("lhs");
+  auto rhsName = def.getValueAsString("rhs");
+  if (const auto *lhsArg = findSeenArg(lhsName))
+    variableTyResolver[rhsName] = {
+        lhsArg, nullptr, def.getValueAsString("transformer")};
+  else if (const auto *attr = findSeenAttr(lhsName))
+    variableTyResolver[rhsName] = { nullptr, attr, llvm::None };
+
+}
 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;
+}
+
 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);
   // Optionals.
   if (curToken.getKind() == Token::l_paren)
     return parseOptional(element, isTopLevel);
   // Variables.