[mlir][ODS] Add new SymbolNameAttr and add support for in assemblyFormat

[rriddle]

Summary

Symbol names are a special form of StringAttr that get treated specially in certain areas, such as formatting. This revision adds a special derived attr for them in ODS and adds support in the assemblyFormat for formatting them properly.

Depends On D84719

[rriddle]

format

[rriddle]

Rebase

[mehdi_amini]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

399

can->should?

[rriddle]

Resolve comments

[rriddle]

Rebase

mlir/include/mlir/IR/OpBase.td

class StringBasedAttr<Pred condition, string descr> : Attr<condition, descr> {
   let storageType = [{ ::mlir::StringAttr }];
   let returnType = [{ ::llvm::StringRef }];
   let valueType = NoneType;
 }
 
 def StrAttr : StringBasedAttr<CPred<"$_self.isa<::mlir::StringAttr>()">,
                               "string attribute">;
 
+// A string attribute that represents the name of a symbol.
+def SymbolNameAttr : StringBasedAttr<CPred<"$_self.isa<::mlir::StringAttr>()">,
+                                     "string attribute">;
+
 // String attribute that has a specific value type.
 class TypedStrAttr<Type ty>
     : StringBasedAttr<CPred<"$_self.isa<::mlir::StringAttr>()">,
                             "string attribute"> {
   let valueType = ty;
 }
 
 // Base class for attributes containing types. Example:

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

def FormatOptAttrAOp : TEST_Op<"format_opt_attr_op_a"> {
   let arguments = (ins OptionalAttr<I64Attr>:$opt_attr);
   let assemblyFormat = "(`(` $opt_attr^ `)` )? attr-dict";
 }
 def FormatOptAttrBOp : TEST_Op<"format_opt_attr_op_b"> {
   let arguments = (ins OptionalAttr<I64Attr>:$opt_attr);
   let assemblyFormat = "($opt_attr^)? attr-dict";
 }
 
+// Test that we format symbol name attributes properly.
+def FormatSymbolNameAttrOp : TEST_Op<"format_symbol_name_attr_op"> {
+  let arguments = (ins SymbolNameAttr:$attr);
+  let assemblyFormat = "$attr attr-dict";
+}
+
+// Test that we format optional symbol name attributes properly.
+def FormatOptSymbolNameAttrOp : TEST_Op<"format_opt_symbol_name_attr_op"> {
+  let arguments = (ins OptionalAttr<SymbolNameAttr>:$opt_attr);
+  let assemblyFormat = "($opt_attr^)? attr-dict";
+}
+
 // Test that we elide attributes that are within the syntax.
 def FormatAttrDictWithKeywordOp : TEST_Op<"format_attr_dict_w_keyword"> {
   let arguments = (ins I64Attr:$attr, OptionalAttr<I64Attr>:$opt_attr);
   let assemblyFormat = "attr-dict-with-keyword";
 }
 
 // Test that we don't need to provide types in the format if they are buildable.
 def FormatBuildableTypeOp : TEST_Op<"format_buildable_type_op"> {

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

 test.format_opt_attr_op_a(10)
 test.format_opt_attr_op_a
 
 // CHECK: test.format_opt_attr_op_b 10
 // CHECK-NOT: {opt_attr
 test.format_opt_attr_op_b 10
 test.format_opt_attr_op_b
 
+// CHECK: test.format_symbol_name_attr_op @name
+// CHECK-NOT: {attr
+test.format_symbol_name_attr_op @name
+
+// CHECK: test.format_symbol_name_attr_op @opt_name
+// CHECK-NOT: {attr
+test.format_symbol_name_attr_op @opt_name
+test.format_opt_symbol_name_attr_op
+
 // CHECK: test.format_attr_dict_w_keyword attributes {attr = 10 : i64}
 test.format_attr_dict_w_keyword attributes {attr = 10 : i64}
 
 // CHECK: test.format_attr_dict_w_keyword attributes {attr = 10 : i64, opt_attr = 10 : i64}
 test.format_attr_dict_w_keyword attributes {attr = 10 : i64, opt_attr = 10 : i64}
 
 // CHECK: test.format_buildable_type_op %[[I64]]
 %ignored = test.format_buildable_type_op %i64

mlir/tools/mlir-tblgen/OpFormatGen.cpp

static bool canFormatEnumAttr(const NamedAttribute *attr) {
   if (!enumAttr)
     return false;
 
   // The attribute must have a valid underlying type and a constant builder.
   return !enumAttr->getUnderlyingType().empty() &&
          !enumAttr->getConstBuilderTemplate().empty();
 }
 
+/// Returns if we should format the given attribute as an SymbolNameAttr.
+static bool shouldFormatSymbolNameAttr(const NamedAttribute *attr) {
+  return attr->attr.getBaseAttr().getAttrDefName() == "SymbolNameAttr";
+}
+
 /// The code snippet used to generate a parser call for an attribute.
 ///
 /// {0}: The name of the attribute.
 /// {1}: The type for the attribute.
 const char *const attrParserCode = R"(
   if (parser.parseAttribute({0}Attr{1}, "{0}", result.attributes))
     return failure();
 )";
 const char *const optionalAttrParserCode = R"(
   {
     ::mlir::OptionalParseResult parseResult =
       parser.parseOptionalAttribute({0}Attr{1}, "{0}", result.attributes);
     if (parseResult.hasValue() && failed(*parseResult))
       return failure();
   }
 )";
 
