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

[Mehdi Amini]

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

`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

>From 73f65a28193cf45b3a03915f8396b4e266915841 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Mar 2026 08:20:37 -0700
Subject: [PATCH] [MLIR][TableGen] Fix EnumAttr not recognized as enum in
 OpFormatGen

`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
---
 mlir/lib/TableGen/Attribute.cpp        |  4 +++-
 mlir/test/mlir-tblgen/op-format.td     | 24 ++++++++++++++++++++++++
 mlir/tools/mlir-tblgen/OpFormatGen.cpp | 16 +++++++++++++---
 3 files changed, 40 insertions(+), 4 deletions(-)

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,