[Mlir-commits] [mlir] [MLIR][Python] Make the TypeID allocator global defined in `PassManager.add` (PR #162594)

[llvmlistbot at llvm.org]

https://github.com/PragmaTwice created https://github.com/llvm/llvm-project/pull/162594

Previously, each time we called `PassManager.add(python_pass_callable)`, a new `TypeID` allocator was created and never released afterward. This approach could potentially lead to some issues. In this PR, we introduce a global `TypeIDAllocator` that is shared across all `add` calls to allocate IDs.


>From 7b4be6c4d2a52d2a7b1d1786044c13b25a232e22 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Thu, 9 Oct 2025 11:52:45 +0800
Subject: [PATCH] [MLIR][Python] Make the TypeID allocator global defined in
 PassManager.add

---
 mlir/lib/Bindings/Python/Pass.cpp | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index e489585fd5f50..e371a8cf76a9d 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -52,6 +52,23 @@ class PyPassManager {
   MlirPassManager passManager;
 };
 
+class PyTypeIDAllocator {
+public:
+  PyTypeIDAllocator() : allocator(mlirTypeIDAllocatorCreate()) {}
+  ~PyTypeIDAllocator() {
+    if (!allocator.ptr)
+      mlirTypeIDAllocatorDestroy(allocator);
+  }
+
+  MlirTypeIDAllocator get() { return allocator; }
+  MlirTypeID allocate() { return mlirTypeIDAllocatorAllocateTypeID(allocator); }
+
+private:
+  MlirTypeIDAllocator allocator;
+};
+
+PyTypeIDAllocator globalTypeIDAllocator;
+
 } // namespace
 
 /// Create the `mlir.passmanager` here.
@@ -181,9 +198,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
               name = nb::cast<std::string>(
                   nb::borrow<nb::str>(run.attr("__name__")));
             }
-            MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
-            MlirTypeID passID =
-                mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
+            MlirTypeID passID = globalTypeIDAllocator.allocate();
             MlirExternalPassCallbacks callbacks;
             callbacks.construct = [](void *obj) {
               (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();