[mlir] Example/test code for Python binding of MLIR enum attributes

[jfurtek]

Summary

This diff contains test code that exercises the use of Enum attributes
(integer and bit) from Python code.

The implementation in this diff manually specifies the enum values. An approach
that leverages mlir-tblgen is possible, and is a logical next step. This test
code in this diff is intended to be a proof of concept, and to act as a template
for other dialects, as well as an eventual mlir-tblgen implementation.

Python binding for Enum attributes is necessary to implement fastmath
attributes in the Arithmetic dialect. (Other dialects generate arith
operations from Python, and recent changes to Python bindings require non-None
attribute values when creating these operations.)

[jfurtek]

Adding @mehdi_amini to suggest reviewers, since none were automatically triggered.

[stellaraccident]

Thanks for this. Just one comment of the variety of "have you considered doing it entirely differently) :) I'm not convinced either way -- open to discuss

mlir/test/python/lib/PythonTestCAPI.cpp

25

Just so I'm clear, you are wanting to auto generate this as a next step, right? (Which would be the first but if TableGen generated Capi we have)?

That would alleviate some of my concerns but is still a really heavy api and ask of the duplication comes with a real cost.

It occurs to me that as a class of things, we could supports common API for all enums based on some kind of opaque EnumHandle, just having one enums specific entrypoint to get the handle for a specific enums. If the common code also exposed an atoi and itoa like function, we could make the entire python implementation generic, based on reflection and exposed via an mlir.ir._DefEnum(handle_capsule).

Not that we care a ton, I guess, but pybind code really doesn't scale with either compile time or binary size to O(variants), and if reach attribute/enums just had one handle, we could reduce it to one c function and have a generated implementation that was pure c (as it would be trivial).

mlir/test/python/lib/PythonTestModule.cpp

14

This is pretty weird (directly including the .inc file here). I know why you're doing it...

[stellaraccident]

(also sorry for some auto correct typos - on mobile)

[jfurtek]

mlir/test/python/lib/PythonTestCAPI.cpp

25

you are wanting to auto generate this as a next step, right?

It seems like it could be automated. If "this is the way" (to use enums from Python), and if enough dialects care about such a thing, then it probably should be automated. But I didn't want to presume this would be the right strategy and go straight to yet-another-tblgen output.

have you considered doing it entirely differently

Yes, but probably from a position of relative ignorance. ;-)

My thought process was:

• Python code should have identifiers available for enum values (for readability)
• The C interface used by the Python binding requires passing "plain" integer values (because the tblgen-erated code uses C++ enum classes)
• Attribute creation from Python should fail if an invalid integer is used. (tblgen-erated C++ code already exists to check the integer value for validity, and there are other diffs "in flight" to automatically attach this to the MyEnumAttribute::getChecked() implementation.)

I thought a little about generating the enum definition in Python as opposed to C++/pybind. (Not a major simplification as you are proposing, but I was just trying to get something to work as a first step.) I abandoned this because it seemed a little clumsier to generate Python code with hard coded values (because the C++ enum identifiers are not visible).

we could reduce it to one c function and have a generated implementation that was pure c (as it would
be trivial)

I confess to not having totally grokked the inner workings of the MLIR python bindings. Would this require some sort of "registration" to allow enums/attributes from different dialects to be handled? Or would the capsule provide a pointer to the attribute-specific generator?

I'm (also) not a big fan of the verbosity of the approach in this diff. If you can see a path to something that is clearly better, and if you can spare a few minutes to describe it a little more, I can try to implement.

[stellaraccident]

mlir/test/python/lib/PythonTestCAPI.cpp

25

you are wanting to auto generate this as a next step, right?

It seems like it could be automated. If "this is the way" (to use enums from Python), and if enough dialects care about such a thing, then it probably should be automated. But I didn't want to presume this would be the right strategy and go straight to yet-another-tblgen output.

I think I'm ok with this for this first case. In my experience, though, there is a big hurdle to "tablegen-enabling" a new part of the tree, and until someone does it, we tend to accumulate sub-par workarounds. In this case, neither the C API nor the Python API has generated code like this. I think it is fine to land one such exemplar case, but then we need to do it right...

have you considered doing it entirely differently

Yes, but probably from a position of relative ignorance. ;-)

My thought process was:

