[mlir] Rework subclass construction in PybindAdaptors.h

[ftynse]

Summary

The constructor function was being defined without indicating its "init"
name, which made it interpret it as a regular fuction rather than a
constructor. When overload resolution failed, Pybind would attempt to print the
arguments actually passed to the function, including "self", which is not
initialized since the constructor couldn't be called. This would result in
"repr" being called with "self" referencing an uninitialized MLIR C API
object, which in turn would cause undefined behavior when attempting to print
in C++. Even if the correct name is provided, the mechanism used by
PybindAdaptors.h to bind constructors directly as "init" functions taking
"self" is deprecated by Pybind. The new mechanism does not seem to have access
to a fully-constructed "self" object (i.e., the constructor in C++ takes a
pybind11::detail::value_and_holder that cannot be forwarded back to Python).

Instead, redefine "new" to perform the required checks (there are no
additional initialization needed for attributes and types as they are all
wrappers around a C++ pointer). "new" can call its equivalent on a
superclass without needing "self".

Bump pybind11 dependency to 3.8.0, which is the first version that allows one
to redefine "new".

[ftynse]

Bump pybind11 requirement to 2.8.0.

[ftynse]

Add a standalone test for PybindAdaptors.h

[stellaraccident]

This makes sense to me, and I'm trying to remember why I didn't do it originally (I remember considering the new approach but can't remember why I avoided it). I think I was aliasing the prohibition that regular pybind classes cannot further modify new with these not actually being pybind classes.

In any case, thank you for the heroic debugging and fix.

mlir/cmake/modules/MLIRDetectPythonEnv.cmake

35 ↗

(On Diff #401158)

Can you include in the patch description the reason for the version bump? Based on offline discussion, it sounded like this version was chosen based on presence of something in the changelog. Would be good to capture that.

[ftynse]

This makes sense to me, and I'm trying to remember why I didn't do it originally (I remember considering the new approach but can't remember why I avoided it). I think I was aliasing the prohibition that regular pybind classes cannot further modify new with these not actually being pybind classes.

It was not possible to redefine __new__ when you wrote this, it was only added in September 2021 (https://github.com/pybind/pybind11/issues/3253), hence the version bump to the 3.8 release from October 2021.

mlir/cmake/modules/MLIRDetectPythonEnv.cmake

35 ↗

(On Diff #401158)

Yes, I updated the git commit but arc doesn't pick up new description. Amended manually.

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

mlir_attribute_subclass(py::handle scope, const char *attrClassName,
             py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
                 .attr("Attribute")) {}
 
   /// Subclasses with a provided mlir.ir.Attribute 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_attribute_subclass(py::handle scope, const char *typeClassName,
-                          IsAFunctionTy isaFunction,
-                          const py::object &superClass)
-      : pure_subclass(scope, typeClassName, superClass) {
-    // Casting constructor. Note that defining an __init__ method is special
-    // and not yet generalized on pure_subclass (it requires a somewhat
-    // different cpp_function and other requirements on chaining to super
-    // __init__ make it more awkward to do generally).
+                          IsAFunctionTy isaFunction, const py::object &superCls)
+      : pure_subclass(scope, typeClassName, 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 captureTypeName(
         typeClassName); // As string in case if typeClassName is not static.
-    py::cpp_function initCf(
-        [superClass, isaFunction, captureTypeName](py::object self,
-                                                   py::object otherType) {
-          MlirAttribute rawAttribute = py::cast<MlirAttribute>(otherType);
+    py::cpp_function newCf(
+        [superCls, isaFunction, captureTypeName](py::object cls,
+                                                 py::object otherAttribute) {
+          MlirAttribute rawAttribute = py::cast<MlirAttribute>(otherAttribute);
           if (!isaFunction(rawAttribute)) {
-            auto origRepr = py::repr(otherType).cast<std::string>();
+            auto origRepr = py::repr(otherAttribute).cast<std::string>();
             throw std::invalid_argument(
                 (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
                  " (from " + origRepr + ")")
                     .str());
           }
-          superClass.attr("__init__")(self, otherType);
+          py::object self = superCls.attr("__new__")(cls, otherAttribute);
+          return self;
         },
-        py::arg("cast_from_type"), py::is_method(py::none()),
-        "Casts the passed type to this specific sub-type.");
-    thisClass.attr("__init__") = initCf;
+        py::name("__new__"), py::arg("cls"), py::arg("cast_from_attr"));
+    thisClass.attr("__new__") = newCf;
 
     // 'isinstance' method.
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirAttribute other) { return isaFunction(other); },
         py::arg("other_attribute"));
   }
 };
mlir_type_subclass(py::handle scope, const char *typeClassName,
             scope, typeClassName, isaFunction,
             py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Type")) {}
 
   /// Subclasses with a provided mlir.ir.Type 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_type_subclass(py::handle scope, const char *typeClassName,
-                     IsAFunctionTy isaFunction, const py::object &superClass)
-      : pure_subclass(scope, typeClassName, superClass) {
-    // Casting constructor. Note that defining an __init__ method is special
-    // and not yet generalized on pure_subclass (it requires a somewhat
-    // different cpp_function and other requirements on chaining to super
-    // __init__ make it more awkward to do generally).
+                     IsAFunctionTy isaFunction, const py::object &superCls)
+      : pure_subclass(scope, typeClassName, 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 captureTypeName(
         typeClassName); // As string in case if typeClassName is not static.
-    py::cpp_function initCf(
-        [superClass, isaFunction, captureTypeName](py::object self,
-                                                   py::object otherType) {
+    py::cpp_function newCf(
+        [superCls, isaFunction, captureTypeName](py::object cls,
           MlirType rawType = py::cast<MlirType>(otherType);
           if (!isaFunction(rawType)) {
             auto origRepr = py::repr(otherType).cast<std::string>();
             throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
                                          captureTypeName + " (from " +
                                          origRepr + ")")
                                             .str());
           }
-          superClass.attr("__init__")(self, otherType);
+          py::object self = superCls.attr("__new__")(cls, otherType);
+          return self;
         },
-        py::arg("cast_from_type"), py::is_method(py::none()),
-        "Casts the passed type to this specific sub-type.");
-    thisClass.attr("__init__") = initCf;
+        py::name("__new__"), py::arg("cls"), py::arg("cast_from_type"));
+    thisClass.attr("__new__") = newCf;
 
     // 'isinstance' method.
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirType other) { return isaFunction(other); },
         py::arg("other_type"));
   }
 };
 
 } // namespace adaptors
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_PYBINDADAPTORS_H