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

Author: makslevental

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,

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

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();

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:

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