[Mlir-commits] [mlir] [mlir][python] Expose `PyInsertionPoint`'s reference operation (PR #69082)

[Tomás Longeri]

https://github.com/tlongeri updated https://github.com/llvm/llvm-project/pull/69082

>From 43820c47721cded1c90030ae30191a7dc3da5e25 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= <tlongeri at google.com>
Date: Fri, 13 Oct 2023 23:15:43 +0000
Subject: [PATCH] [mlir][python] Expose PyInsertionPoint's reference operation

Updates the insertion_point.py test with asserts for existing testcases
that check the reference operation and the block.
---
 mlir/lib/Bindings/Python/IRCore.cpp      | 13 ++++++++++++-
 mlir/lib/Bindings/Python/IRModule.h      |  1 +
 mlir/python/mlir/_mlir_libs/_mlir/ir.pyi |  2 ++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index c8373e06f0db776..389a4621c14e594 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3207,7 +3207,18 @@ 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) -> py::object {
+            auto ref_operation = self.getRefOperation();
+            if (ref_operation)
+              return ref_operation->getObject();
+            return py::none();
+          },
+          "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
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index e8f4440d216eeb4..2609117dd220bea 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -755,6 +755,8 @@ class InsertionPoint:
     def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ...
     @property
     def block(self) -> Block: ...
+    @property
+    def ref_operation(self) -> Optional[_OperationBase]: ...
 
 # TODO: Auto-generated. Audit and fix.
 class IntegerAttr(Attribute):