• Python code should have identifiers available for enum values (for readability)
• The C interface used by the Python binding requires passing "plain" integer values (because the tblgen-erated code uses C++ enum classes)
• Attribute creation from Python should fail if an invalid integer is used. (tblgen-erated C++ code already exists to check the integer value for validity, and there are other diffs "in flight" to automatically attach this to the MyEnumAttribute::getChecked() implementation.)

I thought a little about generating the enum definition in Python as opposed to C++/pybind. (Not a major simplification as you are proposing, but I was just trying to get something to work as a first step.) I abandoned this because it seemed a little clumsier to generate Python code with hard coded values (because the C++ enum identifiers are not visible).

we could reduce it to one c function and have a generated implementation that was pure c (as it would
be trivial)

I confess to not having totally grokked the inner workings of the MLIR python bindings. Would this require some sort of "registration" to allow enums/attributes from different dialects to be handled? Or would the capsule provide a pointer to the attribute-specific generator?

I'm (also) not a big fan of the verbosity of the approach in this diff. If you can see a path to something that is clearly better, and if you can spare a few minutes to describe it a little more, I can try to implement.

I do see a hazy path in my mind. And I do think that my concern is that there are a lot of such cases spread up and down the users, and I generally wish we were not exposing so much API surface for a string/int mapping. Perhaps a half step to a better place would be to collapse the C API down to a single entry-point per enum attr:

struct MlirEnumAttributeHandle {
  bool (*isa)(MlirAttribute attr);
  MlirAttribute (*get_from_value)(int32_t value);
  ... Each of the others ...
};

const MlirEnumAttributeHandle *mlirPythonTestTestEnumAttributeGetHandle();

I am pretty sure we can extend something like that (which now gives us a pointer to a thing that lets us operate on the enum) into generic code on the Python side with a lot less coupling. In general, I'm looking to extend the DialectHandle support to be more reflectiony so that we can simplify the integration, and this may layer into that in the future.

Cons for doing it this way is that it deviates from how other attributes are defined, but I would argue that if we were defining another "class" of attributes to expose an API for like this in the future, we could follow this approach (we just haven't encountered that yet).

[jfurtek]

mlir/test/python/lib/PythonTestCAPI.cpp

25

Thanks for the review and for the explanation.

It seems like an approach that collapses enum handling to a single entry point would still benefit from tblgen code generation, and as you pointed out:

neither the C API nor the Python API has generated code like this.

Do you think that tblgen-ing code for the C API and/or Python is a reasonable path to go down?

[stellaraccident]

mlir/test/python/lib/PythonTestCAPI.cpp

25

For the C API, yes, I think we need to do that: we already have some suboptimal things there because we don't before this patch.

For the python pybind API, I'm less convinced but open to it if we have to go there. My primary concern is that the native code is not really set up to scale to per op/type/attribute. Ideally, I would rather expose some more reflection capabilities at the CAPI level so that we could program more of this generically in python. That's where I'm going with the existing DialectHandle bits and if we at least had a C object for the other types of interest (vs a bunch of loose functions), it would go in that direction. That's why I suggested this as a halfway point.

And also, I'm roughly ok with this patch and open coding the arith enum. I'd just like to have a better plan before it spiders.

mlir/python/CMakeLists.txt

declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
     ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
     SOURCES
       PythonTestModule.cpp
     PRIVATE_LINK_LIBS
       LLVMSupport
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
   )
+  include_directories("${MLIR_BINARY_DIR}/test/python/lib")
 endif()
 
 ################################################################################
 # Common CAPI dependency DSO.
 # All python extensions must link through one DSO which exports the CAPI, and
 # this must have a globally unique name amongst all embeddors of the python
 # library since it will effectively have global scope.
 #

mlir/python/mlir/dialects/python_test.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 ._python_test_ops_gen import *
-from .._mlir_libs._mlirPythonTest import TestAttr, TestType
+from .._mlir_libs._mlirPythonTest import (
+  TestEnum, TestBitEnum, TestAttr, TestEnumAttr, TestBitEnumAttr, TestType)
 
 def register_python_test_dialect(context, load=True):
   from .._mlir_libs import _mlirPythonTest
   _mlirPythonTest.register_python_test_dialect(context, load)

mlir/test/python/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS python_test_ops.td)
 mlir_tablegen(lib/PythonTestDialect.h.inc -gen-dialect-decls)
 mlir_tablegen(lib/PythonTestDialect.cpp.inc -gen-dialect-defs)
 mlir_tablegen(lib/PythonTestOps.h.inc -gen-op-decls)
 mlir_tablegen(lib/PythonTestOps.cpp.inc -gen-op-defs)
