[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

306–344

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

45

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

333

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.Value, 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/IRCore.cpp

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/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.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 __str__(self):
+        return super().__str__().replace("Value", "Tensor")
+
+    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)
+
+      # CHECK: False
+      print(tt.is_null())

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()));
+}

mlir/test/python/lib/PythonTestModule.cpp

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("is_null", [](MlirValue &self) { return mlirValueIsNull(self); });
 }