[MLIR][python bindings] implement `PyValue` subclassing to enable operator overloading

[makslevental]

Summary

Refactor python bindings to enable subclassing values and overloading operators on both python and cpp side. The test shows one use-case (tensor wrapper indexing maps to tensor.extract).

There are two ways to make this work (both on display in the sequence of commits);

1. Move PyConcreteValue to a public header(https://reviews.llvm.org/D147758?vs=on&id=511599#toc)
2. (Current) Add mlir_value_subclass but also add a python side constructor to py::class_<PyValue>

I can't comment on the relative merits/tradeoffs; first one "breaks symmetry" with mlir_type_subclass and mlir_attribute_subclass, while second potentially does a bad thing in adding a constructor to py::class_<PyValue> where I believe the intent is to not have one.

Note in the diff demonstrating the first approach there are some stray PYBIND11_EXPORTs (i.e., they are not necessary).

[makslevental]

formatting

[makslevental]

formatting

[makslevental]

cleaner approach

[makslevental]

formatting

[makslevental]

formatting

[makslevental]

formatting

[makslevental]

update cmake

[makslevental]

mlir/lib/Bindings/Python/IRCore.cpp

3263

I'll be honest I'm not 100% sure why this works (cf the Value(t) in the test) since there is no single argument constructor for PyValue. But it does indeed pass on in all my envs.

[ftynse]

Using pure_subclass is the right way, the PyConcrete* doesn't scale properly to out-of-tree projects. We didn't need it for values because value kinds are not user-extensible, there are only block arguments and operation results. That being said, I see the value of letting Python subclass those similarly to how C++ now provides TypedValue wrappers.

Why does this need to move PybindUtils.h? It feels like an implementation detail that may not be reusable.

mlir/lib/Bindings/Python/IRCore.cpp

3263

Implicit copy constructor?

mlir/test/python/dialects/python_test.py

307–345

I'd rather keep this as a separate test for "extensions", I don't think we are ready for IR-building operator overloads being perceived as part of main bindings.

[makslevental]

undo PybindUtils.h move

[makslevental]

mlir/test/python/lib/PythonTestModule.cpp

46

just a small (trivial) illustration for how the subclassing can be used to implement something on the cpp side as well.

[ftynse]

mlir/test/python/dialects/python_test.py

346

Nit: please add the newline.

[makslevental]

fix nits

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

mlir_type_subclass(py::handle scope, const char *typeClassName,
     // 'isinstance' method.
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirType other) { return isaFunction(other); },
         py::arg("other_type"));
   }
 };
 
+/// Creates a custom subclass of mlir.ir.Type, implementing a casting
+/// constructor and type checking methods.
+class mlir_value_subclass : public pure_subclass {
+public:
+  using IsAFunctionTy = bool (*)(MlirValue);
+
+  /// Subclasses by looking up the super-class dynamically.
+  mlir_value_subclass(py::handle scope, const char *valueClassName,
+                     IsAFunctionTy isaFunction)
+      : mlir_value_subclass(
+            scope, valueClassName, isaFunction,
+            py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Value")) {}
+
+  /// Subclasses with a provided mlir.ir.Value super-class. This must
+  /// be used if the subclass is being defined in the same extension module
+  /// as the mlir.ir class (otherwise, it will trigger a recursive
+  /// initialization).
+  mlir_value_subclass(py::handle scope, const char *valueClassName,
+                     IsAFunctionTy isaFunction, const py::object &superCls)
+      : pure_subclass(scope, valueClassName, superCls) {
+    // Casting constructor. Note that it hard, if not impossible, to properly
+    // call chain to parent `__init__` in pybind11 due to its special handling
+    // for init functions that don't have a fully constructed self-reference,
+    // which makes it impossible to forward it to `__init__` of a superclass.
+    // Instead, provide a custom `__new__` and call that of a superclass, which
+    // eventually calls `__init__` of the superclass. Since attribute subclasses
+    // have no additional members, we can just return the instance thus created
+    // without amending it.
+    std::string captureValueName(
+        valueClassName); // As string in case if valueClassName is not static.
+    py::cpp_function newCf(
+        [superCls, isaFunction, captureValueName](py::object cls,
+                                                 py::object otherValue) {
+          MlirValue rawValue = py::cast<MlirValue>(otherValue);
+          if (!isaFunction(rawValue)) {
+            auto origRepr = py::repr(otherValue).cast<std::string>();
+            throw std::invalid_argument((llvm::Twine("Cannot cast value to ") +
+                                         captureValueName + " (from " +
+                                         origRepr + ")")
+                                            .str());
+          }
+          py::object self = superCls.attr("__new__")(cls, otherValue);
+          return self;
+        },
+        py::name("__new__"), py::arg("cls"), py::arg("cast_from_value"));
+    thisClass.attr("__new__") = newCf;
+
+    // 'isinstance' method.
+    def_staticmethod(
+        "isinstance",
+        [isaFunction](MlirValue other) { return isaFunction(other); },
+        py::arg("other_value"));
+  }
+};
+
 } // namespace adaptors
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H

