summaryrefslogtreecommitdiff
path: root/mlir/test/python/ir/context_lifecycle.py
AgeCommit message (Collapse)Author
2025-03-10[MLIR][py] Add PyThreadPool as wrapper around MlirLlvmThreadPool in MLIR ↵vfdev
python bindings (#130109) In some projects like JAX ir.Context are used with disabled multi-threading to avoid caching multiple threading pools: https://github.com/jax-ml/jax/blob/623865fe9538100d877ba9d36f788d0f95a11ed2/jax/_src/interpreters/mlir.py#L606-L611 However, when context has enabled multithreading it also uses locks on the StorageUniquers and this can be helpful to avoid data races in the multi-threaded execution (for example with free-threaded cpython, https://github.com/jax-ml/jax/issues/26272). With this PR user can enable the multi-threading: 1) enables additional locking and 2) set a shared threading pool such that cached contexts can have one global pool.
2023-12-21Revert "[mlir][python] Make the Context/Operation capsule creation methods ↵Alex Zinenko
work as documented. (#76010)" This reverts commit bbc29768683b394b34600347f46be2b8245ddb30. This change seems to be at odds with the non-owning part semantics of MlirOperation in C API. Since downstream clients can only take and return MlirOperation, it does not sound correct to force all returns of MlirOperation transfer ownership. Specifically, this makes it impossible for downstreams to implement IR-traversing functions that, e.g., look at neighbors of an operation. The following patch triggers the exception, and there does not seem to be an alternative way for a downstream binding writer to express this: ``` diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 39757dfad5be..2ce640674245 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -3071,6 +3071,11 @@ void mlir::python::populateIRCore(py::module &m) { py::arg("successors") = py::none(), py::arg("regions") = 0, py::arg("loc") = py::none(), py::arg("ip") = py::none(), py::arg("infer_type") = false, kOperationCreateDocstring) + .def("_get_first_in_block", [](PyOperation &self) -> MlirOperation { + MlirBlock block = mlirOperationGetBlock(self.get()); + MlirOperation first = mlirBlockGetFirstOperation(block); + return first; + }) .def_static( "parse", [](const std::string &sourceStr, const std::string &sourceName, diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py index f59b1a26ba48..6b12b8da5c24 100644 --- a/mlir/test/python/ir/operation.py +++ b/mlir/test/python/ir/operation.py @@ -24,6 +24,25 @@ def expect_index_error(callback): except IndexError: pass +@run +def testCustomBind(): + ctx = Context() + ctx.allow_unregistered_dialects = True + module = Module.parse( + r""" + func.func @f1(%arg0: i32) -> i32 { + %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32 + return %1 : i32 + } + """, + ctx, + ) + add = module.body.operations[0].regions[0].blocks[0].operations[0] + op = add.operation + # This will get a reference to itself. + f1 = op._get_first_in_block() + + # Verify iterator based traversal of the op/region/block hierarchy. # CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators ```
2023-12-20[mlir][python] Make the Context/Operation capsule creation methods work as ↵Stella Laurenzo
documented. (#76010) This fixes a longstanding bug in the `Context._CAPICreate` method whereby it was not taking ownership of the PyMlirContext wrapper when casting to a Python object. The result was minimally that all such contexts transferred in that way would leak. In addition, counter to the documentation for the `_CAPICreate` helper (see `mlir-c/Bindings/Python/Interop.h`) and the `forContext` / `forOperation` methods, we were silently upgrading any unknown context/operation pointer to steal-ownership semantics. This is dangerous and was causing some subtle bugs downstream where this facility is getting the most use. This patch corrects the semantics and will only do an ownership transfer for `_CAPICreate`, and it will further require that it is an ownership transfer (if already transferred, it was just silently succeeding). Removing the mis-aligned behavior made it clear where the downstream was doing the wrong thing. It also adds some `_testing_` functions to create unowned context and operation capsules so that this can be fully tested upstream, reworking the tests to verify the behavior. In some torture testing downstream, I was not able to trigger any memory corruption with the newly enforced semantics. When getting it wrong, a regular exception is raised.
2021-05-03Move MLIR python sources to mlir/python.Stella Laurenzo
* NFC but has some fixes for CMake glitches discovered along the way (things not cleaning properly, co-mingled depends). * Includes previously unsubmitted fix in D98681 and a TODO to fix it more appropriately in a smaller followup. Differential Revision: https://reviews.llvm.org/D101493