Skip to content

Commit 7b4be6c

Browse files
committed
[MLIR][Python] Make the TypeID allocator global defined in PassManager.add
1 parent bf34b2e commit 7b4be6c

File tree

1 file changed

+18
-3
lines changed

1 file changed

+18
-3
lines changed

mlir/lib/Bindings/Python/Pass.cpp

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,23 @@ class PyPassManager {
5252
MlirPassManager passManager;
5353
};
5454

55+
class PyTypeIDAllocator {
56+
public:
57+
PyTypeIDAllocator() : allocator(mlirTypeIDAllocatorCreate()) {}
58+
~PyTypeIDAllocator() {
59+
if (!allocator.ptr)
60+
mlirTypeIDAllocatorDestroy(allocator);
61+
}
62+
63+
MlirTypeIDAllocator get() { return allocator; }
64+
MlirTypeID allocate() { return mlirTypeIDAllocatorAllocateTypeID(allocator); }
65+
66+
private:
67+
MlirTypeIDAllocator allocator;
68+
};
69+
70+
PyTypeIDAllocator globalTypeIDAllocator;
71+
5572
} // namespace
5673

5774
/// Create the `mlir.passmanager` here.
@@ -181,9 +198,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
181198
name = nb::cast<std::string>(
182199
nb::borrow<nb::str>(run.attr("__name__")));
183200
}
184-
MlirTypeIDAllocator typeIDAllocator = mlirTypeIDAllocatorCreate();
185-
MlirTypeID passID =
186-
mlirTypeIDAllocatorAllocateTypeID(typeIDAllocator);
201+
MlirTypeID passID = globalTypeIDAllocator.allocate();
187202
MlirExternalPassCallbacks callbacks;
188203
callbacks.construct = [](void *obj) {
189204
(void)nb::handle(static_cast<PyObject *>(obj)).inc_ref();

0 commit comments

Comments
 (0)