mlir/lib/Bindings/Python/Globals.h

 //===- Globals.h - MLIR Python extension globals --------------------------===//
 //
 // 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 MLIR_BINDINGS_PYTHON_GLOBALS_H
 #define MLIR_BINDINGS_PYTHON_GLOBALS_H
 
 #include <string>
 #include <vector>
 #include <optional>
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 
 namespace mlir {
 namespace python {
 
 /// Globals that are always accessible once the extension has been initialized.

mlir/lib/Bindings/Python/IRAffine.cpp

 //===- IRAffine.cpp - Exports 'ir' module affine related bindings ---------===//
 //
 // 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 <utility>
 
 #include "IRModule.h"
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/IntegerSet.h"
 #include "llvm/ADT/Hashing.h"
 
 namespace py = pybind11;
 using namespace mlir;

mlir/lib/Bindings/Python/IRAttributes.cpp

 //===- IRAttributes.cpp - Exports builtin and standard attributes ---------===//
 //
 // 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 <utility>
 #include <optional>
 
 #include "IRModule.h"
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 

mlir/lib/Bindings/Python/IRCore.cpp

 //===- IRModules.cpp - IR Submodules of pybind module ---------------------===//
 //
 // 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 "IRModule.h"
 
 #include "Globals.h"
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 //#include "mlir-c/Registration.h"
py::class_<PyType>(m, "Type", py::module_local())
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Value.
   //----------------------------------------------------------------------------
   py::class_<PyValue>(m, "Value", py::module_local())
+      .def(py::init<PyValue &>(), py::keep_alive<0, 1>(), py::arg("value"))
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyValue::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyValue::createFromCapsule)
       .def_property_readonly(
           "context",
           [](PyValue &self) { return self.getParentOperation()->getContext(); },
           "Context in which the value lives.")
       .def(
           "dump", [](PyValue &self) { mlirValueDump(self.get()); },

mlir/lib/Bindings/Python/IRModule.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // 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 MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <optional>
 #include <utility>
 #include <vector>
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/IntegerSet.h"
 #include "llvm/ADT/DenseMap.h"
 

mlir/lib/Bindings/Python/IRModule.cpp

 //===- IRModule.cpp - IR pybind module ------------------------------------===//
 //
 // 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 "IRModule.h"
 #include "Globals.h"
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include <vector>
 #include <optional>
 
 #include "mlir-c/Bindings/Python/Interop.h"
 
 namespace py = pybind11;
 using namespace mlir;

mlir/lib/Bindings/Python/IRTypes.cpp

 //===- IRTypes.cpp - Exports builtin and standard types -------------------===//
 //
 // 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 "IRModule.h"
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include <optional>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;

mlir/lib/Bindings/Python/MainModule.cpp

 //===- MainModule.cpp - Main pybind module --------------------------------===//
 //
 // 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 <tuple>
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 #include "Globals.h"
 #include "IRModule.h"
 #include "Pass.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;

mlir/lib/Bindings/Python/Pass.h

 //===- Pass.h - PassManager Submodules of pybind module -------------------===//
 //
 // 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 MLIR_BINDINGS_PYTHON_PASS_H
 #define MLIR_BINDINGS_PYTHON_PASS_H
 
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 namespace mlir {
 namespace python {
 
 void populatePassManagerSubmodule(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_PASS_H

mlir/lib/Bindings/Python/PybindUtils.cpp

 //===- PybindUtils.cpp - Utilities for interop with pybind11 --------------===//
 //
 // 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 "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindUtils.h"
 
 pybind11::error_already_set
 mlir::python::SetPyError(PyObject *excClass, const llvm::Twine &message) {
   auto messageStr = message.str();
   PyErr_SetString(excClass, messageStr.c_str());
   return pybind11::error_already_set();
 }

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 TestAttr, TestType, TestTensorValue
 
 def register_python_test_dialect(context, load=True):
   from .._mlir_libs import _mlirPythonTest
   _mlirPythonTest.register_python_test_dialect(context, load)

mlir/test/python/dialects/python_test.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.python_test as test
+import mlir.dialects.arith as arith
+import mlir.dialects.tensor as tensor
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
 # CHECK-LABEL: TEST: testAttributes
 @run
with Context() as ctx:
     # The following must trigger a TypeError from pybind (therefore, not
     # checking its message) and must not crash.
     try:
       test.TestType(42, 56)
     except TypeError:
       pass
     else:
       raise
+
+
+@run
+# CHECK-LABEL: TEST: testTensorValue
+def testTensorValue():
+  with Context() as ctx, Location.unknown():
+    test.register_python_test_dialect(ctx)
+
+    i8 = IntegerType.get_signless(8)
+
+    class Tensor(test.TestTensorValue):
+      def __getitem__(self, dims):
+        dims = list(dims)
+        for i, d in enumerate(dims):
+          if isinstance(d, int):
+            dims[i] = arith.ConstantOp.create_index(d)
+
+        return tensor.ExtractOp(self, dims).result
+
+    module = Module.create()
+    with InsertionPoint(module.body):
+      t = tensor.EmptyOp([10, 10], i8).result
+
+      # CHECK: Value(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
+      print(Value(t))
+
+      tt = Tensor(t)
+      # CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
+      print(tt)
+
+      el = tt[0, 0]
+
+    # CHECK: module {
+    # CHECK:   %[[T0:.*]] = tensor.empty() : tensor<10x10xi8>
+    # CHECK:   %[[C0:.*]]  = arith.constant 0 : index
+    # CHECK:   %[[C00:.*]] = arith.constant 0 : index
+    # CHECK:   %[[RES:.*]] = tensor.extract %[[T0]][%[[C0]], %[[C00]]] : tensor<10x10xi8>
+    # CHECK: }
+    print(module)
 No newline at end of file

mlir/test/python/lib/PythonTestCAPI.h

 
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirPythonTestTestAttributeGet(MlirContext context);
 
 MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestType(MlirType type);
 
 MLIR_CAPI_EXPORTED MlirType mlirPythonTestTestTypeGet(MlirContext context);
 
+MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_TEST_PYTHON_LIB_PYTHONTESTCAPI_H

mlir/test/python/lib/PythonTestCAPI.cpp

 //===- PythonTestCAPI.cpp - C API 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 "PythonTestDialect.h"
+#include "mlir-c/BuiltinTypes.h"
 #include "mlir/CAPI/Registration.h"
 #include "mlir/CAPI/Wrap.h"
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(PythonTest, python_test,
                                       python_test::PythonTestDialect)
 
 bool mlirAttributeIsAPythonTestTestAttribute(MlirAttribute attr) {
   return unwrap(attr).isa<python_test::TestAttrAttr>();
 }
 
 MlirAttribute mlirPythonTestTestAttributeGet(MlirContext context) {
   return wrap(python_test::TestAttrAttr::get(unwrap(context)));
 }
 
 bool mlirTypeIsAPythonTestTestType(MlirType type) {
   return unwrap(type).isa<python_test::TestTypeType>();
 }
 
 MlirType mlirPythonTestTestTypeGet(MlirContext context) {
   return wrap(python_test::TestTypeType::get(unwrap(context)));
 }
+
+bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value) {
+  return mlirTypeIsATensor(wrap(unwrap(value).getType()));
+}
 No newline at end of file

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 "mlir/Bindings/Python/PybindUtils.h"
 
 namespace py = pybind11;
 using namespace mlir::python::adaptors;
 
 PYBIND11_MODULE(_mlirPythonTest, m) {
   m.def(
       "register_python_test_dialect",
       [](MlirContext context, bool load) {
mlir_attribute_subclass(m, "TestAttr",
           py::arg("cls"), py::arg("context") = py::none());
   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_value_subclass(m, "TestTensorValue",
+                      mlirTypeIsAPythonTestTestTensorValue)
+      .def("__str__", [](MlirValue &self) {
+        mlir::PyPrintAccumulator printAccum;
+        printAccum.parts.append("Tensor(");
+        mlirValuePrint(self, printAccum.getCallback(),
+                       printAccum.getUserData());
+        printAccum.parts.append(")");
+        return printAccum.join();
+      });
 }