[mlir][python bindings] generate all the enums

[makslevental]

Summary

This PR implements python enum bindings for *all* the enums - this includes I*Attrs (including positional/bit) and Dialect/EnumAttr.

There are a few parts to this:

1. CMake: a small addition to declare_mlir_dialect_python_bindings and declare_mlir_dialect_extension_python_bindings to generate the enum, a boolean arg GEN_ENUM_BINDINGS to make it opt-in (even though it works for basically all of the dialects), and an optional GEN_ENUM_BINDINGS_TD_FILE for handling corner cases.
2. EnumPythonBindingGen.cpp: there are two weedy aspects here that took investigation:
   1. If an enum attribute is not a Dialect/EnumAttr then the EnumAttrInfo record is canonical, as far as both the cases of the enum and the AttrDefName. On the otherhand, if an enum is a Dialect/EnumAttr then the EnumAttr record has the correct AttrDefName ("load bearing", i.e., populates ods.ir.AttributeBuilder('<NAME>')) but its enum field contains the cases, which is an instance of EnumAttrInfo. The solution is to generate an one enum class for both Dialect/EnumAttr and "independent" EnumAttrInfo but to make that class interopable with two builder registrations that both do the right thing (see next sub-bullet).
   2. Because we don't have a good connection to cpp EnumAttr, i.e., only the enum class getters are exposed (like DimensionAttr::get(Dimension value)), we have to resort to parsing e.g., Attribute.parse(f'#gpu<dim {x}>'). This means that the set of supported assemblyFormats (for the enum) is fixed at compile of MLIR (currently 2, the only 2 I saw). There might be some things that could be done here but they would require quite a bit more C API work to support generically (e.g., casting ints to enum cases and binding all the getters or going generically through the symbolize* methods, like symbolizeDimension(uint32_t) or symbolizeDimension(StringRef)).

A few small changes:

1. In addition, since this patch registers default builders for attributes where people might've had their own builders already written, I added a replace param to AttributeBuilder.insert (False by default).
2. makePythonEnumCaseName can't handle all the different ways in which people write their enum cases, e.g., llvm.CConv.Intel_OCL_BI, which gets turned into INTEL_O_C_L_B_I (because llvm::convertToSnakeFromCamelCase doesn't look for runs of caps). So I dropped it. On the otherhand regularization does need to done because some enums have None as a case (and others might have other python keywords).
3. I turned on llvm dialect generation here in order to test nvvm.WGMMAScaleIn, which is an enum with no explicit discriminator for the neg case.

Note, dialects that didn't get a GEN_ENUM_BINDINGS don't have any enums to generate.

Let me know if I should add more tests (the three trivial ones I added exercise both the supported assemblyFormats and replace=True).

[makslevental]

add top-level imports

[makslevental]

newline

[makslevental]

refactor

[makslevental]

guard no enums emitted

[makslevental]

refactor

[makslevental]

remove cruft

[makslevental]

update test

[makslevental]

treat EnumAttrs correctly

[makslevental]

treat EnumAttrs even more correctly

[makslevental]

rename

[makslevental]

rename

[makslevental]

rename

[makslevental]

rename

[makslevental]

add emptiness check

[makslevental]

remove unnecessary includes

[rkayaith]

The generated bindings: multiple dialects can emit tablegen records for the same enum attributes

I think it would be better to allow specifying a .td file that's different from the one that contains the ops in this case (iirc most dialects have a separate file for attrs/types/etc for this reason) and/or add an extra flag specifying which dialect should be generated (similar to https://github.com/llvm/llvm-project/blob/cb6fe61be3dee67bb8b080c73dd2c48f2d0ce2c7/mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp#L920-L922)

[makslevental]

refactor

[makslevental]

fix test

[makslevental]

add doc string to cmake

[makslevental]

fix doc in LinalgTransformEnums

[makslevental]

put back import mlir.dialects.gpu.passes

[makslevental]

update doc strings

[makslevental]

refactor

[makslevental]

update test

[makslevental]

restore emitAttributeBuilder

[makslevental]

fix register_builder name

[makslevental]

update test

[rkayaith]

mlir/cmake/modules/AddMLIRPython.cmake

276

Can we drop the "relative to ROOT_DIR" part, and point directly at the "real" td file? (the python bindings having their own wrapper td files should be removed, see TODO below)

mlir/python/CMakeLists.txt

219

To avoid accidentally picking up extra stuff in the future, I think these should all point to the same td file that's used for -gen-enum-decls etc., i.e. this one should point to VectorTransformsBase.td: https://github.com/llvm/llvm-project/blob/463e7cb89278936a86fef401abc20f171fcb6363/mlir/include/mlir/Dialect/Vector/Transforms/CMakeLists.txt#L1-L3

mlir/test/mlir-tblgen/enums-python-bindings.td

7

is the Enum import still needed?

7

can you update one of the tests here to show how auto gets used

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

52

Can this be merged with isPythonReserved https://github.com/llvm/llvm-project/blob/463e7cb89278936a86fef401abc20f171fcb6363/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp#L281-L291?
A slightly different list is used there, not sure if the differences are intentional.

116

what was the issue with bit enums before?

154

nit: make this a separate loop over getAllDerivedDefinitionsIfDefined("EnumAttr")? (unless there's some subtle difference? i'm not too familiar with these methods)

