[Mlir-commits] [mlir] [MLIR][TableGen] Fix EnumAttr not recognized as enum in OpFormatGen (PR #189046)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`Attribute::isEnumAttr()` checked `isSubClassOf("EnumAttrInfo")`, but the newer `EnumAttr` class (from `mlir/include/mlir/IR/EnumAttr.td`) extends `AttrDef`, not `EnumAttrInfo`. This caused `canFormatEnumAttr` in `OpFormatGen.cpp` to return `false` for `EnumAttr`-backed attributes, forcing them to use the generic format path instead of the enum keyword format path.

The fix has two parts:
1. Update `Attribute::isEnumAttr()` to also check `isSubClassOf("EnumAttr")`.
2. Add `getEnumInfoRecord()` helper in `OpFormatGen.cpp` that, for `EnumAttr`-based attributes, retrieves the `enum` sub-field (which is the actual `EnumInfo` record) rather than the attribute def itself. Update `canFormatEnumAttr`, `genEnumAttrParser`, and `genEnumAttrPrinter` to use this helper.

Add a test to `op-format.td` verifying that an `EnumAttr`-wrapped enum generates the enum keyword format (symbolize/stringify calls).

Fixes #<!-- -->180455

Assisted-by: Claude Code

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


3 Files Affected:

- (modified) mlir/lib/TableGen/Attribute.cpp (+3-1) 
- (modified) mlir/test/mlir-tblgen/op-format.td (+24) 
- (modified) mlir/tools/mlir-tblgen/OpFormatGen.cpp (+13-3) 


``````````diff
diff --git a/mlir/lib/TableGen/Attribute.cpp b/mlir/lib/TableGen/Attribute.cpp
index d53ccea576b8b..ec7f4faafeb4e 100644
--- a/mlir/lib/TableGen/Attribute.cpp
+++ b/mlir/lib/TableGen/Attribute.cpp
@@ -52,7 +52,9 @@ bool Attribute::isSymbolRefAttr() const {
   return isSubClassOf("SymbolRefAttr") || isSubClassOf("FlatSymbolRefAttr");
 }
 
-bool Attribute::isEnumAttr() const { return isSubClassOf("EnumAttrInfo"); }
+bool Attribute::isEnumAttr() const {
+  return isSubClassOf("EnumAttrInfo") || isSubClassOf("EnumAttr");
+}
 
 StringRef Attribute::getStorageType() const {
   const auto *init = def->getValueInit("storageType");
diff --git a/mlir/test/mlir-tblgen/op-format.td b/mlir/test/mlir-tblgen/op-format.td
index 1790737a3a349..aa3cd7dde3685 100644
--- a/mlir/test/mlir-tblgen/op-format.td
+++ b/mlir/test/mlir-tblgen/op-format.td
@@ -1,6 +1,7 @@
 // RUN: mlir-tblgen -gen-op-defs -I %S/../../include %s | FileCheck %s
 
 include "mlir/IR/OpBase.td"
+include "mlir/IR/EnumAttr.td"
 
 def TestDialect : Dialect {
   let name = "test";
@@ -50,6 +51,29 @@ def CustomStringLiteralD : TestFormat_Op<[{
   custom<Foo>(prop-dict) attr-dict
 }]>;
 
+//===----------------------------------------------------------------------===//
+// EnumAttr formatting
+//===----------------------------------------------------------------------===//
+
+// Test that EnumAttr (backed by EnumInfo, not EnumAttrInfo) is recognized as
+// an enum attribute and uses the enum-keyword format path.
+
+def TestEnumCase0 : I32EnumCase<"Case0", 0>;
+def TestEnumCase1 : I32EnumCase<"Case1", 1>;
+
+def TestEnum : I32Enum<"TestEnum", "a test enum", [TestEnumCase0, TestEnumCase1]> {
+  let cppNamespace = "::test";
+}
+
+def TestEnumAttr : EnumAttr<TestDialect, TestEnum, "enum">;
+
+// CHECK-LABEL: EnumAttrOp::parse
+// CHECK: symbolizeTestEnum
+// CHECK-LABEL: EnumAttrOp::print
+// CHECK: stringifyTestEnum
+def EnumAttrOp : TestFormat_Op<"$attr attr-dict">,
+    Arguments<(ins TestEnumAttr:$attr)>;
+
 //===----------------------------------------------------------------------===//
 // Optional Groups
 //===----------------------------------------------------------------------===//
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index ff51fb403ffc8..f53076253b5a3 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -428,13 +428,23 @@ struct OperationFormat {
 // Parser Gen
 //===----------------------------------------------------------------------===//
 
+/// Returns the Record to use when constructing an EnumInfo for the given
+/// attribute. For legacy EnumAttrInfo-based attributes, this is the attribute
+/// def itself (which extends both EnumInfo and Attr). For newer EnumAttr-based
+/// attributes (which extend AttrDef), this is the `enum` sub-field.
+static const llvm::Record *getEnumInfoRecord(const Attribute &attr) {
+  if (attr.isSubClassOf("EnumAttr"))
+    return attr.getDef().getValueAsDef("enum");
+  return &attr.getDef();
+}
+
 /// Returns true if we can format the given attribute as an enum in the
 /// parser format.
 static bool canFormatEnumAttr(const NamedAttribute *attr) {
   Attribute baseAttr = attr->attr.getBaseAttr();
   if (!baseAttr.isEnumAttr())
     return false;
-  EnumInfo enumInfo(&baseAttr.getDef());
+  EnumInfo enumInfo(getEnumInfoRecord(baseAttr));
 
   // The attribute must have a valid underlying type and a constant builder.
   return !enumInfo.getUnderlyingType().empty() &&
@@ -1162,7 +1172,7 @@ static void genEnumAttrParser(const NamedAttribute *var, MethodBody &body,
                               FmtContext &attrTypeCtx, bool parseAsOptional,
                               bool useProperties, StringRef opCppClassName) {
   Attribute baseAttr = var->attr.getBaseAttr();
-  EnumInfo enumInfo(&baseAttr.getDef());
+  EnumInfo enumInfo(getEnumInfoRecord(baseAttr));
   std::vector<EnumCase> cases = enumInfo.getAllCases();
 
   // Generate the code for building an attribute for this enum.
@@ -2263,7 +2273,7 @@ static MethodBody &genTypeOperandPrinter(FormatElement *arg, const Operator &op,
 static void genEnumAttrPrinter(const NamedAttribute *var, const Operator &op,
                                MethodBody &body) {
   Attribute baseAttr = var->attr.getBaseAttr();
-  const EnumInfo enumInfo(&baseAttr.getDef());
+  const EnumInfo enumInfo(getEnumInfoRecord(baseAttr));
   std::vector<EnumCase> cases = enumInfo.getAllCases();
 
   body << formatv(enumAttrBeginPrinterCode,

``````````

</details>


https://github.com/llvm/llvm-project/pull/189046