[Mlir-commits] [mlir] [MLIR][ODS] Add support for wrapping enums with std::optional in Type/Attr definitions (PR #117719)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/117719

None

>From 5dc8e92153bdba84f458a860802560d9503ec737 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Tue, 26 Nov 2024 06:58:25 -0800
Subject: [PATCH] [MLIR][ODS] Add support for wrapping enums with std::optional
 in Type/Attr definitions

---
 mlir/test/lib/Dialect/Test/TestAttrDefs.td    |  6 ++++++
 .../attr-or-type-format-roundtrip.mlir        |  8 ++++++-
 mlir/tools/mlir-tblgen/EnumsGen.cpp           | 21 +++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/mlir/test/lib/Dialect/Test/TestAttrDefs.td b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
index 907184b2e1ce4c..0fd272f85d39bd 100644
--- a/mlir/test/lib/Dialect/Test/TestAttrDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
@@ -226,6 +226,12 @@ def TestAttrWithOptionalUnsigned : Test_Attr<"TestAttrWithOptionalUnsigned"> {
   let mnemonic = "attr_with_optional_unsigned";
 }
 
+def TestAttrWithOptionalEnum : Test_Attr<"TestAttrWithOptionalEnum"> {
+  let parameters = (ins OptionalParameter<"std::optional<SimpleEnum>">:$value);
+  let assemblyFormat = "`<` $value `>`";
+  let mnemonic = "attr_with_optional_enum";
+}
+
 def TestAttrUgly : Test_Attr<"TestAttrUgly"> {
   let parameters = (ins "::mlir::Attribute":$attr);
 
diff --git a/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir b/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
index f0d1571d583877..77b94698a278e9 100644
--- a/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
+++ b/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
@@ -26,7 +26,13 @@ attributes {
   // CHECK: #test.attr_with_optional_signed<-9223372036854775808>
   attr9 = #test.attr_with_optional_signed<-9223372036854775808>,
   // CHECK: #test.attr_with_optional_unsigned<18446744073709551615>
-  attr_10 = #test.attr_with_optional_unsigned<18446744073709551615>
+  attr_10 = #test.attr_with_optional_unsigned<18446744073709551615>,
+  // CHECK: #test.attr_with_optional_enum<>
+  attr_11 = #test.attr_with_optional_enum<>,
+  // CHECK: #test.attr_with_optional_enum<a>
+  attr_12 = #test.attr_with_optional_enum<a>,
+  // CHECK: #test.attr_with_optional_enum<b>
+  attr_13 = #test.attr_with_optional_enum<b>
 }
 
 // CHECK-LABEL: @test_roundtrip_default_parsers_struct
diff --git a/mlir/tools/mlir-tblgen/EnumsGen.cpp b/mlir/tools/mlir-tblgen/EnumsGen.cpp
index 47ffd291cae687..d11aa9b27c2d86 100644
--- a/mlir/tools/mlir-tblgen/EnumsGen.cpp
+++ b/mlir/tools/mlir-tblgen/EnumsGen.cpp
@@ -101,6 +101,27 @@ struct FieldParser<{0}, {0}> {{
     return parser.emitError(loc, "invalid {2} specification: ") << enumKeyword;
   }
 };
+
+/// Support for std::optional, useful in attribute/type definition where the enum is
+/// used as:
+///
+///    let parameters = (ins OptionalParameter<"std::optional<TheEnumName>">:$value);
+template<>
+struct FieldParser<std::optional<{0}>, std::optional<{0}>> {{
+  template <typename ParserT>
+  static FailureOr<std::optional<{0}>> parse(ParserT &parser) {{
+    // Parse the keyword/string containing the enum.
+    std::string enumKeyword;
+    auto loc = parser.getCurrentLocation();
+    if (failed(parser.parseOptionalKeywordOrString(&enumKeyword)))
+      return std::optional<{0}>{{};
+
+    // Symbolize the keyword.
+    if (::std::optional<{0}> attr = {1}::symbolizeEnum<{0}>(enumKeyword))
+      return attr;
+    return parser.emitError(loc, "invalid {2} specification: ") << enumKeyword;
+  }
+};
 } // namespace mlir
 
 namespace llvm {