[mlir][python,CAPI] expose Op::isBeforeInBlock (#150271)

This commit is contained in:
Maksim Levental 2025-07-23 12:33:42 -05:00 committed by GitHub
parent e89e678839
commit a36508483e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 37 additions and 0 deletions

View File

@ -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,

View File

@ -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) {

View File

@ -624,6 +624,13 @@ public:
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();

View File

@ -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:

View File

@ -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