[Mlir-commits] [mlir] [MLIR][Python] bind InsertionPointAfter (PR #157156)

[Maksim Levental]

https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/157156

None

>From 14053c1370ed9f6b560fd0ed3e0365fc3159aac4 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Fri, 5 Sep 2025 11:22:04 -0700
Subject: [PATCH] [MLIR][Python] bind InsertionPointAfter

---
 mlir/lib/Bindings/Python/IRCore.cpp    | 11 +++++++++++
 mlir/lib/Bindings/Python/IRModule.h    |  3 +++
 mlir/test/python/ir/insertion_point.py | 27 ++++++++++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index bf4950fc1a070..aaf13a67af2f6 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2073,6 +2073,15 @@ PyInsertionPoint PyInsertionPoint::atBlockTerminator(PyBlock &block) {
   return PyInsertionPoint{block, std::move(terminatorOpRef)};
 }
 
+PyInsertionPoint PyInsertionPoint::after(PyOperation &op) {
+  MlirOperation nextOperation = mlirOperationGetNextInBlock(op.get());
+  if (mlirOperationIsNull(nextOperation))
+    throw nb::value_error(
+        "Operation is terminator; can't set insertion point after terminator");
+  PyOperation nextOp{op.getContext(), nextOperation};
+  return PyInsertionPoint{nextOp};
+}
+
 nb::object PyInsertionPoint::contextEnter(nb::object insertPoint) {
   return PyThreadContextEntry::pushInsertionPoint(insertPoint);
 }
@@ -3861,6 +3870,8 @@ void mlir::python::populateIRCore(nb::module_ &m) {
                   nb::arg("block"), "Inserts at the beginning of the block.")
       .def_static("at_block_terminator", &PyInsertionPoint::atBlockTerminator,
                   nb::arg("block"), "Inserts before the block terminator.")
+      .def_static("after", &PyInsertionPoint::after, nb::arg("operation"),
+                  "Inserts after the operation.")
       .def("insert", &PyInsertionPoint::insert, nb::arg("operation"),
            "Inserts an operation.")
       .def_prop_ro(
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 0cc0459ebc9a0..24112985c5d27 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -829,6 +829,9 @@ class PyInsertionPoint {
   static PyInsertionPoint atBlockBegin(PyBlock &block);
   /// Shortcut to create an insertion point before the block terminator.
   static PyInsertionPoint atBlockTerminator(PyBlock &block);
+  /// Shortcut to create an insertion point to the node after the specified
+  /// operation.
+  static PyInsertionPoint after(PyOperation &op);
 
   /// Inserts an operation.
   void insert(PyOperationBase &operationBase);
diff --git a/mlir/test/python/ir/insertion_point.py b/mlir/test/python/ir/insertion_point.py
index 5eb861a2c0891..c6e6535bb014e 100644
--- a/mlir/test/python/ir/insertion_point.py
+++ b/mlir/test/python/ir/insertion_point.py
@@ -63,6 +63,33 @@ def test_insert_before_operation():
 run(test_insert_before_operation)
 
 
+# CHECK-LABEL: TEST: test_insert_before_operation
+def test_insert_before_operation():
+    ctx = Context()
+    ctx.allow_unregistered_dialects = True
+    with Location.unknown(ctx):
+        module = Module.parse(
+            r"""
+      func.func @foo() -> () {
+        "custom.op1"() : () -> ()
+        "custom.op2"() : () -> ()
+      }
+    """
+        )
+        entry_block = module.body.operations[0].regions[0].blocks[0]
+        ip = InsertionPoint.after(entry_block.operations[1])
+        assert ip.block == entry_block
+        assert ip.ref_operation == entry_block.operations[1]
+        ip.insert(Operation.create("custom.op3"))
+        # CHECK: "custom.op1"
+        # CHECK: "custom.op2"
+        # CHECK: "custom.op3"
+        module.operation.print()
+
+
+run(test_insert_before_operation)
+
+
 # CHECK-LABEL: TEST: test_insert_at_block_begin
 def test_insert_at_block_begin():
     ctx = Context()