+/// The code snippet used to generate a parser call for a symbol name attribute.
+///
+/// {0}: The name of the attribute.
+const char *const symbolNameAttrParserCode = R"(
+  if (parser.parseSymbolName({0}Attr, "{0}", result.attributes))
+    return failure();
+)";
+const char *const optionalSymbolNameAttrParserCode = R"(
+  // Parsing an optional symbol name doesn't fail, so no need to check the
+  // result.
+  (void)parser.parseOptionalSymbolName({0}Attr, "{0}", result.attributes);
+)";
+
 /// The code snippet used to generate a parser call for an enum attribute.
 ///
 /// {0}: The name of the attribute.
 /// {1}: The c++ namespace for the enum symbolize functions.
 /// {2}: The function to symbolize a string of the enum.
 /// {3}: The constant builder call to create an attribute of the enum type.
 const char *const enumAttrParserCode = R"(
   {
if (canFormatEnumAttr(var)) {
 
       body << formatv(var->attr.isOptional() ? optionalEnumAttrParserCode
                                              : enumAttrParserCode,
                       var->name, enumAttr.getCppNamespace(),
                       enumAttr.getStringToSymbolFnName(), attrBuilderStr);
       return;
     }
 
+    // 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);
     }
 
if (auto *attr = dyn_cast<AttributeVariable>(element)) {
     // If we are formatting as an enum, symbolize the attribute as a string.
     if (canFormatEnumAttr(var)) {
       const EnumAttr &enumAttr = cast<EnumAttr>(var->attr);
       body << "  p << \"\\\"\" << " << enumAttr.getSymbolToStringFnName() << "("
            << var->name << "()) << \"\\\"\";\n";
       return;
     }
 
+    // If we are formatting as a symbol name, handle it as a symbol name.
+    if (shouldFormatSymbolNameAttr(var)) {
+      body << "  p.printSymbolName(" << var->name << "Attr().getValue());\n";
+      return;
+    }
+
     // Elide the attribute type if it is buildable.
     if (attr->getTypeBuilder())
       body << "  p.printAttributeWithoutType(" << var->name << "Attr());\n";
     else
       body << "  p.printAttribute(" << var->name << "Attr());\n";
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     if (operand->getVar()->isOptional()) {
       body << "  if (::mlir::Value value = " << operand->getVar()->name