[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

1261–1262

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

1261–1262

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

1261–1262

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

1261–1262

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

2178

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

1261–1262

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 FormatOptionalDefaultAttr : TEST_Op<"format_optional_default_attr"> {
+  let arguments = (ins DefaultValuedAttr<StrAttr, "\"default\"">:$attr);
+  let assemblyFormat = "($attr^)? 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_attr "foo"
+test.format_optional_default_attr "foo"
+
+// CHECK: test.format_optional_default_attr
+// CHECK-NOT: default
+test.format_optional_default_attr "default"
+
 //===----------------------------------------------------------------------===//
 // 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(getContext()).getStringAttr("default")) {
+//  CHECK-NEXT:   odsPrinter << ' ';
+//  CHECK-NEXT:   odsPrinter.printAttributeWithoutType(getAAttr());
+//  CHECK-NEXT: }
+def OptionalGroupC : TestFormat_Op<[{
+  ($a^)? attr-dict
+}]>, Arguments<(ins DefaultValuedAttr<StrAttr, "\"default\"">:$a)>;

mlir/tools/mlir-tblgen/OpFormatGen.cpp

void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   paramList.emplace_back("::mlir::OperationState &", "result");
 
   auto *method = opClass.addStaticMethod("::mlir::ParseResult", "parse",
                                          std::move(paramList));
   auto &body = method->body();
 
   // Generate variables to store the operands and type within the format. This
   // allows for referencing these variables in the presence of optional
   // groupings.
   for (FormatElement *element : elements)
     genElementParserStorage(element, op, body);
 
   // A format context used when parsing attributes with buildable types.
   FmtContext attrTypeCtx;
   attrTypeCtx.withBuilder("parser.getBuilder()");
 
   // Generate parsers for each of the elements.
auto genElementParsers = [&](FormatElement *firstElement,
                << "\", parser.getBuilder().getUnitAttr());\n";
         }
       }
 
       // Generate the rest of the elements inside an optional group. Elements in
       // an optional group after the guard are parsed as required.
       for (FormatElement *childElement : elements)
         if (childElement != elidedAnchorElement)
-          genElementParser(childElement, body, attrTypeCtx,
-                           GenContext::Optional);
+          genElementParser(childElement, body, attrTypeCtx);
     };
 
     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);
+      genElementParser(attrVar, body, attrTypeCtx, GenContext::Optional);
       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)) {
     // 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()) {
+    if (genCtx == GenContext::Optional) {
       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);
     }
 
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(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");