+mlir_tablegen(lib/PythonTestEnums.h.inc -gen-enum-decls)
+mlir_tablegen(lib/PythonTestEnums.cpp.inc -gen-enum-defs)
 mlir_tablegen(lib/PythonTestAttributes.h.inc -gen-attrdef-decls)
 mlir_tablegen(lib/PythonTestAttributes.cpp.inc -gen-attrdef-defs)
 mlir_tablegen(lib/PythonTestTypes.h.inc -gen-typedef-decls)
 mlir_tablegen(lib/PythonTestTypes.cpp.inc -gen-typedef-defs)
 add_public_tablegen_target(MLIRPythonTestIncGen)
 
 add_subdirectory(lib)

mlir/test/python/dialects/python_test.py

with Context() as ctx:
     test.register_python_test_dialect(ctx)
     a = test.TestAttr.get()
     # CHECK: #python_test.test_attr
     print(a)
 
     # The following cast must not assert.
     b = test.TestAttr(a)
 
+    # Verify pybind-generated enum comparisons.
+    # CHECK: True
+    print(test.TestEnum.OneSquared == test.TestEnum.OneSquared)
+
+    # CHECK: False
+    print(test.TestEnum.TwoSquared == test.TestEnum.ThreeSquared)
+
+    # CHECK: True
+    print(test.TestEnum.ThreeSquared.value == 9)
+
+    # Get a TestEnumAttr using the enum value.
+    c = test.TestEnumAttr.get(test.TestEnum.OneSquared, ctx)
+    # CHECK: #python_test<"test_enum OneSquared">
+    print(c)
+
+    # Get a TestEnumAttr using the associated integer value.
+    d = test.TestEnumAttr.get(4, ctx)
+    # CHECK: #python_test<"test_enum TwoSquared">
+    print(d)
+
+    # Get a TestBitEnumAttr using a single enum value.
+    e = test.TestBitEnumAttr.get(test.TestBitEnum.Bit0, ctx)
+    # CHECK: #python_test<"test_bit_enum Bit0">
+    print(e)
+
+    # Get a TestBitEnumAttr using multiple enum values.
+    f = test.TestBitEnumAttr.get(test.TestBitEnum.Bit0 | test.TestBitEnum.Bit3,
+                                 ctx)
+    # CHECK: #python_test<"test_bit_enum Bit0|Bit3">
+    print(f)
+
+    # Get a TestBitEnumAttr using the associated integer value.
+    g = test.TestBitEnumAttr.get(4 | 16, ctx)
+    # CHECK: #python_test<"test_bit_enum Bit2|Bit4">
+    print(g)
+
     unit = UnitAttr.get()
     try:
       test.TestAttr(unit)
     except ValueError as e:
       assert "Cannot cast attribute to TestAttr" in str(e)
     else:
       raise
 

