[Mlir-commits] [mlir] 24b8892 - [mlir][ODS] Add new SymbolNameAttr and add support for in assemblyFormat

[River Riddle]

Author: River Riddle
Date: 2020-08-31T13:26:23-07:00
New Revision: 24b88920fed750c864e30add6a0fefe58f1a9f54

URL: https://github.com/llvm/llvm-project/commit/24b88920fed750c864e30add6a0fefe58f1a9f54
DIFF: https://github.com/llvm/llvm-project/commit/24b88920fed750c864e30add6a0fefe58f1a9f54.diff

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

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.

Differential Revision: https://reviews.llvm.org/D86759

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/op-format.mlir
    mlir/tools/mlir-tblgen/OpFormatGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 8375f2416062..247d486fedf5 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -979,6 +979,10 @@ class StringBasedAttr<Pred condition, string descr> : Attr<condition, descr> {
 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>()">,

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index bbe246a011af..0e186d0cd29b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1308,6 +1308,18 @@ def FormatOptAttrBOp : TEST_Op<"format_opt_attr_op_b"> {
   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);

diff  --git a/mlir/test/mlir-tblgen/op-format.mlir b/mlir/test/mlir-tblgen/op-format.mlir
index 96a923ba81ec..40eb17ac52dc 100644
--- a/mlir/test/mlir-tblgen/op-format.mlir
+++ b/mlir/test/mlir-tblgen/op-format.mlir
@@ -22,6 +22,15 @@ test.format_opt_attr_op_a
 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}
 

diff  --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index d3fdc9ec323e..684c8ad8cb17 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -395,6 +395,11 @@ static bool canFormatEnumAttr(const NamedAttribute *attr) {
          !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.
@@ -412,6 +417,19 @@ const char *const optionalAttrParserCode = R"(
   }
 )";
 
+/// 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.
@@ -861,6 +879,14 @@ static void genElementParser(Element *element, OpMethodBody &body,
       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;
@@ -1339,6 +1365,12 @@ static void genElementPrinter(Element *element, OpMethodBody &body,
       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";