[Mlir-commits] [mlir] try fixing optional parameter for dialect loading (PR #96667)

[Jeremy Kun]

https://github.com/j2kun created https://github.com/llvm/llvm-project/pull/96667

Description of problem TBD.

>From d0a5d93abf353ef77b2d978d6d23a536424316d3 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun at users.noreply.github.com>
Date: Tue, 25 Jun 2024 10:00:44 -0700
Subject: [PATCH] try fixing optional parameter for dialect loading

---
 .../Polynomial/IR/PolynomialAttributes.cpp        | 15 +++++++++++++++
 mlir/test/IR/parser.mlir                          |  2 +-
 mlir/test/lib/Dialect/Test/TestAttrDefs.td        |  4 ++--
 mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp    |  8 +++++---
 4 files changed, 23 insertions(+), 6 deletions(-)

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();
+        }
       }
     }
   }