[Mlir-commits] [mlir] [mlir:python] Construct PyOperation objects in-place on the Python heap. (PR #123813)

[Peter Hawkins]

https://github.com/hawkinsp created https://github.com/llvm/llvm-project/pull/123813

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.

>From 7d52b9a3646326456df95c6e9a0943d4ffc00402 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins at google.com>
Date: Tue, 21 Jan 2025 20:12:32 +0000
Subject: [PATCH] [mlir:python] Construct PyOperation objects in-place on the
 Python heap.

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.
---
 mlir/lib/Bindings/Python/IRCore.cpp | 28 ++++++++++++++++++++--------
 mlir/lib/Bindings/Python/IRModule.h |  3 ++-
 2 files changed, 22 insertions(+), 9 deletions(-)

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);