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

[llvmlistbot at llvm.org]

Author: Twice
Date: 2025-10-09T19:18:53+08:00
New Revision: b6bf196cdc6a71955dd90968ab96667899cebdb8

URL: https://github.com/llvm/llvm-project/commit/b6bf196cdc6a71955dd90968ab96667899cebdb8
DIFF: https://github.com/llvm/llvm-project/commit/b6bf196cdc6a71955dd90968ab96667899cebdb8.diff

LOG: [MLIR][Python] Make the TypeID allocator globally defined in `PassManager.add` (#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.

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/Globals.h
    mlir/lib/Bindings/Python/Pass.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index 71a051cb3d9f5..1e81f53e465ac 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/mlir/lib/Bindings/Python/Globals.h
@@ -17,6 +17,7 @@
 
 #include "NanobindUtils.h"
 #include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
 #include "mlir/CAPI/Support.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringExtras.h"
@@ -151,6 +152,29 @@ class PyGlobals {
 
   TracebackLoc &getTracebackLoc() { return tracebackLoc; }
 
+  class TypeIDAllocator {
+  public:
+    TypeIDAllocator() : allocator(mlirTypeIDAllocatorCreate()) {}
+    ~TypeIDAllocator() {
+      if (allocator.ptr)
+        mlirTypeIDAllocatorDestroy(allocator);
+    }
+    TypeIDAllocator(const TypeIDAllocator &) = delete;
+    TypeIDAllocator(TypeIDAllocator &&other) : allocator(other.allocator) {
+      other.allocator.ptr = nullptr;
+    }
+
+    MlirTypeIDAllocator get() { return allocator; }
+    MlirTypeID allocate() {
+      return mlirTypeIDAllocatorAllocateTypeID(allocator);
+    }
+
+  private:
+    MlirTypeIDAllocator allocator;
+  };
+
+  MlirTypeID allocateTypeID() { return typeIDAllocator.allocate(); }
+
 private:
   static PyGlobals *instance;
 
@@ -173,6 +197,7 @@ class PyGlobals {
   llvm::StringSet<> loadedDialectModules;
 
   TracebackLoc tracebackLoc;
+  TypeIDAllocator typeIDAllocator;
 };
 
 } // namespace python

diff  --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index e489585fd5f50..3c9f72a9d20d2 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -8,6 +8,7 @@
 
 #include "Pass.h"
 
+#include "Globals.h"
 #include "IRModule.h"
 #include "mlir-c/Pass.h"
 // clang-format off
@@ -181,9 +182,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 = PyGlobals::get().allocateTypeID();
             MlirExternalPassCallbacks callbacks;
             callbacks.construct = [](void *obj) {
               (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();