mlir/test/python/lib/PythonTestCAPI.h

 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(PythonTest, python_test);
 
 MLIR_CAPI_EXPORTED bool
 mlirAttributeIsAPythonTestTestAttribute(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirPythonTestTestAttributeGet(MlirContext context);
 
+MLIR_CAPI_EXPORTED bool
+mlirAttributeIsAPythonTestTestEnumAttribute(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirPythonTestTestEnumAttributeGet(MlirContext context, uint32_t value);
+
+MLIR_CAPI_EXPORTED uint32_t
+mlirPythonTestTestEnumAttributeGetValue(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED bool
+mlirAttributeIsAPythonTestTestBitEnumAttribute(MlirAttribute attr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirPythonTestTestBitEnumAttributeGet(MlirContext context, uint32_t value);
+
+MLIR_CAPI_EXPORTED uint32_t
+mlirPythonTestTestBitEnumAttributeGetValue(MlirAttribute attr);
+
 MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestType(MlirType type);
 
 MLIR_CAPI_EXPORTED MlirType mlirPythonTestTestTypeGet(MlirContext context);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_TEST_PYTHON_LIB_PYTHONTESTCAPI_H

mlir/test/python/lib/PythonTestCAPI.cpp

 bool mlirAttributeIsAPythonTestTestAttribute(MlirAttribute attr) {
   return unwrap(attr).isa<python_test::TestAttrAttr>();
 }
 
 MlirAttribute mlirPythonTestTestAttributeGet(MlirContext context) {
   return wrap(python_test::TestAttrAttr::get(unwrap(context)));
 }
 
+bool mlirAttributeIsAPythonTestTestEnumAttribute(MlirAttribute attr) {
+  return unwrap(attr).isa<python_test::TestEnumAttr>();
+}
+
+MlirAttribute mlirPythonTestTestEnumAttributeGet(MlirContext context,
+                                                 uint32_t value) {
+  return wrap(python_test::TestEnumAttr::get(
+      unwrap(context), static_cast<python_test::TestEnum>(value)));
+}
+
+uint32_t mlirPythonTestTestEnumAttributeGetValue(MlirAttribute attr) {
+  return static_cast<uint32_t>(
+      unwrap(attr).cast<python_test::TestEnumAttr>().getValue());
+}
+
+bool mlirAttributeIsAPythonTestTestBitEnumAttribute(MlirAttribute attr) {
+  return unwrap(attr).isa<python_test::TestBitEnumAttr>();
+}
+
+MlirAttribute mlirPythonTestTestBitEnumAttributeGet(MlirContext context,
+                                                    uint32_t value) {
+  return wrap(python_test::TestBitEnumAttr::get(
+      unwrap(context), static_cast<python_test::TestBitEnum>(value)));
+}
+
+uint32_t mlirPythonTestTestBitEnumAttributeGetValue(MlirAttribute attr) {
+  return static_cast<uint32_t>(
+      unwrap(attr).cast<python_test::TestBitEnumAttr>().getValue());
+}
+
 bool mlirTypeIsAPythonTestTestType(MlirType type) {
   return unwrap(type).isa<python_test::TestTypeType>();
 }
 
 MlirType mlirPythonTestTestTypeGet(MlirContext context) {
   return wrap(python_test::TestTypeType::get(unwrap(context)));
 }

mlir/test/python/lib/PythonTestDialect.h

 #define MLIR_TEST_PYTHON_LIB_PYTHONTESTDIALECT_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 
 #include "PythonTestDialect.h.inc"
 
+#include "PythonTestEnums.h.inc"
+
 #define GET_OP_CLASSES
 #include "PythonTestOps.h.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "PythonTestAttributes.h.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "PythonTestTypes.h.inc"
 
 #endif // MLIR_TEST_PYTHON_LIB_PYTHONTESTDIALECT_H

mlir/test/python/lib/PythonTestDialect.cpp

 //===- PythonTestDialect.cpp - PythonTest dialect definition --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "PythonTestDialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #include "PythonTestDialect.cpp.inc"
+#include "PythonTestEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "PythonTestAttributes.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "PythonTestTypes.cpp.inc"
 
 #define GET_OP_CLASSES

mlir/test/python/lib/PythonTestModule.cpp

 //===- PythonTestModule.cpp - Python extension for the PythonTest dialect -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "PythonTestCAPI.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringRef.h"
+
+#include "PythonTestEnums.h.inc"
 
 namespace py = pybind11;
 using namespace mlir::python::adaptors;
 
+template <typename T> T enum_or(T a, T b) {
+  using utype = typename std::underlying_type<T>::type;
+  utype v = static_cast<utype>(a) | static_cast<utype>(b);
+  return static_cast<T>(v);
+}
+
 PYBIND11_MODULE(_mlirPythonTest, m) {
+  py::enum_<python_test::TestEnum>(m, "TestEnum", py::module_local())
+      .value("OneSquared", python_test::TestEnum::OneSquared)
+      .value("TwoSquared", python_test::TestEnum::TwoSquared)
+      .value("ThreeSquared", python_test::TestEnum::ThreeSquared);
+  py::enum_<python_test::TestBitEnum>(m, "TestBitEnum", py::module_local(),
+                                      py::arithmetic())
+      .value("Bit0", python_test::TestBitEnum::Bit0)
+      .value("Bit1", python_test::TestBitEnum::Bit1)
+      .value("Bit2", python_test::TestBitEnum::Bit2)
+      .value("Bit3", python_test::TestBitEnum::Bit3)
+      .value("Bit4", python_test::TestBitEnum::Bit4)
+      // Pybind does not wish to define the bitwise or operator for
+      // enum classes (because C++ does not by default) - see:
+      // https://github.com/pybind/pybind11/issues/2221
+      .def("__or__",
+           [](python_test::TestBitEnum e1, python_test::TestBitEnum e2) {
+             return enum_or(e1, e2);
+           })
+      .def("__ror__",
+           [](python_test::TestBitEnum e1, python_test::TestBitEnum e2) {
+             return enum_or(e1, e2);
+           });
   m.def(
       "register_python_test_dialect",
       [](MlirContext context, bool load) {
         MlirDialectHandle pythonTestDialect =
             mlirGetDialectHandle__python_test__();
         mlirDialectHandleRegisterDialect(pythonTestDialect, context);
         if (load) {
           mlirDialectHandleLoadDialect(pythonTestDialect, context);
         }
       },
       py::arg("context"), py::arg("load") = true);
 
   mlir_attribute_subclass(m, "TestAttr",
                           mlirAttributeIsAPythonTestTestAttribute)
       .def_classmethod(
           "get",
           [](py::object cls, MlirContext ctx) {
             return cls(mlirPythonTestTestAttributeGet(ctx));
           },
           py::arg("cls"), py::arg("context") = py::none());
+  mlir_attribute_subclass(m, "TestEnumAttr",
+                          mlirAttributeIsAPythonTestTestEnumAttribute)
+      .def_classmethod(
+          "get",
+          [](py::object cls, uint32_t value, MlirContext ctx) {
+            return cls(mlirPythonTestTestEnumAttributeGet(ctx, value));
+          },
+          py::arg("cls"), py::arg("value"), py::arg("context") = py::none())
+      .def_property_readonly("value", [](MlirAttribute self) {
+        auto value = mlirPythonTestTestEnumAttributeGetValue(self);
+        return static_cast<python_test::TestEnum>(value);
+      });
+  mlir_attribute_subclass(m, "TestBitEnumAttr",
+                          mlirAttributeIsAPythonTestTestBitEnumAttribute)
+      .def_classmethod(
+          "get",
+          [](py::object cls, uint32_t value, MlirContext ctx) {
+            return cls(mlirPythonTestTestBitEnumAttributeGet(ctx, value));
+          },
+          py::arg("cls"), py::arg("value"), py::arg("context") = py::none())
+      .def_property_readonly("value", [](MlirAttribute self) {
+        auto value = mlirPythonTestTestBitEnumAttributeGetValue(self);
+        return static_cast<python_test::TestBitEnum>(value);
+      });
   mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType)
       .def_classmethod(
           "get",
           [](py::object cls, MlirContext ctx) {
             return cls(mlirPythonTestTestTypeGet(ctx));
           },
           py::arg("cls"), py::arg("context") = py::none());
 }

mlir/test/python/python_test_ops.td

 //===-- python_test_ops.td - Python test Op definitions ----*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_TEST_OPS
 #define PYTHON_TEST_OPS
 
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 
 def Python_Test_Dialect : Dialect {
   let name = "python_test";
   let cppNamespace = "python_test";
 

 //===----------------------------------------------------------------------===//
 
 def TestType : TestType<"TestType", "test_type">;
 
 //===----------------------------------------------------------------------===//
 // Attribute definitions.
 //===----------------------------------------------------------------------===//
 
+def TestEnum :
+  I32EnumAttr<"TestEnum", "A test enum",
+  [I32EnumAttrCase<"OneSquared", 1>,
+   I32EnumAttrCase<"TwoSquared", 4>,
+   I32EnumAttrCase<"ThreeSquared", 9>, ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "python_test";
+}
+
+def TestBitEnum : I32BitEnumAttr<"TestBitEnum", "A test bit enum",
+  [I32BitEnumAttrCaseBit<"Bit0", 0>,
+   I32BitEnumAttrCaseBit<"Bit1", 1>,
+   I32BitEnumAttrCaseBit<"Bit2", 2>,
+   I32BitEnumAttrCaseBit<"Bit3", 3>,
+   I32BitEnumAttrCaseBit<"Bit4", 4>, ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "python_test";
+}
+
+
 def TestAttr : TestAttr<"TestAttr", "test_attr">;
+def TestEnumAttr : EnumAttr<Python_Test_Dialect, TestEnum, "test_enum">;
+def TestBitEnumAttr : EnumAttr<Python_Test_Dialect, TestBitEnum,
+                               "test_bit_enum">;
 
 //===----------------------------------------------------------------------===//
 // Operation definitions.
 //===----------------------------------------------------------------------===//
 
 def AttributedOp : TestOp<"attributed_op"> {
   let arguments = (ins I32Attr:$mandatory_i32,
                    OptionalAttr<I32Attr>:$optional_i32,