[mlir][python,CAPI] expose Op::isBeforeInBlock (#150271)
This commit is contained in:
parent
e89e678839
commit
a36508483e
@ -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,
|
||||
|
@ -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) {
|
||||
|
@ -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();
|
||||
|
@ -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:
|
||||
|
@ -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
|
||||
|
Loading…
x
Reference in New Issue
Block a user