[Mlir-commits] [mlir] [mlir][python, CAPI] expose Op::isBeforeInBlock (PR #150271)
Maksim Levental
llvmlistbot at llvm.org
Wed Jul 23 10:15:25 PDT 2025
https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/150271
None
>From 936f23553dbb430c9c537723b67cdb0efbef2737 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Wed, 23 Jul 2025 13:14:51 -0400
Subject: [PATCH] [mlir][python,CAPI] expose Op::isBeforeInBlock
---
mlir/include/mlir-c/IR.h | 7 +++++++
mlir/lib/Bindings/Python/IRCore.cpp | 15 +++++++++++++++
mlir/lib/Bindings/Python/IRModule.h | 7 +++++++
mlir/lib/CAPI/IR/IR.cpp | 4 ++++
mlir/test/python/ir/operation.py | 4 ++++
5 files changed, 37 insertions(+)
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 81299c7911d24..71c7d4378677f 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -813,6 +813,13 @@ MLIR_CAPI_EXPORTED void mlirOperationMoveAfter(MlirOperation op,
MLIR_CAPI_EXPORTED void mlirOperationMoveBefore(MlirOperation op,
MlirOperation other);
+/// Given an operation 'other' that is within the same parent block, return
+/// whether the current operation is before 'other' in the operation list
+/// of the parent block.
+/// Note: This function has an average complexity of O(1), but worst case may
+/// take O(N) where N is the number of operations within the parent block.
+MLIR_CAPI_EXPORTED bool mlirOperationIsBeforeInBlock(MlirOperation op,
+ MlirOperation other);
/// Operation walk result.
typedef enum MlirWalkResult {
MlirWalkResultAdvance,
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 7b790e90e0d87..5feed95f96f53 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1454,6 +1454,14 @@ void PyOperationBase::moveBefore(PyOperationBase &other) {
operation.parentKeepAlive = otherOp.parentKeepAlive;
}
+bool PyOperationBase::isBeforeInBlock(PyOperationBase &other) {
+ PyOperation &operation = getOperation();
+ PyOperation &otherOp = other.getOperation();
+ operation.checkValid();
+ otherOp.checkValid();
+ return mlirOperationIsBeforeInBlock(operation, otherOp);
+}
+
bool PyOperationBase::verify() {
PyOperation &op = getOperation();
PyMlirContext::ErrorCapture errors(op.getContext());
@@ -3409,6 +3417,13 @@ void mlir::python::populateIRCore(nb::module_ &m) {
.def("move_before", &PyOperationBase::moveBefore, nb::arg("other"),
"Puts self immediately before the other operation in its parent "
"block.")
+ .def("is_before_in_block", &PyOperationBase::isBeforeInBlock,
+ nb::arg("other"),
+ "Given an operation 'other' that is within the same parent block, "
+ "return"
+ "whether the current operation is before 'other' in the operation "
+ "list"
+ "of the parent block.")
.def(
"clone",
[](PyOperationBase &self, nb::object ip) {
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 0fdd2d1a7eff6..9c22dea157c06 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -624,6 +624,13 @@ class PyOperationBase {
void moveAfter(PyOperationBase &other);
void moveBefore(PyOperationBase &other);
+ /// Given an operation 'other' that is within the same parent block, return
+ /// whether the current operation is before 'other' in the operation list
+ /// of the parent block.
+ /// Note: This function has an average complexity of O(1), but worst case may
+ /// take O(N) where N is the number of operations within the parent block.
+ bool isBeforeInBlock(PyOperationBase &other);
+
/// Verify the operation. Throws `MLIRError` if verification fails, and
/// returns `true` otherwise.
bool verify();
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index fbc66bcf5c2d0..8491553dab76f 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -850,6 +850,10 @@ void mlirOperationMoveBefore(MlirOperation op, MlirOperation other) {
return unwrap(op)->moveBefore(unwrap(other));
}
+bool mlirOperationIsBeforeInBlock(MlirOperation op, MlirOperation other) {
+ return unwrap(op)->isBeforeInBlock(unwrap(other));
+}
+
static mlir::WalkResult unwrap(MlirWalkResult result) {
switch (result) {
case MlirWalkResultAdvance:
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index ede1571f940f6..c6b5dafe792ca 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -978,8 +978,12 @@ def testModuleMerge():
foo = m1.body.operations[0]
bar = m2.body.operations[0]
qux = m2.body.operations[1]
+ assert bar.is_before_in_block(qux)
bar.move_before(foo)
+ assert bar.is_before_in_block(foo)
qux.move_after(foo)
+ assert bar.is_before_in_block(qux)
+ assert foo.is_before_in_block(qux)
# CHECK: module
# CHECK: func private @bar
More information about the Mlir-commits
mailing list