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

[llvmlistbot at llvm.org]

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

>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 1/3] [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();

>From 3b3c3f6b4f85eda3f8b1f0fb0659202a83ce8204 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Thu, 9 Oct 2025 14:11:31 +0800
Subject: [PATCH 2/3] move to PyGlobals

---
 mlir/lib/Bindings/Python/Globals.h | 20 ++++++++++++++++++++
 mlir/lib/Bindings/Python/Pass.cpp  | 20 ++------------------
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index 71a051cb3d9f5..6b6b4b598e366 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/mlir/lib/Bindings/Python/Globals.h
@@ -151,6 +151,25 @@ class PyGlobals {
 
   TracebackLoc &getTracebackLoc() { return tracebackLoc; }
 
+  class TypeIDAllocator {
+  public:
+    TypeIDAllocator() : allocator(mlirTypeIDAllocatorCreate()) {}
+    ~TypeIDAllocator() {
+      if (!allocator.ptr)
+        mlirTypeIDAllocatorDestroy(allocator);
+    }
+
+    MlirTypeIDAllocator get() { return allocator; }
+    MlirTypeID allocate() {
+      return mlirTypeIDAllocatorAllocateTypeID(allocator);
+    }
+
+  private:
+    MlirTypeIDAllocator allocator;
+  };
+
+  MlirTypeID allocateTypeID() { return typeIDAllocator.allocate(); }
+
 private:
   static PyGlobals *instance;
 
@@ -173,6 +192,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 e371a8cf76a9d..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
@@ -52,23 +53,6 @@ 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.
@@ -198,7 +182,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
               name = nb::cast<std::string>(
                   nb::borrow<nb::str>(run.attr("__name__")));
             }
-            MlirTypeID passID = globalTypeIDAllocator.allocate();
+            MlirTypeID passID = PyGlobals::get().allocateTypeID();
             MlirExternalPassCallbacks callbacks;
             callbacks.construct = [](void *obj) {
               (void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();

>From c336e54d8bb5be3e7773108abbbb841108e81423 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Thu, 9 Oct 2025 16:20:07 +0800
Subject: [PATCH 3/3] fix dtor

---
 mlir/lib/Bindings/Python/Globals.h | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index 6b6b4b598e366..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"
@@ -155,9 +156,13 @@ class PyGlobals {
   public:
     TypeIDAllocator() : allocator(mlirTypeIDAllocatorCreate()) {}
     ~TypeIDAllocator() {
-      if (!allocator.ptr)
+      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() {