[mlir][python] Expose PyInsertionPoint's reference operation

[tlongeri]

The reason I want this is that I am writing my own Python bindings and would like to use the insertion point from PyThreadContextEntry::getDefaultInsertionPoint() to call C++ functions that take an OpBuilder (I don't need to expose it in Python but it also seems appropriate). AFAICT, there is currently no way to translate a PyInsertionPoint into an OpBuilder because the operation is inaccessible.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Tomás Longeri (tlongeri)

Changes

The reason I want this is that I am writing my own Python bindings and would like to use the insertion point from PyThreadContextEntry::getDefaultInsertionPoint() to call C++ functions that take an OpBuilder (I don't need to expose it in Python but it also seems appropriate). AFAICT, there is currently no way to translate a PyInsertionPoint into an OpBuilder because the operation is inaccessible.

---

Full diff: https://github.com/llvm/llvm-project/pull/69082.diff

2 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRCore.cpp (+7-1)
• (modified) mlir/lib/Bindings/Python/IRModule.h (+1)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index c8373e06f0db776..4d86ab2d9583834 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -3207,7 +3207,13 @@ void mlir::python::populateIRCore(py::module &m) { "Inserts an operation.") .def_property_readonly( "block", [](PyInsertionPoint &self) { return self.getBlock(); }, - "Returns the block that this InsertionPoint points to."); + "Returns the block that this InsertionPoint points to.") + .def_property_readonly( + "ref_operation", + [](PyInsertionPoint &self) { return self.getRefOperation(); }, + "The reference operation before which new operations are " + "inserted, or None if the insertion point is at the end of " + "the block"); //---------------------------------------------------------------------------- // Mapping of PyAttribute. diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h index 3ca7dd851961a46..c5412e735dddcb5 100644 --- a/mlir/lib/Bindings/Python/IRModule.h +++ b/mlir/lib/Bindings/Python/IRModule.h @@ -833,6 +833,7 @@ class PyInsertionPoint { const pybind11::object &excTb); PyBlock &getBlock() { return block; } + std::optional<PyOperationRef> &getRefOperation() { return refOperation; } private: // Trampoline constructor that avoids null initializing members while 

[ftynse]

Consider updating https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi#L743 as well to provide type hints.

[tlongeri]

Consider updating https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi#L743 as well to provide type hints.

Done!

Also fixed the lambda in def_property_readonly (it was broken and I hadn't tested it properly)

[tlongeri]

Updated the insertion_point.py test as well. It now has asserts checking block and ref_operation on the previously existing checks. I was going to add separate testcases for it but I realized the MLIR would've just been identical to the already existing ones.