[mlir][IR] Fix enum attribute handling by using parseKeywordOrString instead of parseKeyword (#156662)

Change enum attribute parsing to handle special characters and
multi-word
identifiers. This allows enum attrs to use symbols like "+" and strings
with separators like "dash-separated-sentence" instead of being limited
to
valid identifiers. 

This also aligns enum attribute parsing with how enums are already
handled
by the `FieldParser`:

https://github.com/llvm/llvm-project/blob/main/mlir/tools/mlir-tblgen/EnumsGen.cpp#L108

Signed-off-by: Fabian Mora <[email protected]>

Author: fabianmcg

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);

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 = []

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";
 }

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

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.