[Mlir-commits] [mlir] a365084 - [mlir][python, CAPI] expose Op::isBeforeInBlock (#150271)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 23 10:33:45 PDT 2025
Author: Maksim Levental
Date: 2025-07-23T12:33:42-05:00
New Revision: a36508483e1a36612175500d2dc4754f9aadacc2
URL: https://github.com/llvm/llvm-project/commit/a36508483e1a36612175500d2dc4754f9aadacc2
DIFF: https://github.com/llvm/llvm-project/commit/a36508483e1a36612175500d2dc4754f9aadacc2.diff
LOG: [mlir][python,CAPI] expose Op::isBeforeInBlock (#150271)
Added:
Modified:
mlir/include/mlir-c/IR.h
mlir/lib/Bindings/Python/IRCore.cpp
mlir/lib/Bindings/Python/IRModule.h
mlir/lib/CAPI/IR/IR.cpp
mlir/test/python/ir/operation.py
Removed:
################################################################################
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