[mlir] fix crash 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++.

Fix this by specifying the correct name.

This in turn uncovers the fact the the mechanism used by PybindAdaptors.h to
bind constructors directly as "init" functions taking "self" is deprecated
by Pybind. Instead, leverage the fact that the adaptors are intended for
attrbutes/types that cannot have additional data members and are all ultimately
instances of "PyAttribute"/"PyType" C++ class. In constructors of derived
classes, construct an instance of the base class first, then steal its internal
pointer to the C++ object to construct the instance of the derived class.

On top of that, the definition of the function was incorrectly indicated as the
method on the "None" object instead of being the method of its parent class.
This would result in a second problem when Pybind would attempt to print
warnings pointing to the parent class since the "None" does not have a
"name" field or its C API equivalent.

Fix this by specifying the correct parent class by looking it up by name in the
parent module.

[stellaraccident]

This is gross - I'm sorry it had caused some hard debugging for you, but thank you.

This use case may be worth an issue/discussion with the pybind folks: I don't have the details anymore but I didn't just make this up, having come across the approach as a suggestion from core pybind devs. Once we leave the fix, let's file an issue with them, describe what we are trying to do and get advice.

[ftynse]

This didn't work out in some variants. Let's try a different mechanism.

[ftynse]

Use pybind11 details to implement superclass initialization

[stellaraccident]

Ugh, ugh. Let's definitely ask the pybind folks after landing. I'm also wondering if this part should just be implemented with the python c API. The concept is simpler than the code.

[phawkins]

I've started seeing a nondeterministic crash in the JAX test suite that I bisected to this revision.

Here's a typical failure from our CI:
https://source.cloud.google.com/results/invocations/914b1d68-3b59-401c-8417-4080cc48cc35/targets/jax%2Ftesting%2Fcpu%2Fpresubmit_github/log

If you go and look at the python source lines that pytest reports the crash at, one is a call to .verify() and another is a call to .print(). I don't have a more reduced reproduction yet: it's going to be hard to get one because it's a rare crash that I only see when running a large test suite and not when running any tests in isolation. My guess is there's some sort of memory corruption problem.

Can we revert this in the meantime?

[stellaraccident]

In D117325#3251231, @phawkins wrote:

I've started seeing a nondeterministic crash in the JAX test suite that I bisected to this revision.

Here's a typical failure from our CI:
https://source.cloud.google.com/results/invocations/914b1d68-3b59-401c-8417-4080cc48cc35/targets/jax%2Ftesting%2Fcpu%2Fpresubmit_github/log

If you go and look at the python source lines that pytest reports the crash at, one is a call to .verify() and another is a call to .print(). I don't have a more reduced reproduction yet: it's going to be hard to get one because it's a rare crash that I only see when running a large test suite and not when running any tests in isolation. My guess is there's some sort of memory corruption problem.

Can we revert this in the meantime?

I'm fine reverting on such evidence.

My suspicion is a refcount/object lifetime issue, but I can't quite see the mechanic. Jax adds a lot of GC pressure and gil context switches that don't show up in isolation, and this would not be the first time it has triggered such a non deterministic failure due to a latent refcount bug. When I have reduced these in the past, I've done it by adding tests with explicit reference release and gc (look for "gc" in the tests to see examples).

There are only a few ways that this patch intersects with what Jax does: by building custom MHLO attributes (which iirc is just for convolution at the moment). I would isolate those code sequences and put them in a unit test with explicit GC used in anger to fuzz it.

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

pure_subclass &def_classmethod(const char *name, Func &&f,
     thisClass.attr(cf.name()) =
         py::reinterpret_borrow<py::object>(PyClassMethod_New(cf.ptr()));
     return *this;
   }
 
   py::object get_class() const { return thisClass; }
 
 protected:
+  /// Defers to the constructor of the superClass the configuration of the
+  /// pybind11 object from the given arguments. Pybind11 has a special handling
+  /// for constructors such that they don't accept a "self" reference, unlike
+  /// Python "__init__" calls. Therefore, one cannot just call the "__init__" of
+  /// the parent class, which would require access to "self". Instead, create an
+  /// instance of the superclass and take its instance pointer to the base C++
+  /// object to populate the instance pointer of the constructed object. Since
+  /// we only deal with _pure_ subclasses, this should be sufficient as derived
+  /// classes cannot have more data fields.
+  template <typename... Args>
+  static void deferToSuperclassConstructor(py::detail::value_and_holder &vh,
+                                           py::object superClass,
+                                           Args &&...args) {
+    py::object super = superClass(std::forward<Args>(args)...);
+    py::detail::type_info *ti =
+        py::detail::get_type_info((PyTypeObject *)superClass.ptr());
+    auto *instance = reinterpret_cast<py::detail::instance *>(super.ptr());
+
+    // Take ownership of the value pointer from the base class.
+    vh.value_ptr() = instance->get_value_and_holder(ti, true).value_ptr();
+    super.release();
+  }
+
   py::object superClass;
   py::object thisClass;
 };
 
 /// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
 /// constructor and type checking methods.
 class mlir_attribute_subclass : public pure_subclass {
 public:
public:
   /// 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).
+    // __init__ make it more awkward to do generally). It is marked as
+    // `is_new_style_constructor` to suppress the deprecation warning from
+    // pybind11 related to placement-new since we are not doing any allocation
+    // here but relying on the superclass constructor that does "new-style"
+    // allocation for pybind11.
     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) {
+        [superClass, isaFunction, captureTypeName](
+            py::detail::value_and_holder &vh, py::object otherType) {
           MlirAttribute rawAttribute = py::cast<MlirAttribute>(otherType);
           if (!isaFunction(rawAttribute)) {
             auto origRepr = py::repr(otherType).cast<std::string>();
             throw std::invalid_argument(
                 (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
                  " (from " + origRepr + ")")
                     .str());
           }
-          superClass.attr("__init__")(self, otherType);
+          pure_subclass::deferToSuperclassConstructor(vh, superClass,
+                                                      otherType);
         },
-        py::arg("cast_from_type"), py::is_method(py::none()),
+        py::name("__init__"), py::arg("cast_from_type"),
+        py::is_method(scope.attr(typeClassName)),
+        py::detail::is_new_style_constructor(),
         "Casts the passed type to this specific sub-type.");
     thisClass.attr("__init__") = initCf;
 
     // 'isinstance' method.
     def_staticmethod(
         "isinstance",
         [isaFunction](MlirAttribute other) { return isaFunction(other); },
         py::arg("other_attribute"));
public:
   /// 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).
+    // __init__ make it more awkward to do generally). It is marked as
+    // `is_new_style_constructor` to suppress the deprecation warning from
+    // pybind11 related to placement-new since we are not doing any allocation
+    // here but relying on the superclass constructor that does "new-style"
+    // allocation for pybind11.
     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) {
+        [superClass, isaFunction, captureTypeName](
+            py::detail::value_and_holder &vh, py::object otherType) {
           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);
+          pure_subclass::deferToSuperclassConstructor(vh, superClass,
+                                                      otherType);
         },
-        py::arg("cast_from_type"), py::is_method(py::none()),
+        py::name("__init__"), py::arg("cast_from_type"),
+        py::is_method(scope.attr(typeClassName)),
+        py::detail::is_new_style_constructor(),
         "Casts the passed type to this specific sub-type.");
     thisClass.attr("__init__") = initCf;
 
     // '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