[mlir] load dialect in parser for optional parameters

[j2kun]

#96242 fixed an issue where the auto-generated parsers were not loading dialects whose namespaces are not present in the textual IR. This required the attribute parameter to be a tablegen def with its dialect information attached.

This fails when using parameter wrapper classes like OptionalParameter. This came up because RingAttr uses OptionalParameter for its second and third attributes. OptionalParameter takes as input the C++ type as a string instead of the tablegen def, and so it doesn't have a dialect member value to trigger the fix from #96242. The docs on this topic say the appropriate solution as overloading FieldParser for a particular type.

This PR updates FieldParser for generic attributes to load the dialect on demand. This requires mlir-tblgen to emit a dialectName static field on the generated attribute class, and check for it with template metaprogramming, since not all attribute types go through mlir-tblgen.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Jeremy Kun (j2kun)

Changes

Description of problem TBD.

---

Full diff: https://github.com/llvm/llvm-project/pull/96667.diff

4 Files Affected:

• (modified) mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp (+15)
• (modified) mlir/test/IR/parser.mlir (+1-1)
• (modified) mlir/test/lib/Dialect/Test/TestAttrDefs.td (+2-2)
• (modified) mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp (+5-3)

diff --git a/mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp b/mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp
index cc7d3172b1a1d..c13073e807adc 100644
--- a/mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp
+++ b/mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp
@@ -205,4 +205,19 @@ Attribute FloatPolynomialAttr::parse(AsmParser &parser, Type type) {
 }
 
 } // namespace polynomial
+
+/// Specialize FieldParser for parsing IntPolynomialAttr. This is required to
+/// support using IntPolynomialAttr in an OptionalParameter.
+template <>
+struct FieldParser<polynomial::IntPolynomialAttr> {
+  static FailureOr<polynomial::IntPolynomialAttr> parse(AsmParser &parser) {
+    parser.getContext()
+        ->getOrLoadDialect<mlir::polynomial::PolynomialDialect>();
+    Attribute polynomialAttr = polynomial::IntPolynomialAttr::parse(parser, {});
+    if (!polynomialAttr)
+      return failure();
+    return cast<polynomial::IntPolynomialAttr>(polynomialAttr);
+  }
+};
+
 } // namespace mlir
diff --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index e0d7f2a5dd0cb..dfdfe069d1914 100644
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -1470,7 +1470,7 @@ test.format_optional_result_d_op : f80
 
 // This is a testing that a non-qualified attribute in a custom format
 // correctly preload the dialect before creating the attribute.
-#attr = #test.nested_polynomial<<1 + x**2>>
+#attr = #test.nested_polynomial<poly=<1 + x**2>>
 // CHECK-lABLE: @parse_correctly
 llvm.func @parse_correctly() {
   test.containing_int_polynomial_attr #attr
diff --git a/mlir/test/lib/Dialect/Test/TestAttrDefs.td b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
index 9e25acf5f5ba4..471b78e489532 100644
--- a/mlir/test/lib/Dialect/Test/TestAttrDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
@@ -354,9 +354,9 @@ def TestCustomFloatAttr : Test_Attr<"TestCustomFloat"> {
 
 def NestedPolynomialAttr : Test_Attr<"NestedPolynomialAttr"> {
   let mnemonic = "nested_polynomial";
-  let parameters = (ins Polynomial_IntPolynomialAttr:$poly);
+  let parameters = (ins OptionalParameter<"::mlir::polynomial::IntPolynomialAttr">:$poly);
   let assemblyFormat = [{
-    `<` $poly `>`
+    `<` struct(params) `>`
   }];
 }
 
diff --git a/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp b/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
index 50378feeb14ed..91c4ce35889b5 100644
--- a/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
@@ -424,9 +424,11 @@ void DefFormat::genVariableParser(ParameterElement *el, FmtContext &ctx,
         Dialect dialect(dialectInit->getDef());
         auto cppNamespace = dialect.getCppNamespace();
         std::string name = dialect.getCppClassName();
-        dialectLoading = ("\nodsParser.getContext()->getOrLoadDialect<" +
-                          cppNamespace + "::" + name + ">();")
-                             .str();
+        if (name != "BuiltinDialect" || cppNamespace != "::mlir") {
+          dialectLoading = ("\nodsParser.getContext()->getOrLoadDialect<" +
+                            cppNamespace + "::" + name + ">();")
+                               .str();
+        }
       }
     }
   }

[github-actions]

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Discourse for more information.

[ftynse]

Should we complain when the field is missing? This would be the case for attributes defined outside tablegen.

[j2kun]

You mean a diagnostic? It seems that quite a few builtin attributes would trigger it (BoolAttr, e.g.,), so I'm not sure what one could do if they saw such a diagnostic... If we want to ensure all attributes have a dialectName string for this purpose, we could add it as a virtual method on Attribute. I don't have a strong opinion, I just tried to do the least invasive solution I could manage.

[ftynse]

I wonder if it would be possible to get the TypeID of the attribute class (there should be a static method) and somehow use it to find the dialect instead of relying on the textual name. Not sure if it is possible, I suspect we populate attribute TypeID maps when loading the dialect and here the dialect hasn't been loaded yet...

[j2kun]

AFAICT TypeID is not static on Attribute, nor part of the mlir-tblgen generated code.

[joker-eph]

The way you get TypeID is through TypeID::get<DialectClass>().
I would think that if you can generate a dialectName field, you could generate a TypeID one instead?

[j2kun]

If I had a generated TypeID, I'm still not sure how I would use it to get the dialect. This method seems to still need the dialect namespace as a string.

[ftynse]

I don't quite understand the second part of the condition. Why don't we want to load the dialect if its namespace starts with ::mlir? These shouldn't be anyhow special.

[j2kun]

Overall this condition is the logical negation of name == "BuiltinDialect" && namespace == "::mlir". For some unholy reason, someone might make an out-of-tree dialect called BuiltinDialect in a different namespace. In that case, we would still want to load it.

[j2kun]

I could also alternatively have getOrLoadDialect be aware of BuiltinDialect and handle it specially.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[j2kun]

@joker-eph @ftynse anything else needed here?

[joker-eph]

TypeID would be better than String lookup, but this isn't in the critical path of anything so string may be acceptable here.

[j2kun]

I will write a followup PR with the TypeID fix.