[mlir:python] Construct PyOperation objects in-place on the Python heap.

[hawkinsp]

Currently we make two memory allocations for each PyOperation: a Python object, and the PyOperation class itself. With some care we can allocate the PyOperation inline inside the Python object, saving us a malloc() call per object and perhaps improving cache locality.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Peter Hawkins (hawkinsp)

Changes

Currently we make two memory allocations for each PyOperation: a Python object, and the PyOperation class itself. With some care we can allocate the PyOperation inline inside the Python object, saving us a malloc() call per object and perhaps improving cache locality.

---

Full diff: https://github.com/llvm/llvm-project/pull/123813.diff

2 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRCore.cpp (+20-8)
• (modified) mlir/lib/Bindings/Python/IRModule.h (+2-1)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 53806ca9f04a49..2ded1a8df7c3e6 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1195,21 +1195,33 @@ PyOperation::~PyOperation() {
   }
 }
 
+namespace {
+
+// Constructs a new object of type T in-place on the Python heap, returning a
+// PyObjectRef to it, loosely analogous to std::make_shared<T>().
+template <typename T, class... Args>
+PyObjectRef<T> makeObjectRef(Args &&...args) {
+  nb::handle type = nb::type<T>();
+  nb::object instance = nb::inst_alloc(type);
+  T *ptr = nb::inst_ptr<T>(instance);
+  new (ptr) T(std::forward<Args>(args)...);
+  nb::inst_mark_ready(instance);
+  return PyObjectRef<T>(ptr, std::move(instance));
+}
+
+} // namespace
+
 PyOperationRef PyOperation::createInstance(PyMlirContextRef contextRef,
                                            MlirOperation operation,
                                            nb::object parentKeepAlive) {
   // Create.
-  PyOperation *unownedOperation =
-      new PyOperation(std::move(contextRef), operation);
-  // Note that the default return value policy on cast is automatic_reference,
-  // which does not take ownership (delete will not be called).
-  // Just be explicit.
-  nb::object pyRef = nb::cast(unownedOperation, nb::rv_policy::take_ownership);
-  unownedOperation->handle = pyRef;
+  PyOperationRef unownedOperation =
+      makeObjectRef<PyOperation>(std::move(contextRef), operation);
+  unownedOperation->handle = unownedOperation.getObject();
   if (parentKeepAlive) {
     unownedOperation->parentKeepAlive = std::move(parentKeepAlive);
   }
-  return PyOperationRef(unownedOperation, std::move(pyRef));
+  return unownedOperation;
 }
 
 PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef,
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index d1fb4308dbb77c..05c0e6f3a1ee0c 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -705,8 +705,9 @@ class PyOperation : public PyOperationBase, public BaseContextObject {
   /// Clones this operation.
   nanobind::object clone(const nanobind::object &ip);
 
-private:
   PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
+
+private:
   static PyOperationRef createInstance(PyMlirContextRef contextRef,
                                        MlirOperation operation,
                                        nanobind::object parentKeepAlive);

[hawkinsp]

I suggested upstreaming the helper here named makeObjectRef to nanobind here wjakob/nanobind#872, since it's a utility I've written several times at this point. We'll see what upstream says.

[jpienaar]

I'm not sure why this one was private ...