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( nb::borrow(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(obj)).inc_ref();