[mlir][ods] Support default-valued attributes in optional groups

[rkayaith]

Summary

Add support for default-valued attributes as optional-group anchors. The
attribute is considered present if it holds a non-default value.

[rkayaith]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1275–1276

I didn't really understand the existing logic, it seemed odd to me to parse as optional under Normal context, but not under Optional context, but maybe I'm misunderstanding what "context" means here. I switched the use of GenContext so that genContext == Optional whenever the attribute should be optionally parsed (currently this is only for the first element of the then group).

[Mogball]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1275–1276

GenContext::Optional means the element is inside an optional group, not whether the element should be optionally parsed.

[rkayaith]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1275–1276

I see, is there supposed to be a difference to how things are parsed in optional groups (other than the first element)? The only place I see genCtx used is this line, and it seems like it's allowing optional attributes to be optionally parsed when outside an optional group, but that just ends up breaking in the printer: https://github.com/llvm/llvm-project/issues/58064

[Mogball]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1275–1276

Not really. Only that certain elements are prohibited in optional groups (e.g. nested optional groups). The issue you linked just looks like a bug to me. I think the feature you are adding is useful but I'd prefer not to fiddle with the gen context

[rkayaith]

preserve genCtx logic

[Mogball]

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

2187

there's a DefaultValuedStrAttr that lets you drop the quotes

[rkayaith]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1096

on main genCtx ends up being ignored on this path, should this be parseAsOptional instead?

and similarly for the isOptional() check inside genEnumAttrParser

1275–1276

ok, updated the diff and tried to preserve existing genCtx logic

[rkayaith]

fix default-valued symbol and enum attrs, use DefaultValuedStrAttr

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1096

i switched these over to use parseAsOptional; otherwise default-valued symbols and enums break

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

def FormatOptionalUnitAttrNoElide
   let assemblyFormat = "($is_optional^)? attr-dict";
 }
 
 def FormatOptionalEnumAttr : TEST_Op<"format_optional_enum_attr"> {
   let arguments = (ins OptionalAttr<SomeI64Enum>:$attr);
   let assemblyFormat = "($attr^)? attr-dict";
 }
 
