[PSA] Stop using StrEnumAttr (Do use EnumAttr)

MLIR Deprecation & Important Refactoring
1

StrEnumAttr is going to be deprecated soon in favour of EnumAttr, which should work as a drop-in replacement.

StrEnumAttr is slow and memory inefficient; there aren’t many good reasons to use strings to store enum values over integers in MLIR other than for prettier assembly formats. However, with AttrDefs and attribute assembly formats, this is redundant.

Here’s how you can convert a StrEnumAttr to use the new EnumAttr

// Before
def MyEnumAttr : StrEnumAttr<"MyEnum", "", [
  StrEnumAttrCase<"A">,
  StrEnumAttrCase<"B">
]>;

// After (pick an integer enum of your choice)
def MyEnum : I32EnumAttr<"MyEnum", "", [
  I32EnumAttrCase<"A", 0>,
  I32EnumAttrCase<"B", 1>
]> {
  // Don't generate a C++ class! We want to use the AttrDef
  let genSpecializedAttr = 0;
}
// Define the AttrDef
def MyEnum : EnumAttr<MyDialect, MyEnum, "my_enum">;

The assembly format will change as well:

// For `assemblyFormat = "$my_enum attr-dict"`
my_dialect.my_op "A"
// Becomes
my_dialect.my_op A

And no more string comparisons in C++!

if (op.getEnum() == "A") { ... }
// Becomes
if (op.getEnum() == MyEnum::A) { ... }

At the moment, it requires slightly more boilerplate code to define an EnumAttr. However, once StrEnumAttr (and later BitEnumAttr) are deprecated, much of this boilerplate code can go away.

2 Likes