157

Can you add/use the methods on the llvm::Record wrappers for these? I think most should already exist on mlir::tblgen::AttrDef.

159–176

nit

[makslevental]

incorporate comments

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

52

I didn't know about that. I'm not sure how it was constructed (it's missing None, True, False, which are keywords) but I did python -c"import keyword; print(set(keyword.kwlist))". I was actually thinking of shelling out (since, presumably if you're generating you have python in your env) but that's probably overkill. I'll just update that list.

116

generating the 1<< case doesn't work because there's no wrapper about them in the python bindings (there's no CRTP for those attributes). So purely IntEnum didn't work but parsing works. Maybe a todo to just make the CRTP for BitEnums.

[makslevental]

point to the correct td enum file

[makslevental]

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

116

Actually this is wrong - it's just that all current eg I32BitEnumAttrs are wrapped in EnumAttr which required parsing - ie if you parse the mnemonic for the case it works.

Not sure if this is the reason EnumAttrs exist but currently you can't actually expose a bare I32BitEnumAttr because there's no way to go back and forth between an enum case and an APInt or int64_t (so the parsers will complain about missing stringify and symbolize methods). I have a draft patch (minus the tablegen parts) https://reviews.llvm.org/D158236 that implements this.

[makslevental]

implement bitenums correctly

[makslevental]

update tests

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

116

third time's the charm: implemented here correctly (no need for another patch)

[makslevental]

support py3.9

[rkayaith]

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

87–88

nit

90–94

doesn't IntFlag already have these defined? https://docs.python.org/3/library/enum.html#enum.Flag

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

18 ↗

(On Diff #551376)

include can be removed

[makslevental]

support py3.11

[makslevental]

support 3.7

[makslevental]

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

89

on >= py3.11 these things are implemented for us

167

shockingly, on py3.7 f'{x}' doesn't call __str__.

[makslevental]

mlir/test/mlir-tblgen/enums-python-bindings.td

67–71

for >=py3.11 this is handled by IntFlag

90

shockingly, on py3.7, f'{x}' doesn't call __str__

[makslevental]

mlir/test/python/dialects/arith_dialect.py

31–34 ↗

(On Diff #551383)

demo of bitenum union (note FastMathFlags uses , as the separator)

[makslevental]

comments

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

90–94

for <3.11, len() gives you the number of possible cases, not the current number of flags set

see

https://github.com/python/cpython/blob/3.10/Lib/enum.py#L449

vs

https://github.com/python/cpython/blob/3.11/Lib/enum.py#L1478-L1479

[makslevental]

put stringset back

[makslevental]

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

18 ↗

(On Diff #551376)

nope, there's another stringset down here isODSReserved

[rkayaith]

Seems generally fine to me, but I'd wait to see if anyone has opinions on the approach (using Attribute.parse) or enum name changes.

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

90–94

Ah I see. Importing private functions (_decompose) seems fragile though, can we just generate explicit bit tests? And if the logic works in later versions as well, seems simpler to use it regardless of python version and just leave a TODO to simplify if the min python version is raised to 3.11.

[jpienaar]

mlir/python/mlir/dialects/gpu/__init__.py

6

OOC is there a required order?

[makslevental]

mlir/python/mlir/dialects/gpu/__init__.py

6

no don't think so - neither module imports from the other and ops_gen only depends on enum_gen in so far as eg @_register_attribute_builder("GPU_AddressSpaceAttr") needs to be run at some point.

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

90–94

i agree it's brittel-ish but i doubt cpython will drop this for a point release? and i tested for 3.11>v>=3.8

[makslevental]

generate members for bitenum

[makslevental]

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

90–94

changed to generate explicitly

[stellaraccident]

This looks pretty reasonable to me, although I've commented on one simplification that I think we should do before landing this and having people start adapting their code.

I'm not completely comfortable with the fragility noted on enum builder calls, but neither can I see a straight-forward way out of it. Overall, I think this is a good improvement.

mlir/python/CMakeLists.txt

148

I think you can remove this entire GEN_ENUM_BINDINGS_TD_FILE carve-out if you instead just added to TransformOps.td:

include "mlir/Dialect/Transform/IR/TransformAttrs.td"

Basically, that Python level td file was intended to be the "container" for all of the records needed to generate bindings, and extending it in these cases is what it was made for.

Would eliminate a lot of special case stuff here and a chunk of CMake complexity.

[stellaraccident]

Chatted offline with Maks. I'm ok with this as-is with the caveat that we should:

• Address the TODO about refactoring TD_FILE and the wrappers quickly so as not to have the transient GEN_ENUM_BINDINGS_TD_FILE baked into downstreams (changed CMake arguments like that are hard to track down).
• Add a comment to GEN_ENUM_BINDINGS_TD_FILE indicating that downstreams should not use this as it will be removed shortly/without warning upon satisfying the TODO.

I think that satisfying the TODO involves:

• TD_FILES to directly reference all TD files.
• Add a GEN_DIR argument
• Keep TD_FILE with the current semantics for downstreams.

[makslevental]

add GEN_ENUM_BINDINGS_TD_FILE warning

[makslevental]

update structured_ext test

[rkayaith]

In D157934#4610971, @stellaraccident wrote:

• Address the TODO about refactoring TD_FILE and the wrappers quickly so as not to have the transient GEN_ENUM_BINDINGS_TD_FILE baked into downstreams (changed CMake arguments like that are hard to track down).
• Add a comment to GEN_ENUM_BINDINGS_TD_FILE indicating that downstreams should not use this as it will be removed shortly/without warning upon satisfying the TODO.

I didn't quite follow how we'd be able to get rid of GEN_ENUM_BINDINGS_TD_FILE. The issue is that:

• the "main" td file (with ops etc.) may include EnumAttr and EnumAttrInfo definitions for unrelated dialects, and
• EnumAttrInfos aren't associated with a specific dialect (they don't have a Dialect argument in tablegen), so I'm not sure how we'd filter out the unrelated ones.

So we need to be able to specify a .td file that only contains the relevant defs. And as far as I can tell, the C++ generation for these has the same restriction.

mlir/cmake/modules/AddMLIRPython.cmake

 #   ROOT_DIR: Same as for declare_mlir_python_sources().
 #   ADD_TO_PARENT: Same as for declare_mlir_python_sources(). Unique names
 #     for the subordinate source groups are derived from this.
 #   TD_FILE: Tablegen file to generate source for (relative to ROOT_DIR).
 #   DIALECT_NAME: Python name of the dialect.
 #   SOURCES: Same as declare_mlir_python_sources().
 #   SOURCES_GLOB: Same as declare_mlir_python_sources().
 #   DEPENDS: Additional dependency targets.
+#   ENUM_GEN: Generate enum mappings.
 #
 # TODO: Right now `TD_FILE` can't be the actual dialect tablegen file, since we
 #       use its path to determine where to place the generated python file. If
 #       we made the output path an additional argument here we could remove the
 #       need for the separate "wrapper" .td files
 function(declare_mlir_dialect_python_bindings)
   cmake_parse_arguments(ARG
-    ""
+    "ENUM_GEN"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
     "SOURCES;SOURCES_GLOB;DEPENDS"
     ${ARGN})
   # Sources.
   set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
   declare_mlir_python_sources(${_dialect_target}
     ROOT_DIR "${ARG_ROOT_DIR}"
     ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
if(ARG_TD_FILE)
     set(dialect_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_ops_gen.py")
     set(LLVM_TARGET_DEFINITIONS ${td_file})
     mlir_tablegen("${dialect_filename}"
       -gen-python-op-bindings -bind-dialect=${ARG_DIALECT_NAME}
       DEPENDS ${ARG_DEPENDS}
     )
     add_public_tablegen_target(${tblgen_target})
 
+    set(_sources ${dialect_filename})
+    if(ARG_ENUM_GEN)
+      set(enum_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_enum_gen.py")
+      mlir_tablegen(${enum_filename} -gen-python-enum-bindings)
+      list(APPEND _sources ${enum_filename})
+    endif()
+
     # Generated.
     declare_mlir_python_sources("${_dialect_target}.ops_gen"
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_dialect_target}"
-      SOURCES "${dialect_filename}"
+      SOURCES ${_sources}
     )
   endif()
 endfunction()
 
 # Function: declare_mlir_dialect_extension_python_bindings
 # Helper to generate source groups for dialect extensions, including both
 # static source files and a TD_FILE to generate wrappers.
 #
 # This will generate a source group named ${ADD_TO_PARENT}.${EXTENSION_NAME}.
 #
 # Arguments:
 #   ROOT_DIR: Same as for declare_mlir_python_sources().
 #   ADD_TO_PARENT: Same as for declare_mlir_python_sources(). Unique names
 #     for the subordinate source groups are derived from this.
 #   TD_FILE: Tablegen file to generate source for (relative to ROOT_DIR).
 #   DIALECT_NAME: Python name of the dialect.
 #   EXTENSION_NAME: Python name of the dialect extension.
 #   SOURCES: Same as declare_mlir_python_sources().
 #   SOURCES_GLOB: Same as declare_mlir_python_sources().
 #   DEPENDS: Additional dependency targets.
 function(declare_mlir_dialect_extension_python_bindings)
   cmake_parse_arguments(ARG
-    ""
+    "ENUM_GEN"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME"
     "SOURCES;SOURCES_GLOB;DEPENDS"
     ${ARGN})
   # Source files.
   set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}")
   declare_mlir_python_sources(${_extension_target}
     ROOT_DIR "${ARG_ROOT_DIR}"
     ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
if(ARG_TD_FILE)
     mlir_tablegen("${output_filename}" -gen-python-op-bindings
                   -bind-dialect=${ARG_DIALECT_NAME}
                   -dialect-extension=${ARG_EXTENSION_NAME})
     add_public_tablegen_target(${tblgen_target})
     if(ARG_DEPENDS)
       add_dependencies(${tblgen_target} ${ARG_DEPENDS})
     endif()
 
+    set(_sources ${output_filename})
+    if(ARG_ENUM_GEN)
+      set(enum_filename "${relative_td_directory}/_${ARG_EXTENSION_NAME}_enum_gen.py")
+      mlir_tablegen(${enum_filename} -gen-python-enum-bindings)
+      list(APPEND _sources ${enum_filename})
+    endif()
+
     declare_mlir_python_sources("${_extension_target}.ops_gen"
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_extension_target}"
-      SOURCES "${output_filename}"
+      SOURCES ${_sources}
     )
   endif()
 endfunction()
 
 # Function: mlir_python_setup_extension_rpath
 # Sets RPATH properties on a target, assuming that it is being output to
 # an _mlir_libs directory with all other libraries. For static linkage,
 # the RPATH will just be the origin. If linking dynamically, then the LLVM

mlir/lib/Bindings/Python/IRCore.cpp

 
 struct PyAttrBuilderMap {
   static bool dunderContains(const std::string &attributeKind) {
     return PyGlobals::get().lookupAttributeBuilder(attributeKind).has_value();
   }
   static py::function dundeGetItemNamed(const std::string &attributeKind) {
     auto builder = PyGlobals::get().lookupAttributeBuilder(attributeKind);
     if (!builder)
-      throw py::key_error();
+      throw py::key_error(attributeKind);
     return *builder;
   }
   static void dundeSetItemNamed(const std::string &attributeKind,
                                 py::function func) {
     PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func));
   }
 
   static void bind(py::module &m) {

mlir/lib/Bindings/Python/IRModule.cpp

void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
   // Note: Iterator cannot be shared from prior to loading, since re-entrancy
   // may have occurred, which may do anything.
   loadedDialectModulesCache.insert(dialectNamespace);
 }
 
 void PyGlobals::registerAttributeBuilder(const std::string &attributeKind,
                                          py::function pyFunc) {
   py::object &found = attributeBuilderMap[attributeKind];
-  if (found) {
+  if (found && !found.is_none()) {
     throw std::runtime_error((llvm::Twine("Attribute builder for '") +
-                              attributeKind + "' is already registered")
+                              attributeKind +
+                              "' is already registered with func: " +
+                              py::str(found).operator std::string())
                                  .str());
   }
   found = std::move(pyFunc);
 }
 
 void PyGlobals::registerTypeCaster(MlirTypeID mlirTypeID,
                                    pybind11::function typeCaster,
                                    bool replace) {

mlir/python/CMakeLists.txt

 ################################################################################
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/AMDGPUOps.td
   SOURCES
     dialects/amdgpu.py
-  DIALECT_NAME amdgpu)
+  DIALECT_NAME amdgpu
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/AsyncOps.td
   SOURCES_GLOB dialects/async_dialect/*.py
   DIALECT_NAME async_dialect)
 
SOURCES
     dialects/_func_ops_ext.py
   DIALECT_NAME func)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/GPUOps.td
   SOURCES_GLOB dialects/gpu/*.py
-  DIALECT_NAME gpu)
+  DIALECT_NAME gpu
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/LinalgOps.td
   SOURCES
     dialects/_linalg_ops_ext.py
   SOURCES_GLOB
     dialects/linalg/*.py
   DIALECT_NAME linalg
-  DEPENDS LinalgOdsGen)
+  DEPENDS LinalgOdsGen
+  ENUM_GEN)
 
 declare_mlir_dialect_extension_python_bindings(
 ADD_TO_PARENT MLIRPythonSources.Dialects
 ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformPDLExtensionOps.td
   SOURCES
     dialects/_transform_pdl_extension_ops_ext.py
     dialects/transform/pdl.py
   DIALECT_NAME transform
   EXTENSION_NAME transform_pdl_extension)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformOps.td
   SOURCES
     dialects/_transform_ops_ext.py
     dialects/transform/__init__.py
     _mlir_libs/_mlir/dialects/transform/__init__.pyi
-  DIALECT_NAME transform)
-
-set(LLVM_TARGET_DEFINITIONS "${CMAKE_CURRENT_SOURCE_DIR}/mlir/dialects/TransformOps.td")
-mlir_tablegen("dialects/_transform_enum_gen.py" -gen-python-enum-bindings)
-add_public_tablegen_target(MLIRTransformDialectPyEnumGen)
-declare_mlir_python_sources(
-  MLIRPythonSources.Dialects.transform.enum_gen
-  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
-  ADD_TO_PARENT MLIRPythonSources.Dialects.transform
-  SOURCES "dialects/_transform_enum_gen.py")
+  DIALECT_NAME transform
+  ENUM_GEN)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/BufferizationTransformOps.td
   SOURCES
     dialects/_bufferization_transform_ops_ext.py
     dialects/transform/bufferization.py
   DIALECT_NAME transform
-  EXTENSION_NAME bufferization_transform)
-
-set(LLVM_TARGET_DEFINITIONS "${CMAKE_CURRENT_SOURCE_DIR}/mlir/dialects/BufferizationTransformOps.td")
-mlir_tablegen("dialects/_bufferization_transform_enum_gen.py" -gen-python-enum-bindings)
-add_public_tablegen_target(MLIRBufferizationTransformDialectPyEnumGen)
-declare_mlir_python_sources(
-  MLIRPythonSources.Dialects.bufferization_transform.enum_gen
-  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
-  ADD_TO_PARENT MLIRPythonSources.Dialects.bufferization_transform
-  SOURCES "dialects/_bufferization_transform_enum_gen.py")
+  EXTENSION_NAME bufferization_transform
+  ENUM_GEN)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/GPUTransformOps.td
   SOURCES
     dialects/_gpu_transform_ops_ext.py
     dialects/transform/gpu.py

 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/LinalgStructuredTransformOps.td
   SOURCES
     dialects/_structured_transform_ops_ext.py
     dialects/transform/structured.py
   DIALECT_NAME transform
-  EXTENSION_NAME structured_transform)
+  EXTENSION_NAME structured_transform
+  ENUM_GEN)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TensorTransformOps.td
   SOURCES
     dialects/_tensor_transform_ops_ext.py
     dialects/transform/tensor.py
   DIALECT_NAME transform
   EXTENSION_NAME tensor_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/VectorTransformOps.td
   SOURCES
     dialects/transform/vector.py
   DIALECT_NAME transform
-  EXTENSION_NAME vector_transform)
-
-set(LLVM_TARGET_DEFINITIONS "${CMAKE_CURRENT_SOURCE_DIR}/mlir/dialects/VectorTransformOps.td")
-mlir_tablegen("dialects/_vector_transform_enum_gen.py" -gen-python-enum-bindings)
-add_public_tablegen_target(MLIRVectorTransformPyEnumGen)
-declare_mlir_python_sources(
-  MLIRPythonSources.Dialects.vector_transform.enum_gen
-  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
-  ADD_TO_PARENT MLIRPythonSources.Dialects.vector_transform
-  SOURCES "dialects/_vector_transform_enum_gen.py" )
+  EXTENSION_NAME vector_transform
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MathOps.td
   SOURCES dialects/math.py
-  DIALECT_NAME math)
+  DIALECT_NAME math
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/ArithOps.td
   SOURCES
     dialects/arith.py
     dialects/_arith_ops_ext.py
-  DIALECT_NAME arith)
+  DIALECT_NAME arith
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MemRefOps.td
   SOURCES
     dialects/memref.py
     dialects/_memref_ops_ext.py
-  DIALECT_NAME memref)
+  DIALECT_NAME memref
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MLProgramOps.td
   SOURCES
     dialects/ml_program.py
     dialects/_ml_program_ops_ext.py
   DIALECT_NAME ml_program)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/NVGPUOps.td
   SOURCES
     dialects/nvgpu.py
-  DIALECT_NAME nvgpu)
+  DIALECT_NAME nvgpu
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/NVVMOps.td
   SOURCES
     dialects/nvvm.py
-  DIALECT_NAME nvvm)
+  DIALECT_NAME nvvm
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/ROCDLOps.td
   SOURCES
     dialects/rocdl.py
   DIALECT_NAME rocdl)
 
 declare_mlir_python_sources(
   MLIRPythonSources.Dialects.quant
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ENUM_GEN
   SOURCES
     dialects/quant.py
     _mlir_libs/_mlir/dialects/quant.pyi)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/PDLOps.td
declare_mlir_dialect_python_bindings(
   SOURCES dialects/shape.py
   DIALECT_NAME shape)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/SparseTensorOps.td
   SOURCES dialects/sparse_tensor.py
-  DIALECT_NAME sparse_tensor)
+  DIALECT_NAME sparse_tensor
+  ENUM_GEN)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TensorOps.td
   SOURCES
     dialects/tensor.py
     dialects/_tensor_ops_ext.py
   DIALECT_NAME tensor)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TosaOps.td
   SOURCES dialects/tosa.py
   DIALECT_NAME tosa)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/VectorOps.td
   SOURCES dialects/vector.py
-  DIALECT_NAME vector)
+  DIALECT_NAME vector
+  ENUM_GEN)
 
 ################################################################################
 # Python extensions.
 # The sources for these are all in lib/Bindings/Python, but since they have to
 # be rebuilt for each package and integrate with the source setup here, we
 # just reference them here instead of having ordered, cross package target
 # dependencies.
 ################################################################################

mlir/python/mlir/dialects/amdgpu.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._amdgpu_ops_gen import *
+from ._amdgpu_enum_gen import *

mlir/python/mlir/dialects/arith.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._arith_ops_gen import *
+from ._arith_enum_gen import *

mlir/python/mlir/dialects/gpu/__init__.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._gpu_ops_gen import *
+from .._gpu_enum_gen import *

mlir/python/mlir/dialects/linalg/__init__.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Re-export the objects provided by pybind.
 from ..._mlir_libs._mlirDialectsLinalg import *
 
 # These are the backing OpView classes generated from the linalg tablegen
 # definitions following these steps:
 #   DSL -> YAML -> tblgen -> pytblgen -> build/.../_linalg_ops_gen.py.
 from .._linalg_ops_gen import *
+from .._linalg_enum_gen import *
 
 # These are the ground truth functions defined as:
 # ```
 #    @linalg_structured_op
 #    def matmul(A=TensorDef(T1, S.M, S.K),
 #               B=TensorDef(T2, S.K, S.N),
 #               C=TensorDef(U, S.M, S.N, output=True)):
 # ```

mlir/python/mlir/dialects/math.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._math_ops_gen import *
+from ._math_enum_gen import *

mlir/python/mlir/dialects/memref.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._memref_ops_gen import *
+from ._memref_enum_gen import *

mlir/python/mlir/dialects/nvgpu.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._nvgpu_ops_gen import *
+from ._nvgpu_enum_gen import *

mlir/python/mlir/dialects/nvvm.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._nvvm_ops_gen import *
+from ._nvvm_enum_gen import *

mlir/python/mlir/dialects/sparse_tensor.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._sparse_tensor_ops_gen import *
+from ._sparse_tensor_enum_gen import *
 from .._mlir_libs._mlirDialectsSparseTensor import *
 from .._mlir_libs import _mlirSparseTensorPasses as _cextSparseTensorPasses

mlir/python/mlir/dialects/transform/structured.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._structured_transform_ops_gen import *
+from .._structured_transform_enum_gen import *

mlir/python/mlir/dialects/vector.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._vector_ops_gen import *
+from ._vector_enum_gen import *

mlir/test/mlir-tblgen/enums-python-bindings.td

 // RUN: mlir-tblgen -gen-python-enum-bindings %s -I %S/../../include | FileCheck %s
 
 include "mlir/IR/EnumAttr.td"
 
 // CHECK: Autogenerated by mlir-tblgen; don't manually edit.
 
-// CHECK: from enum import Enum
+// CHECK: from enum import IntEnum, Enum
 // CHECK: from ._ods_common import _cext as _ods_cext
 // CHECK: _ods_ir = _ods_cext.ir
 
 def One : I32EnumAttrCase<"CaseOne", 1, "one">;
 def Two : I32EnumAttrCase<"CaseTwo", 2, "two">;
 
 def MyEnum : I32EnumAttr<"MyEnum", "An example 32-bit enum", [One, Two]>;
 // CHECK: def _register_attribute_builder(kind):
 // CHECK:     def decorator_builder(func):
-// CHECK:         _ods_ir.AttrBuilder.insert(kind, func)
+// CHECK:         if not _ods_ir.AttrBuilder.contains(kind):
+// CHECK:             _ods_ir.AttrBuilder.insert(kind, func)
 // CHECK:         return func
 // CHECK:     return decorator_builder
 
-// CHECK-LABEL: class MyEnum(Enum):
+// CHECK-LABEL: class MyEnum(IntEnum):
 // CHECK:     """An example 32-bit enum"""
 
 // CHECK:     CASE_ONE = 1
 // CHECK:     CASE_TWO = 2
 
-// CHECK:     def _as_int(self):
-// CHECK:         if self is MyEnum.CASE_ONE:
-// CHECK:             return 1
-// CHECK:         if self is MyEnum.CASE_TWO:
-// CHECK:             return 2
-// CHECK:         assert False, "Unknown MyEnum enum entry."
-
 def One64 : I64EnumAttrCase<"CaseOne64", 1, "one">;
 def Two64 : I64EnumAttrCase<"CaseTwo64", 2, "two">;
 
 def MyEnum64 : I64EnumAttr<"MyEnum64", "An example 64-bit enum", [One64, Two64]>;
 // CHECK: @_register_attribute_builder("MyEnum")
 // CHECK: def _my_enum(x, context):
-// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), x._as_int())
+// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
 
-// CHECK-LABEL: class MyEnum64(Enum):
+// CHECK-LABEL: class MyEnum64(IntEnum):
 // CHECK:     """An example 64-bit enum"""
 
 // CHECK:     CASE_ONE64 = 1
 // CHECK:     CASE_TWO64 = 2
 
-// CHECK:     def _as_int(self):
-// CHECK:         if self is MyEnum64.CASE_ONE64:
-// CHECK:             return 1
-// CHECK:         if self is MyEnum64.CASE_TWO64:
-// CHECK:             return 2
-// CHECK:         assert False, "Unknown MyEnum64 enum entry."
-
 // CHECK: @_register_attribute_builder("MyEnum64")
 // CHECK: def _my_enum64(x, context):
-// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(64, context=context), x._as_int())
+// CHECK:     return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(64, context=context), int(x))

mlir/test/python/dialects/gpu.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
-import mlir.dialects.gpu
-import mlir.dialects.gpu.passes
+import mlir.dialects.gpu as gpu
 from mlir.passmanager import *
 
 
 def run(f):
     print("\nTEST:", f.__name__)
+    with Context(), Location.unknown():
         f()
+    return f
 
 
+# CHECK-LABEL: testGPUPass
+#       CHECK: SUCCESS
+@run
 def testGPUPass():
-    with Context() as context:
-        PassManager.parse("any(gpu-kernel-outlining)")
     print("SUCCESS")
 
 
-# CHECK-LABEL: testGPUPass
-#       CHECK: SUCCESS
-run(testGPUPass)
+# CHECK-LABEL: testMMAElementWiseAttr
+@run
+def testMMAElementWiseAttr():
+    module = Module.create()
+    with InsertionPoint(module.body):
+        gpu.BlockDimOp(gpu.Dimension.Y)
+    # CHECK: %0 = gpu.block_dim  y
+    print(module)
+    pass

mlir/test/python/dialects/vector.py

with InsertionPoint(module.body):
             func.ReturnOp([])
 
     # CHECK: @transfer_read(%[[MEM:.*]]: memref<?x?xf32>, %[[IDX:.*]]: index,
     # CHECK: %[[PAD:.*]]: f32, %[[MASK:.*]]: vector<2x3xi1>)
     # CHECK: vector.transfer_read %[[MEM]][%[[IDX]], %[[IDX]]], %[[PAD]], %[[MASK]]
     # CHECK: vector.transfer_read %[[MEM]][%[[IDX]], %[[IDX]]], %[[PAD]]
     # CHECK-NOT: %[[MASK]]
     print(module)
+
+
+# CHECK-LABEL: TEST: testBitEnumCombiningKind
+@run
+def testBitEnumCombiningKind():
+    module = Module.create()
+    with InsertionPoint(module.body):
+        f32 = F32Type.get()
+        vector_type = VectorType.get([16], f32)
+
+        @func.FuncOp.from_py_func(vector_type)
+        def reduction(arg):
+            v = vector.ReductionOp(f32, vector.CombiningKind.ADD, arg)
+            return v
+
+    # CHECK: func.func @reduction(%[[VEC:.*]]: vector<16xf32>) -> f32 {
+    # CHECK: %0 = vector.reduction <add>, %[[VEC]] : vector<16xf32> into f32
+    print(module)

mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp

 //===- EnumPythonBindingGen.cpp - Generator of Python API for ODS enums ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // EnumPythonBindingGen uses ODS specification of MLIR enum attributes to
 // generate the corresponding Python binding classes.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/TableGen/Argument.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/GenInfo.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/TableGen/Record.h"
 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 /// File header and includes.
 constexpr const char *fileHeader = R"Py(
 # Autogenerated by mlir-tblgen; don't manually edit.
 
-from enum import Enum
+from enum import IntEnum, Enum
 from ._ods_common import _cext as _ods_cext
 _ods_ir = _ods_cext.ir
 
-# Convenience decorator for registering user-friendly Attribute builders.
+# Convenience decorator for (safely) registering user-friendly Attribute builders.
 def _register_attribute_builder(kind):
     def decorator_builder(func):
+        if not _ods_ir.AttrBuilder.contains(kind):
             _ods_ir.AttrBuilder.insert(kind, func)
         return func
 
     return decorator_builder
 
 )Py";
 
 /// Makes enum case name Python-compatible, i.e. UPPER_SNAKE_CASE.
 static std::string makePythonEnumCaseName(StringRef name) {
+  if (llvm::all_of(name, [](char c) {
+        return std::isupper(c) || c == '_' || std::isdigit(c);
+      }))
+    return name.str();
   return StringRef(llvm::convertToSnakeFromCamelCase(name)).upper();
 }
 
+/// Emits the Python class for the given bit enum.
+static void emitBitEnumClass(StringRef enumName, StringRef description,
+                             ArrayRef<EnumAttrCase> cases, raw_ostream &os) {
+  os << llvm::formatv("class {0}(str, Enum):\n", enumName);
+  if (!description.empty())
+    os << llvm::formatv("    \"\"\"{0}\"\"\"\n", description);
+  os << "    def __str__(self): return self.value\n";
+  os << "\n";
+
+  for (const EnumAttrCase &enumCase : cases) {
+    os << llvm::formatv("    {0} = \"{1}\"\n",
+                        makePythonEnumCaseName(enumCase.getSymbol()),
+                        enumCase.getStr());
+  }
+
+  os << "\n";
+}
+
 /// Emits the Python class for the given enum.
 static void emitEnumClass(StringRef enumName, StringRef description,
                           ArrayRef<EnumAttrCase> cases, raw_ostream &os) {
-  os << llvm::formatv("class {0}(Enum):\n", enumName);
+  os << llvm::formatv("class {0}(IntEnum):\n", enumName);
   if (!description.empty())
     os << llvm::formatv("    \"\"\"{0}\"\"\"\n", description);
   os << "\n";
 
   for (const EnumAttrCase &enumCase : cases) {
     os << llvm::formatv("    {0} = {1}\n",
                         makePythonEnumCaseName(enumCase.getSymbol()),
                         enumCase.getValue());
   }
 
   os << "\n";
-  os << llvm::formatv("    def _as_int(self):\n");
-  for (const EnumAttrCase &enumCase : cases) {
-    os << llvm::formatv("        if self is {0}.{1}:\n", enumName,
-                        makePythonEnumCaseName(enumCase.getSymbol()));
-    os << llvm::formatv("            return {0}\n", enumCase.getValue());
-  }
-  os << llvm::formatv("        assert False, \"Unknown {0} enum entry.\"\n\n\n",
-                      enumName);
 }
 
 /// Attempts to extract the bitwidth B from string "uintB_t" describing the
 /// type. This bitwidth information is not readily available in ODS. Returns
 /// `false` on success, `true` on failure.
 static bool extractUIntBitwidth(StringRef uintType, int64_t &bitwidth) {
   if (!uintType.consume_front("uint"))
     return true;
   if (!uintType.consume_back("_t"))
     return true;
   return uintType.getAsInteger(/*Radix=*/10, bitwidth);
 }
 
 /// Emits an attribute builder for the given enum attribute to support automatic
 /// conversion between enum values and attributes in Python. Returns
 /// `false` on success, `true` on failure.
-static bool emitAttributeBuilder(const EnumAttr &enumAttr, raw_ostream &os) {
+static bool emitAttributeBuilder(StringRef attrDefName,
+                                 const EnumAttr &enumAttr, raw_ostream &os) {
   int64_t bitwidth;
   if (extractUIntBitwidth(enumAttr.getUnderlyingType(), bitwidth)) {
     llvm::errs() << "failed to identify bitwidth of "
                  << enumAttr.getUnderlyingType();
     return true;
   }
 
-  os << llvm::formatv("@_register_attribute_builder(\"{0}\")\n",
-                      enumAttr.getAttrDefName());
+  os << llvm::formatv("@_register_attribute_builder(\"{0}\")\n", attrDefName);
   os << llvm::formatv(
       "def _{0}(x, context):\n",
       llvm::convertToSnakeFromCamelCase(enumAttr.getAttrDefName()));
   os << llvm::formatv(
       "    return "
       "_ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless({0}, "
-      "context=context), x._as_int())\n\n",
+      "context=context), int(x))\n\n",
       bitwidth);
   return false;
 }
 
+/// Emits an attribute builder for the given bit enum attribute to support
+/// automatic conversion between enum values and attributes in Python. Returns
+/// `false` on success, `true` on failure.
+static bool emitDialectOrBitEnumAttributeBuilder(StringRef attrDefName,
+                                                 StringRef formatString,
+                                                 raw_ostream &os) {
+  os << llvm::formatv("@_register_attribute_builder(\"{0}\")\n", attrDefName);
+  os << llvm::formatv("def _{0}(x, context):\n",
+                      llvm::convertToSnakeFromCamelCase(attrDefName));
+  os << llvm::formatv("    return "
+                      "_ods_ir.Attribute.parse(f'{0}', context=context)\n\n",
+                      formatString);
+  return false;
+}
+
 /// Emits Python bindings for all enums in the record keeper. Returns
 /// `false` on success, `true` on failure.
 static bool emitPythonEnums(const llvm::RecordKeeper &recordKeeper,
                             raw_ostream &os) {
   os << fileHeader;
-  std::vector<llvm::Record *> defs =
-      recordKeeper.getAllDerivedDefinitionsIfDefined("EnumAttrInfo");
-  for (const llvm::Record *def : defs) {
-    EnumAttr enumAttr(*def);
-    if (enumAttr.isBitEnum()) {
-      llvm::errs() << "bit enums not supported\n";
-      return true;
+
+  // Find the records corresponding to enums with load bearing names that
+  // correspond directly to op attribute names.
+  llvm::SmallMapVector<llvm::StringRef, llvm::Record *, 16> correctlyNamedEnums;
+  for (auto &it : recordKeeper.getDefs()) {
+    if (!it.second->isSubClassOf("EnumAttr") and
+        !it.second->isSubClassOf("EnumAttrInfo"))
+      continue;
+    EnumAttr *enumAttr;
+    if (it.second->isSubClassOf("EnumAttrInfo")) {
+      enumAttr = new EnumAttr(*it.second);
+      if (!correctlyNamedEnums.contains(enumAttr->getAttrDefName())) {
+        correctlyNamedEnums[enumAttr->getAttrDefName()] = &*it.second;
+      }
+    } else {
+      enumAttr = new EnumAttr(it.second->getValueAsDef("enum"));
+      correctlyNamedEnums[enumAttr->getAttrDefName()] = &*it.second;
     }
+  }
+
+  // Emit the enum classes and registration hooks with the correct names.
+  for (auto &it : correctlyNamedEnums) {
+    Attribute attr(it.second);
+    if (it.second->isSubClassOf("EnumAttrInfo")) {
+      EnumAttr enumAttr(*it.second);
       emitEnumClass(enumAttr.getEnumClassName(), enumAttr.getSummary(),
                     enumAttr.getAllCases(), os);
-    emitAttributeBuilder(enumAttr, os);
+      emitAttributeBuilder(attr.getAttrDefName(), enumAttr, os);
+    } else {
+      EnumAttr enumAttr(it.second->getValueAsDef("enum"));
+      emitBitEnumClass(enumAttr.getEnumClassName(), enumAttr.getSummary(),
+                       enumAttr.getAllCases(), os);
+      StringRef dialect =
+          attr.getDef().getValueAsDef("dialect")->getValueAsString("name");
+      StringRef mnemonic = attr.getDef().getValueAsString("mnemonic");
+      StringRef assemblyFormat =
+          attr.getDef().getValueAsString("assemblyFormat");
+
+      if ((assemblyFormat != "`<` $value `>`" && assemblyFormat != "$value") ||
+          attr.getDef().getValueAsBit("hasCustomAssemblyFormat")) {
+        llvm::errs() << "unsupported assembly format for enum";
+        return true;
+      }
+
+      if (assemblyFormat == "`<` $value `>`") {
+        emitDialectOrBitEnumAttributeBuilder(
+            attr.getAttrDefName(),
+            llvm::formatv("#{0}.{1}<{{x}>", dialect, mnemonic).str(), os);
+      } else if (assemblyFormat == "$value") {
+        emitDialectOrBitEnumAttributeBuilder(
+            attr.getAttrDefName(),
+            llvm::formatv("#{0}<{1} {{x}>", dialect, mnemonic).str(), os);
       }
+    }
+  }
+
   return false;
 }
 
 // Registers the enum utility generator to mlir-tblgen.
 static mlir::GenRegistration
     genPythonEnumBindings("gen-python-enum-bindings",
                           "Generate Python bindings for enum attributes",
                           &emitPythonEnums);