[all-commits] [llvm/llvm-project] 922330: [mlir][python bindings] generate all the enums

[Maksim Levental via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 92233062c17590d3157bdc6db430fcdfc54312fe
      https://github.com/llvm/llvm-project/commit/92233062c17590d3157bdc6db430fcdfc54312fe
  Author: max <maksim.levental at gmail.com>
  Date:   2023-08-23 (Wed, 23 Aug 2023)

  Changed paths:
    M mlir/cmake/modules/AddMLIRPython.cmake
    M mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td
    M mlir/lib/Bindings/Python/Globals.h
    M mlir/lib/Bindings/Python/IRCore.cpp
    M mlir/lib/Bindings/Python/IRModule.cpp
    M mlir/python/CMakeLists.txt
    A mlir/python/mlir/dialects/LLVMOps.td
    M mlir/python/mlir/dialects/amdgpu.py
    M mlir/python/mlir/dialects/arith.py
    M mlir/python/mlir/dialects/bufferization.py
    M mlir/python/mlir/dialects/gpu/__init__.py
    M mlir/python/mlir/dialects/linalg/__init__.py
    A mlir/python/mlir/dialects/llvm.py
    M mlir/python/mlir/dialects/nvgpu.py
    M mlir/python/mlir/dialects/nvvm.py
    M mlir/python/mlir/dialects/sparse_tensor.py
    M mlir/python/mlir/dialects/transform/bufferization.py
    M mlir/python/mlir/dialects/transform/structured.py
    M mlir/python/mlir/dialects/vector.py
    M mlir/python/mlir/ir.py
    M mlir/test/mlir-tblgen/enums-python-bindings.td
    M mlir/test/python/dialects/arith_dialect.py
    M mlir/test/python/dialects/gpu.py
    A mlir/test/python/dialects/llvm.py
    M mlir/test/python/dialects/nvvm.py
    M mlir/test/python/dialects/transform.py
    M mlir/test/python/dialects/transform_bufferization_ext.py
    M mlir/test/python/dialects/transform_gpu_ext.py
    M mlir/test/python/dialects/transform_loop_ext.py
    M mlir/test/python/dialects/transform_memref_ext.py
    M mlir/test/python/dialects/transform_structured_ext.py
    M mlir/test/python/dialects/transform_tensor_ext.py
    M mlir/test/python/dialects/transform_vector_ext.py
    M mlir/test/python/dialects/vector.py
    M mlir/tools/mlir-tblgen/EnumPythonBindingGen.cpp
    M mlir/tools/mlir-tblgen/OpGenHelpers.cpp
    M mlir/tools/mlir-tblgen/OpGenHelpers.h
    M mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

  Log Message:
  -----------
  [mlir][python bindings] generate all the enums

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 `assemblyFormat`s (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 [[ https://github.com/llvm/llvm-project/blob/d7e26b56207cbd8995296c5bb7c11ce676b649da/mlir/include/mlir/IR/EnumAttr.td#L22-L25 | 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 `assemblyFormat`s and `replace=True`).

Reviewed By: stellaraccident

Differential Revision: https://reviews.llvm.org/D157934