From 47812756c8d25bc7d7e1f1ee0f4d33314799f436 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fabian.mora-cordero@amd.com>
Date: Wed, 3 Sep 2025 12:51:43 +0000
Subject: [PATCH] fix enum attr handling

Signed-off-by: Fabian Mora <fabian.mora-cordero@amd.com>
---
 mlir/include/mlir/IR/EnumAttr.td                 |  8 ++++----
 mlir/test/IR/array-of-attr.mlir                  |  4 ++--
 mlir/test/lib/Dialect/Test/TestEnumDefs.td       |  4 +++-
 .../attr-or-type-format-roundtrip.mlir           |  6 +++++-
 mlir/test/mlir-tblgen/attr-or-type-format.td     | 16 +++++++++++++++-
 5 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/mlir/include/mlir/IR/EnumAttr.td b/mlir/include/mlir/IR/EnumAttr.td
index 7eba68f05f4c3..4bc94809ed3ce 100644
--- a/mlir/include/mlir/IR/EnumAttr.td
+++ b/mlir/include/mlir/IR/EnumAttr.td
@@ -314,8 +314,8 @@ class IntEnumAttr<I intType, string name, string summary,
   // symbol is not valid.
   let parameterParser = [{[&]() -> ::mlir::FailureOr<}] # cppType # [{> {
     auto loc = $_parser.getCurrentLocation();
-    ::llvm::StringRef enumKeyword;
-    if (::mlir::failed($_parser.parseKeyword(&enumKeyword)))
+    std::string enumKeyword;
+    if (::mlir::failed($_parser.parseKeywordOrString(&enumKeyword)))
       return ::mlir::failure();
     auto maybeEnum = }] # cppNamespace # "::" #
                           stringToSymbolFnName # [{(enumKeyword);
@@ -436,9 +436,9 @@ class BitEnumAttr<I intType, string name, string summary,
   let parameterParser = [{[&]() -> ::mlir::FailureOr<}] # cppType # [{> {
     }] # cppType # [{ flags = {};
     auto loc = $_parser.getCurrentLocation();
-    ::llvm::StringRef enumKeyword;
+    std::string enumKeyword;
     do {
-      if (::mlir::failed($_parser.parseKeyword(&enumKeyword)))
+      if (::mlir::failed($_parser.parseKeywordOrString(&enumKeyword)))
         return ::mlir::failure();
       auto maybeEnum = }] # cppNamespace # "::" #
                             stringToSymbolFnName # [{(enumKeyword);
diff --git a/mlir/test/IR/array-of-attr.mlir b/mlir/test/IR/array-of-attr.mlir
index c2a6075965826..ad046eb8d1a07 100644
--- a/mlir/test/IR/array-of-attr.mlir
+++ b/mlir/test/IR/array-of-attr.mlir
@@ -6,8 +6,8 @@ test.array_of_attr_op
     a = [begin 0 : index end, begin 2 : index end],
     // CHECK-SAME: [0, 1, -42, 42]
     b = [0, 1, -42, 42],
-    // CHECK-SAME: [a, b, b, a]
-    c = [a, b, b, a]
+    // CHECK-SAME: [a, b, b, a, "+"]
+    c = [a, b, b, a, "+"]
 
 // CHECK: test.array_of_attr_op
 // CHECK-SAME: a = [], b = [], c = []
diff --git a/mlir/test/lib/Dialect/Test/TestEnumDefs.td b/mlir/test/lib/Dialect/Test/TestEnumDefs.td
index 51938d4dc68bb..70cf94e3458df 100644
--- a/mlir/test/lib/Dialect/Test/TestEnumDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestEnumDefs.td
@@ -49,7 +49,9 @@ def TestEnum
 
 def TestSimpleEnum : I32Enum<"SimpleEnum", "", [
     I32EnumCase<"a", 0>,
-    I32EnumCase<"b", 1>
+    I32EnumCase<"b", 1>,
+    I32EnumCase<"Plus", 2, "+">,
+    I32EnumCase<"LongString", 3, "dash-separated-sentence">,
   ]> {
   let cppNamespace = "::test";
 }
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 77b94698a278e..35554b20fbece 100644
--- a/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
+++ b/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
@@ -32,7 +32,11 @@ attributes {
   // 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>
+  attr_13 = #test.attr_with_optional_enum<b>,
+  // CHECK: #test<simple_enum"+">
+  attr_14 = #test<simple_enum "+">,
+  // CHECK: #test<simple_enum"dash-separated-sentence">
+  attr_15 = #test<simple_enum "dash-separated-sentence">
 }
 
 // CHECK-LABEL: @test_roundtrip_default_parsers_struct
diff --git a/mlir/test/mlir-tblgen/attr-or-type-format.td b/mlir/test/mlir-tblgen/attr-or-type-format.td
index 0f6b0c401a4e6..70c335f2f826c 100644
--- a/mlir/test/mlir-tblgen/attr-or-type-format.td
+++ b/mlir/test/mlir-tblgen/attr-or-type-format.td
@@ -200,7 +200,7 @@ def AttrE : TestAttr<"TestH"> {
 
 def TestEnum : I32EnumAttr<"TestEnum", "TestEnumType", [
   I32EnumAttrCase<"first", 0>,
-  I32EnumAttrCase<"second", 1>
+  I32EnumAttrCase<"second", 1>,
 ]> {
   let genSpecializedAttr = 0;
 }
@@ -215,6 +215,20 @@ def EnumAttrA : EnumAttr<Test_Dialect, TestEnum, "EnumAttrA"> {
   let assemblyFormat = "custom<Foo>($value)";
 }
 
+def TestEnumB : I32EnumAttr<"TestEnumB", "TestEnumType", [
+  I32EnumAttrCase<"Plus", 0, "+">,
+  I32EnumAttrCase<"LongString", 1, "dash-separated-sentence">,
+  I32EnumAttrCase<"Other", 2>
+]> {
+  let genSpecializedAttr = 0;
+}
+
+// ATTR-LABEL: TestEnumBAttr::parse
+// ATTR: parseKeywordOrString(
+def EnumAttrB : EnumAttr<Test_Dialect, TestEnumB, "EnumAttrB"> {
+  let assemblyFormat = "$value";
+}
+
 /// Test type parser and printer that mix variables and struct are generated
 /// correctly.