+def FormatOptionalDefaultAttrs : TEST_Op<"format_optional_default_attrs"> {
+  let arguments = (ins DefaultValuedStrAttr<StrAttr, "default">:$str,
+                       DefaultValuedStrAttr<SymbolNameAttr, "default">:$sym,
+                       DefaultValuedAttr<SomeI64Enum, "SomeI64Enum::case5">:$e);
+  let assemblyFormat = "($str^)? ($sym^)? ($e^)? attr-dict";
+}
+
 def FormatOptionalWithElse : TEST_Op<"format_optional_else"> {
   let arguments = (ins UnitAttr:$isFirstBranchPresent);
   let assemblyFormat = "(`then` $isFirstBranchPresent^):(`else`)? attr-dict";
 }
 
 def FormatCompoundAttr : TEST_Op<"format_compound_attr"> {
   let arguments = (ins CompoundAttrA:$compound);
   let assemblyFormat = "$compound attr-dict-with-keyword";

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

 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidE : TestFormat_Op<[{
   (`,` $attr^ type(operands))? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: only one element can be marked as the anchor of an optional group
 def OptionalInvalidF : TestFormat_Op<[{
   ($attr^ $attr2^)? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr, OptionalAttr<I64Attr>:$attr2)>;
-// CHECK: error: only optional attributes can be used to anchor an optional group
+// CHECK: error: only optional or default-valued attributes can be used to anchor an optional group
 def OptionalInvalidG : TestFormat_Op<[{
   ($attr^)? attr-dict
 }]>, Arguments<(ins I64Attr:$attr)>;
 // CHECK: error: only variable length operands can be used within an optional group
 def OptionalInvalidH : TestFormat_Op<[{
   ($arg^)? attr-dict
 }]>, Arguments<(ins I64:$arg)>;
 // CHECK: error: only literals, types, and variables can be used within an optional group

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

 
 // CHECK: test.format_optional_enum_attr case5
 test.format_optional_enum_attr case5
 
 // CHECK: test.format_optional_enum_attr
 // CHECK-NOT: "case5"
 test.format_optional_enum_attr
 
+// CHECK: test.format_optional_default_attrs "foo" @foo case10
+test.format_optional_default_attrs "foo" @foo case10
+
+// CHECK: test.format_optional_default_attr
+// CHECK-NOT: "default"
+// CHECK-NOT: @default
+// CHECK-NOT: case5
+test.format_optional_default_attrs "default" @default case5
+
 //===----------------------------------------------------------------------===//
 // Format optional operands and results
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_optional_operand_result_a_op(%[[I64]] : i64) : i64
 test.format_optional_operand_result_a_op(%i64 : i64) : i64
 
 // CHECK: test.format_optional_operand_result_a_op( : ) : i64

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

 // CHECK-LABEL: OptionalGroupB::print
 // CHECK: if (!(*this)->getAttr("a"))
 // CHECK-NEXT: odsPrinter << ' ' << "foo"
 // CHECK-NEXT: else
 // CHECK-NEXT: odsPrinter << ' ' << "bar"
 def OptionalGroupB : TestFormat_Op<[{
   (`foo`) : (`bar` $a^)? attr-dict
 }]>, Arguments<(ins UnitAttr:$a)>;
+
+// Optional group anchored on a default-valued attribute:
+// CHECK-LABEL: OptionalGroupC::parse
+//       CHECK: if ((*this)->getAttr("a") != ::mlir::OpBuilder((*this)->getContext()).getStringAttr("default")) {
+//  CHECK-NEXT:   odsPrinter << ' ';
+//  CHECK-NEXT:   odsPrinter.printAttributeWithoutType(getAAttr());
+//  CHECK-NEXT: }
+def OptionalGroupC : TestFormat_Op<[{
+  ($a^)? attr-dict
+}]>, Arguments<(ins DefaultValuedStrAttr<StrAttr, "default">:$a)>;

mlir/tools/mlir-tblgen/OpFormatGen.cpp

for (FormatElement *param : dir->getArguments()) {
     }
   }
 
   body << "  }\n";
 }
 
 /// Generate the parser for a enum attribute.
 static void genEnumAttrParser(const NamedAttribute *var, MethodBody &body,
-                              FmtContext &attrTypeCtx) {
+                              FmtContext &attrTypeCtx, bool parseAsOptional) {
   Attribute baseAttr = var->attr.getBaseAttr();
   const EnumAttr &enumAttr = cast<EnumAttr>(baseAttr);
   std::vector<EnumAttrCase> cases = enumAttr.getAllCases();
 
   // Generate the code for building an attribute for this enum.
   std::string attrBuilderStr;
   {
     llvm::raw_string_ostream os(attrBuilderStr);
     os << tgfmt(enumAttr.getConstBuilderTemplate(), &attrTypeCtx,
                 "attrOptional.value()");
   }
 
   // Build a string containing the cases that can be formatted as a keyword.
   std::string validCaseKeywordsStr = "{";
   llvm::raw_string_ostream validCaseKeywordsOS(validCaseKeywordsStr);
   for (const EnumAttrCase &attrCase : cases)
     if (canFormatStringAsKeyword(attrCase.getStr()))
       validCaseKeywordsOS << '"' << attrCase.getStr() << "\",";
   validCaseKeywordsOS.str().back() = '}';
 
   // If the attribute is not optional, build an error message for the missing
   // attribute.
   std::string errorMessage;
-  if (!var->attr.isOptional()) {
+  if (!parseAsOptional) {
     llvm::raw_string_ostream errorMessageOS(errorMessage);
     errorMessageOS
         << "return parser.emitError(loc, \"expected string or "
            "keyword containing one of the following enum values for attribute '"
         << var->name << "' [";
     llvm::interleaveComma(cases, errorMessageOS, [&](const auto &attrCase) {
       errorMessageOS << attrCase.getStr();
     });
     errorMessageOS << "]\");";
   }
 
   body << formatv(enumAttrParserCode, var->name, enumAttr.getCppNamespace(),
                   enumAttr.getStringToSymbolFnName(), attrBuilderStr,
                   validCaseKeywordsStr, errorMessage);
 }
 
+// Generate the parser for an attribute.
+static void genAttrParser(AttributeVariable *attr, MethodBody &body,
+                          FmtContext &attrTypeCtx, bool parseAsOptional) {
+  const NamedAttribute *var = attr->getVar();
+
+  // Check to see if we can parse this as an enum attribute.
+  if (canFormatEnumAttr(var))
+    return genEnumAttrParser(var, body, attrTypeCtx, parseAsOptional);
+
+  // Check to see if we should parse this as a symbol name attribute.
+  if (shouldFormatSymbolNameAttr(var)) {
+    body << formatv(parseAsOptional ? optionalSymbolNameAttrParserCode
+                                    : symbolNameAttrParserCode,
+                    var->name);
+    return;
+  }
+
+  // If this attribute has a buildable type, use that when parsing the
+  // attribute.
+  std::string attrTypeStr;
+  if (Optional<StringRef> typeBuilder = attr->getTypeBuilder()) {
+    llvm::raw_string_ostream os(attrTypeStr);
+    os << tgfmt(*typeBuilder, &attrTypeCtx);
+  } else {
+    attrTypeStr = "::mlir::Type{}";
+  }
+  if (parseAsOptional) {
+    body << formatv(optionalAttrParserCode, var->name, attrTypeStr);
+  } else {
+    if (attr->shouldBeQualified() ||
+        var->attr.getStorageType() == "::mlir::Attribute")
+      body << formatv(genericAttrParserCode, var->name, attrTypeStr);
+    else
+      body << formatv(attrParserCode, var->name, attrTypeStr);
+  }
+}
+
 void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   SmallVector<MethodParameter> paramList;
   paramList.emplace_back("::mlir::OpAsmParser &", "parser");
   paramList.emplace_back("::mlir::OperationState &", "result");
 
   auto *method = opClass.addStaticMethod("::mlir::ParseResult", "parse",
                                          std::move(paramList));
   auto &body = method->body();
if (auto *optional = dyn_cast<OptionalElement>(element)) {
 
     ArrayRef<FormatElement *> thenElements =
         optional->getThenElements(/*parseable=*/true);
 
     // Generate a special optional parser for the first element to gate the
     // parsing of the rest of the elements.
     FormatElement *firstElement = thenElements.front();
     if (auto *attrVar = dyn_cast<AttributeVariable>(firstElement)) {
-      genElementParser(attrVar, body, attrTypeCtx);
+      genAttrParser(attrVar, body, attrTypeCtx, /*parseAsOptional=*/true);
       body << "  if (" << attrVar->getVar()->name << "Attr) {\n";
     } else if (auto *literal = dyn_cast<LiteralElement>(firstElement)) {
       body << "  if (::mlir::succeeded(parser.parseOptional";
       genLiteralParser(literal->getSpelling(), body);
       body << ")) {\n";
     } else if (auto *opVar = dyn_cast<OperandVariable>(firstElement)) {
       genElementParser(opVar, body, attrTypeCtx);
       body << "  if (!" << opVar->getVar()->name << "Operands.empty()) {\n";
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     genLiteralParser(literal->getSpelling(), body);
     body << ")\n    return ::mlir::failure();\n";
 
     /// Whitespaces.
   } else if (isa<WhitespaceElement>(element)) {
     // Nothing to parse.
 
     /// Arguments.
   } else if (auto *attr = dyn_cast<AttributeVariable>(element)) {
-    const NamedAttribute *var = attr->getVar();
-
-    // Check to see if we can parse this as an enum attribute.
-    if (canFormatEnumAttr(var))
-      return genEnumAttrParser(var, body, attrTypeCtx);
-
-    // Check to see if we should parse this as a symbol name attribute.
-    if (shouldFormatSymbolNameAttr(var)) {
-      body << formatv(var->attr.isOptional() ? optionalSymbolNameAttrParserCode
-                                             : symbolNameAttrParserCode,
-                      var->name);
-      return;
-    }
-
-    // If this attribute has a buildable type, use that when parsing the
-    // attribute.
-    std::string attrTypeStr;
-    if (Optional<StringRef> typeBuilder = attr->getTypeBuilder()) {
-      llvm::raw_string_ostream os(attrTypeStr);
-      os << tgfmt(*typeBuilder, &attrTypeCtx);
-    } else {
-      attrTypeStr = "::mlir::Type{}";
-    }
-    if (genCtx == GenContext::Normal && var->attr.isOptional()) {
-      body << formatv(optionalAttrParserCode, var->name, attrTypeStr);
-    } else {
-      if (attr->shouldBeQualified() ||
-          var->attr.getStorageType() == "::mlir::Attribute")
-        body << formatv(genericAttrParserCode, var->name, attrTypeStr);
-      else
-        body << formatv(attrParserCode, var->name, attrTypeStr);
-    }
+    bool parseAsOptional =
+        (genCtx == GenContext::Normal && attr->getVar()->attr.isOptional());
+    genAttrParser(attr, body, attrTypeCtx, parseAsOptional);
 
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     ArgumentLengthKind lengthKind = getArgumentLengthKind(operand->getVar());
     StringRef name = operand->getVar()->name;
     if (lengthKind == ArgumentLengthKind::VariadicOfVariadic)
       body << llvm::formatv(
           variadicOfVariadicOperandParserCode, name,
           operand->getVar()->constraint.getVariadicOfVariadicSegmentSizeAttr());
TypeSwitch<FormatElement *>(anchor)
         body << "!" << name << "().empty()";
       })
       .Case<TypeDirective>([&](TypeDirective *element) {
         genOptionalGroupPrinterAnchor(element->getArg(), op, body);
       })
       .Case<FunctionalTypeDirective>([&](FunctionalTypeDirective *element) {
         genOptionalGroupPrinterAnchor(element->getInputs(), op, body);
       })
-      .Case<AttributeVariable>([&](AttributeVariable *attr) {
-        body << "(*this)->getAttr(\"" << attr->getVar()->name << "\")";
+      .Case<AttributeVariable>([&](AttributeVariable *element) {
+        Attribute attr = element->getVar()->attr;
+        body << "(*this)->getAttr(\"" << element->getVar()->name << "\")";
+        if (attr.isOptional())
+          return; // done
+        if (attr.hasDefaultValue()) {
+          // Consider a default-valued attribute as present if it's not the
+          // default value.
+          FmtContext fctx;
+          fctx.withBuilder("::mlir::OpBuilder((*this)->getContext())");
+          body << " != "
+               << tgfmt(attr.getConstBuilderTemplate(), &fctx,
+                        attr.getDefaultValue());
+          return;
+        }
+        llvm_unreachable("attribute must be optional or default-valued");
       });
 }
 
 void collect(FormatElement *element,
              SmallVectorImpl<VariableElement *> &variables) {
   TypeSwitch<FormatElement *>(element)
       .Case([&](VariableElement *var) { variables.emplace_back(var); })
       .Case([&](CustomDirective *ele) {

 
 LogicalResult OpFormatParser::verifyOptionalGroupElement(SMLoc loc,
                                                          FormatElement *element,
                                                          bool isAnchor) {
   return TypeSwitch<FormatElement *, LogicalResult>(element)
       // All attributes can be within the optional group, but only optional
       // attributes can be the anchor.
       .Case([&](AttributeVariable *attrEle) {
-        if (isAnchor && !attrEle->getVar()->attr.isOptional())
-          return emitError(loc, "only optional attributes can be used to "
-                                "anchor an optional group");
+        Attribute attr = attrEle->getVar()->attr;
+        if (isAnchor && !(attr.isOptional() || attr.hasDefaultValue()))
+          return emitError(loc, "only optional or default-valued attributes "
+                                "can be used to anchor an optional group");
         return success();
       })
       // Only optional-like(i.e. variadic) operands can be within an optional
       // group.
       .Case([&](OperandVariable *ele) {
         if (!ele->getVar()->isVariableLength())
           return emitError(loc, "only variable length operands can be used "
                                 "within an optional group");