[mlir][python] allow for detaching operations from a block

[ftynse]

Summary

Provide support for removing an operation from the block that contains it and
moving it back to detached state. This allows for the operation to be moved to
a different block, a common IR manipulation for, e.g., module merging.

Also fix a potential one-past-end iterator dereference in Operation::moveAfter
discovered in the process.

[jpienaar]

Nice

mlir/include/mlir-c/IR.h

350

Could the same comment "style" as in line 342 be used here? (I think if you s/Creates/Removes/ and s/inserted/destroyed/ it works here and then it is formulaic how one reads these :-))

mlir/lib/Bindings/Python/IRCore.cpp

2214

Why not removeFromParent to keep it consistent with python and C++ method on Operation?

mlir/lib/Bindings/Python/IRModule.h

435

operation

[mehdi_amini]

This LGTM, but we may not want to encourage this for the use-case here necessarily: seems a bit low-level.
For example, if it is was native, the code sample in the test should use instead op->moveBefore(Operation *): it seems slightly less low-level, and also allows for inserting anywhere in the block (in particular when all you have is another operation, it also works with block->end())

Another one that is very suitable to "merging" is the splice API exposed by Blocks. It seems more efficient (and provides again a higher-level / more encapsulated API) to just splice a block into another block than manually iterating and moving operations.

[ftynse]

Address review.

[ftynse]

In D112700#3094993, @mehdi_amini wrote:

This LGTM, but we may not want to encourage this for the use-case here necessarily: seems a bit low-level.
For example, if it is was native, the code sample in the test should use instead op->moveBefore(Operation *): it seems slightly less low-level, and also allows for inserting anywhere in the block (in particular when all you have is another operation, it also works with block->end())

IMO, moveBefore/moveAfter is at exactly same level of abstraction as appending to the block. I added those for convenience, but append is still necessary to be able to add operations to an empty block. There is no equivalent of block->end() in C, let alone Python, because we are not modeling iterators as objects. In Python, it would be quite confusing to have C++-style iterators because Python has its own, different notion of iterators.

Another one that is very suitable to "merging" is the splice API exposed by Blocks. It seems more efficient (and provides again a higher-level / more encapsulated API) to just splice a block into another block than manually iterating and moving operations.

My personal, possibility controversial, opinion is that splice is a horrible API that it is also _lower_ level than append/moveAfter/moveBefore. It requires the caller to know that Blocks are linked lists, to access the underlying list object of the block (Block.getOperations(), no equivalent in C or Python, one is probably undesirable in Python), to get the list iterator out of operation (no equivalent in C or Python, again probably undesirable in Python), and has several undocumented overloads. Its semantics is to move linked list nodes pointed to by iterators from one list to another, which is at a lower level of abstraction than moving one operation before/after another that only reasons about operations, not their storage. I've never been able to call it correctly from the first attempt, and I would really like to avoid propagating it to Python. It should be possible to design a more pythonic approach of operating on slices (in Python sense) of consecutive operations, but it requires careful consideration of what we want to expose at the C level to support that, which goes beyond what I want here.

Merging modules specifically requires to iterate over individual operations anyway for symbol name conflict resolution purposes.

mlir/lib/Bindings/Python/IRCore.cpp

2214

Because bindings and their comments use the term "detached operation" to indicate operations that do not belong to a block (and are therefore owned by the caller).

[jpienaar]

mlir/lib/Bindings/Python/IRCore.cpp

2214

Isnt Python a binding here too?

[mehdi_amini]

In D112700#3096768, @ftynse wrote:

IMO, moveBefore/moveAfter is at exactly same level of abstraction as appending to the block.

I think you're missing something that I see as important: removeFromParent. That makes a move a higher abstraction than "detach-reattach".
(It's really the fact that the user must detach manually that triggers me here)

I added those for convenience, but append is still necessary to be able to add operations to an empty block.

Sure, but you can have append be more like op->moveToBlockEnd(block) which provides a superset of the same feature by also handling the detach.

There is no equivalent of block->end() in C, let alone Python, because we are not modeling iterators as objects. In Python, it would be quite confusing to have C++-style iterators because Python has its own, different notion of iterators.

Yes, I wasn't suggesting adding iterators, I rather read something like moveToBlockEnd anyway.

Another one that is very suitable to "merging" is the splice API exposed by Blocks. It seems more efficient (and provides again a higher-level / more encapsulated API) to just splice a block into another block than manually iterating and moving operations.

My personal, possibility controversial, opinion is that splice is a horrible API that it is also _lower_ level than append/moveAfter/moveBefore.

I don't quite get your point about lower-level here? Splicing two block is conceptually a loop that calls these other APIs, which makes it a strictly higher level construct to me.

It requires the caller to know that Blocks are linked lists

How so?
The caller of: block->getOperations().splice(block->end(), block->getOperations()); does not need to know anything about linked list. We could change the block storage to be a vector and the API would stand.

to access the underlying list object of the block (Block.getOperations(), no equivalent in C or Python, one is probably undesirable in Python), to get the list iterator out of operation (no equivalent in C or Python, again probably undesirable in Python), and has several undocumented overloads.

I think there are two things that gets mixed up here: the concept of splice and the way it is exposed in C++ on the Block class.

I am more interested in the concept than the C++ API which could benefit a forwarding API on block.

Its semantics is to move linked list nodes pointed to by iterators from one list to another,

No: this is the implementation, not the semantics. The semantics of a splice should be "Transfers elements container A, inserting them at position B."

The fact is that because it is a higher level API than moving individual element, it allows for a faster implementation leveraging internal details of the containers (which is what you describe).

which is at a lower level of abstraction than moving one operation before/after another that only reasons about operations, not their storage

Again: sorry but I don't understand you point about lower/higher level of abstraction.

It should be possible to design a more pythonic approach of operating on slices (in Python sense) of consecutive operations, but it requires careful consideration of what we want to expose at the C level to support that, which goes beyond what I want here.

I understand what you want, but I object that this is not what is desirable for a Python abstraction, so I'd really like a more careful consideration of the Pythonic API based on higher-level constructs (move, splice/concat, etc.) than "detach" and "attach".

Merging modules specifically requires to iterate over individual operations anyway for symbol name conflict resolution purposes.

Right, but an ideal merging module API for a user would be: module.mergeInto(otherModule) anyway.

[mehdi_amini]

I've been wondering if we use different definitions of higher-level / low-level, may be worth clarifying otherwise we'll talk past each others. So I went back to basic and check that I'm looking at it roughly from the general intro of the wikipedia page here: https://en.wikipedia.org/wiki/High-_and_low-level

High-level describe those operations that are more abstract in nature; wherein the overall goals and systemic features are typically more concerned with the wider, macro system as a whole.
Low-level describes more specific individual components of a systematic operation, focusing on the details of rudimentary micro functions rather than macro, complex processes. Low-level classification is typically more concerned with individual components within the system and how they operate.

There are two axis of "low-level" vs "high-level" that triggered me here:

1. Having use remove_from_parent and then append instead of a single API call.

Here you expose to python the Operation::remove() method, while I don't have an objection to expose this method, we should note that it has exactly 2 uses in tree right now. It is extremely unusual that we have to just manually "detach" an operation like that. The fact that a user in Python would have to do this as a prerequisite to accomplish something in MLIR is an important red flag to me. The only use of this API should be to detach without knowing where to re-attach (for example extracting a nested module to make it a standalone entity).

2. Having to loop over a block to move individually each operation instead of operating on the range on a single API call.

That's the splice discussion, and maybe it got derailed because C++ is a language that design API with leaky low-level implementation choice (not necessarily for bad reason, but still).
I suspect the discussion got derailed because you anchored to the "leaky" aspect of "splice" instead of the abstract way.

Let's forget the name "splice" or the C++ API on ilist and just look at the "concepts" we expose to our users: "merge blocks"

Also, my intuition would be to set a different bar between the level of APIs exposed to the C API users compared to Python users: adding a high-level API to the C API surface likely needs a good justification, but there is rarely a reason to not wrap a sequence of C API calls behind a higher-level API exposed by the Python bindings to its users. (made up example: we may or may not want to add mlirBlockMergeIntoOtherBlock(Block *src, Block *dst) but it does not mean we can't expose a block.mergeInto(destBlock) Python-level API on the python Block class).

mlir/lib/Bindings/Python/IRCore.cpp

2214

Because bindings and their comments use the term "detached operation" t

So why remove_from_parent instead of detach_from_parent here?

[ftynse]

Address some review.

[ftynse]

As a matter of fact, we seem to agree on most points. Having a possibility to move+append in one call similarly to moveBefore/After makes sense, I updated the code.

By higher-level API in this particular case, I understand operating on the concept of a block without having to know how it is organized. The two-argument version of the splice API you referenced is okay-ish in that sense, the others (https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/ADT/ilist.h#L333-L346) less so. The first issue, relevant for all versions of splice, is having to operate on block->getOperations() instead of block directly. The second, relevant for the overloads with 3+ arguments, is the need to pass the operations ilist of the block that currently contains the operations. This is only necessary because we operate on ilist, whose element iterators don't have a reference back to the container presumably to reduce memory consumption. An Operation actually has such a reference. The truly high-level API would let us indeed write block->mergeInto(other) instead of other->getOperations().splice(other.end(), block->getOperations()) and something like block->mergeIntoBefore(operation, other->front()) instead of other->getOperations().splice(operation.getIterator(), other->getOperations(), other->front()). This adds an abstraction on top of what is provided by ilist, otherwise the caller is really operating on ilist of operations, not on a block.

I agree that we need some sort of bulk operation movement API in Python, but, like I mentioned above, it needs careful consideration and is much better off as a patch and discussion. One concern there is the amount of non-trivial code that we put into bindings, and the runtime+memory complexity it might entail, as opposed to writing it in C++ or in Python from primitives.

[mehdi_amini]

mlir/lib/Bindings/Python/IRCore.cpp

2425

Nit: spurious braces?

[ftynse]

Address remaining review.

mlir/include/mlir-c/IR.h

 
 /// Creates a deep copy of an operation. The operation is not inserted and
 /// ownership is transferred to the caller.
 MLIR_CAPI_EXPORTED MlirOperation mlirOperationClone(MlirOperation op);
 
 /// Takes an operation owned by the caller and destroys it.
 MLIR_CAPI_EXPORTED void mlirOperationDestroy(MlirOperation op);
 
+/// Removes the given operation from its parent block. The operation is not
+/// destroyed. The ownership of the operation is transferred to the caller.
+MLIR_CAPI_EXPORTED void mlirOperationRemoveFromParent(MlirOperation op);
+
 /// Checks whether the underlying operation is null.
 static inline bool mlirOperationIsNull(MlirOperation op) { return !op.ptr; }
 
 /// Checks whether two operation handles point to the same operation. This does
 /// not perform deep comparison.
 MLIR_CAPI_EXPORTED bool mlirOperationEqual(MlirOperation op,
                                            MlirOperation other);
 
MLIR_CAPI_EXPORTED void mlirOperationPrintWithFlags(MlirOperation op,
                                                     void *userData);
 
 /// Prints an operation to stderr.
 MLIR_CAPI_EXPORTED void mlirOperationDump(MlirOperation op);
 
 /// Verify the operation and return true if it passes, false if it fails.
 MLIR_CAPI_EXPORTED bool mlirOperationVerify(MlirOperation op);
 
+/// Moves the given operation immediately after the other operation in its
+/// parent block. The given operation may be owned by the caller or by its
+/// current block. The other operation must belong to a block. In any case, the
+/// ownership is transferred to the block of the other operation.
+MLIR_CAPI_EXPORTED void mlirOperationMoveAfter(MlirOperation op,
+                                               MlirOperation other);
+
+/// Moves the given operation immediately before the other operation in its
+/// parent block. The given operation may be owner by the caller or by its
+/// current block. The other operation must belong to a block. In any case, the
+/// ownership is transferred to the block of the other operation.
+MLIR_CAPI_EXPORTED void mlirOperationMoveBefore(MlirOperation op,
+                                                MlirOperation other);
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//
 
 /// Creates a new empty region and transfers ownership to the caller.
 MLIR_CAPI_EXPORTED MlirRegion mlirRegionCreate();
 
 /// Takes a region owned by the caller and destroys it.

mlir/lib/Bindings/Python/IRCore.cpp

print(fileObject, /*binary=*/binary,
         /*enableDebugInfo=*/enableDebugInfo,
         /*prettyDebugInfo=*/prettyDebugInfo,
         /*printGenericOpForm=*/printGenericOpForm,
         /*useLocalScope=*/useLocalScope);
 
   return fileObject.attr("getvalue")();
 }
 
+void PyOperationBase::moveAfter(PyOperationBase &other) {
+  PyOperation &operation = getOperation();
+  PyOperation &otherOp = other.getOperation();
+  operation.checkValid();
+  otherOp.checkValid();
+  mlirOperationMoveAfter(operation, otherOp);
+  operation.parentKeepAlive = otherOp.parentKeepAlive;
+}
+
+void PyOperationBase::moveBefore(PyOperationBase &other) {
+  PyOperation &operation = getOperation();
+  PyOperation &otherOp = other.getOperation();
+  operation.checkValid();
+  otherOp.checkValid();
+  mlirOperationMoveBefore(operation, otherOp);
+  operation.parentKeepAlive = otherOp.parentKeepAlive;
+}
+
 llvm::Optional<PyOperationRef> PyOperation::getParentOperation() {
   checkValid();
   if (!isAttached())
     throw SetPyError(PyExc_ValueError, "Detached operations have no parent");
   MlirOperation operation = mlirOperationGetParentOperation(get());
   if (mlirOperationIsNull(operation))
     return {};
   return PyOperation::forOperation(getContext(), operation);
py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false, kOperationGetAsmDocstring)
       .def(
           "verify",
           [](PyOperationBase &self) {
             return mlirOperationVerify(self.getOperation());
           },
           "Verify the operation and return true if it passes, false if it "
-          "fails.");
+          "fails.")
+      .def("move_after", &PyOperationBase::moveAfter, py::arg("other"),
+           "Puts self immediately after the other operation in its parent "
+           "block.")
+      .def("move_before", &PyOperationBase::moveBefore, py::arg("other"),
+           "Puts self immediately before the other operation in its parent "
+           "block.")
+      .def(
+          "detach_from_parent",
+          [](PyOperationBase &self) {
+            PyOperation &operation = self.getOperation();
+            operation.checkValid();
+            if (!operation.isAttached())
+              throw py::value_error("Detached operation has no parent.");
+
+            operation.detachFromParent();
+            return operation.createOpView();
+          },
+          "Detaches the operation from its parent block.");
 
   py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
       .def_static("create", &PyOperation::create, py::arg("name"),
                   py::arg("results") = py::none(),
                   py::arg("operands") = py::none(),
                   py::arg("attributes") = py::none(),
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
py::class_<PyBlock>(m, "Block", py::module_local())
           "__str__",
           [](PyBlock &self) {
             self.checkValid();
             PyPrintAccumulator printAccum;
             mlirBlockPrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             return printAccum.join();
           },
-          "Returns the assembly form of the block.");
+          "Returns the assembly form of the block.")
+      .def(
+          "append",
+          [](PyBlock &self, PyOperationBase &operation) {
+            if (operation.getOperation().isAttached())
+              operation.getOperation().detachFromParent();
+
+            MlirOperation mlirOperation = operation.getOperation().get();
+            mlirBlockAppendOwnedOperation(self.get(), mlirOperation);
+            operation.getOperation().setAttached(
+                self.getParentOperation().getObject());
+          },
+          "Appends an operation to this block. If the operation is currently "
+          "in another block, it will be moved.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyInsertionPoint.
   //----------------------------------------------------------------------------
 
   py::class_<PyInsertionPoint>(m, "InsertionPoint", py::module_local())
       .def(py::init<PyBlock &>(), py::arg("block"),
            "Inserts after the last operation but still inside the block.")

mlir/lib/Bindings/Python/IRModule.h

public:
   void print(pybind11::object fileObject, bool binary,
              llvm::Optional<int64_t> largeElementsLimit, bool enableDebugInfo,
              bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope);
   pybind11::object getAsm(bool binary,
                           llvm::Optional<int64_t> largeElementsLimit,
                           bool enableDebugInfo, bool prettyDebugInfo,
                           bool printGenericOpForm, bool useLocalScope);
 
+  /// Moves the operation before or after the other operation.
+  void moveAfter(PyOperationBase &other);
+  void moveBefore(PyOperationBase &other);
+
   /// Each must provide access to the raw Operation.
   virtual PyOperation &getOperation() = 0;
 };
 
 /// Wrapper around PyOperation.
 /// Operations exist in either an attached (dependent) or detached (top-level)
 /// state. In the detached state (as on creation), an operation is owned by
 /// the creator and its lifetime extends either until its reference count
forOperation(PyMlirContextRef contextRef, MlirOperation operation,
                pybind11::object parentKeepAlive = pybind11::object());
 
   /// Creates a detached operation. The operation must not be associated with
   /// any existing live operation.
   static PyOperationRef
   createDetached(PyMlirContextRef contextRef, MlirOperation operation,
                  pybind11::object parentKeepAlive = pybind11::object());
 
+  /// Detaches the operation from its parent block and updates its state
+  /// accordingly.
+  void detachFromParent() {
+    mlirOperationRemoveFromParent(getOperation());
+    setDetached();
+    parentKeepAlive = pybind11::object();
+  }
+
   /// Gets the backing operation.
   operator MlirOperation() const { return get(); }
   MlirOperation get() const {
     checkValid();
     return operation;
   }
 
   PyOperationRef getRef() {
     return PyOperationRef(
         this, pybind11::reinterpret_borrow<pybind11::object>(handle));
   }
 
   bool isAttached() { return attached; }
-  void setAttached() {
+  void setAttached(pybind11::object parent = pybind11::object()) {
     assert(!attached && "operation already attached");
     attached = true;
   }
+  void setDetached() {
+    assert(attached && "operation already detached");
+    attached = false;
+  }
   void checkValid() const;
 
   /// Gets the owning block or raises an exception if the operation has no
   /// owning block.
   PyBlock getBlock();
 
   /// Gets the parent operation or raises an exception if the operation has
   /// no parent.
private:
   // Module.
   // TODO: As implemented, this facility is only sufficient for modeling the
   // trivial module parent back-reference. Generalize this to also account for
   // transitions from detached to attached and address TODOs in the
   // ir_operation.py regarding testing corresponding lifetime guarantees.
   pybind11::object parentKeepAlive;
   bool attached = true;
   bool valid = true;
+
+  friend class PyOperationBase;
 };
 
 /// A PyOpView is equivalent to the C++ "Op" wrappers: these are the basis for
 /// providing more instance-specific accessors and serve as the base class for
 /// custom ODS-style operation classes. Since this class is subclass on the
 /// python side, it must present an __init__ method that operates in pure
 /// python types.
 class PyOpView : public PyOperationBase {

mlir/lib/CAPI/IR/IR.cpp

 }
 
 MlirOperation mlirOperationClone(MlirOperation op) {
   return wrap(unwrap(op)->clone());
 }
 
 void mlirOperationDestroy(MlirOperation op) { unwrap(op)->erase(); }
 
+void mlirOperationRemoveFromParent(MlirOperation op) { unwrap(op)->remove(); }
+
 bool mlirOperationEqual(MlirOperation op, MlirOperation other) {
   return unwrap(op) == unwrap(other);
 }
 
 MlirContext mlirOperationGetContext(MlirOperation op) {
   return wrap(unwrap(op)->getContext());
 }
 

 }
 
 void mlirOperationDump(MlirOperation op) { return unwrap(op)->dump(); }
 
 bool mlirOperationVerify(MlirOperation op) {
   return succeeded(verify(unwrap(op)));
 }
 
+void mlirOperationMoveAfter(MlirOperation op, MlirOperation other) {
+  return unwrap(op)->moveAfter(unwrap(other));
+}
+
+void mlirOperationMoveBefore(MlirOperation op, MlirOperation other) {
+  return unwrap(op)->moveBefore(unwrap(other));
+}
+
 //===----------------------------------------------------------------------===//
 // Region API.
 //===----------------------------------------------------------------------===//
 
 MlirRegion mlirRegionCreate() { return wrap(new Region); }
 
 bool mlirRegionEqual(MlirRegion region, MlirRegion other) {
   return unwrap(region) == unwrap(other);

mlir/lib/IR/Operation.cpp

void Operation::moveAfter(Operation *existingOp) {
   moveAfter(existingOp->getBlock(), existingOp->getIterator());
 }
 
 /// Unlink this operation from its current block and insert it right after
 /// `iterator` in the specified block.
 void Operation::moveAfter(Block *block,
                           llvm::iplist<Operation>::iterator iterator) {
   assert(iterator != block->end() && "cannot move after end of block");
-  moveBefore(&*std::next(iterator));
+  moveBefore(block, std::next(iterator));
 }
 
 /// This drops all operand uses from this operation, which is an essential
 /// step in breaking cyclic dependences between references when they are to
 /// be deleted.
 void Operation::dropAllReferences() {
   for (auto &op : getOpOperands())
     op.drop();

mlir/test/python/ir/operation.py

 def testOperationLoc():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with ctx:
     loc = Location.name("loc")
     op = Operation.create("custom.op", loc=loc)
     assert op.location == loc
     assert op.operation.location == loc
+
+# CHECK-LABEL: TEST: testModuleMerge
+@run
+def testModuleMerge():
+  with Context():
+    m1 = Module.parse("func private @foo()")
+    m2 = Module.parse("""
+      func private @bar()
+      func private @qux()
+    """)
+    foo = m1.body.operations[0]
+    bar = m2.body.operations[0]
+    qux = m2.body.operations[1]
+    bar.move_before(foo)
+    qux.move_after(foo)
+
+    # CHECK: module
+    # CHECK: func private @bar
+    # CHECK: func private @foo
+    # CHECK: func private @qux
+    print(m1)
+
+    # CHECK: module {
+    # CHECK-NEXT: }
+    print(m2)
+
+
+# CHECK-LABEL: TEST: testAppendMoveFromAnotherBlock
+@run
+def testAppendMoveFromAnotherBlock():
+  with Context():
+    m1 = Module.parse("func private @foo()")
+    m2 = Module.parse("func private @bar()")
+    func = m1.body.operations[0]  
+    m2.body.append(func)
+
+    # CHECK: module
+    # CHECK: func private @bar
+    # CHECK: func private @foo
+
+    print(m2)
+    # CHECK: module {
+    # CHECK-NEXT: }
+    print(m1)
+
+
+# CHECK-LABEL: TEST: testDetachFromParent
+@run
+def testDetachFromParent():
+  with Context():
+    m1 = Module.parse("func private @foo()")
+    func = m1.body.operations[0].detach_from_parent()
+
+    try:
+      func.detach_from_parent()
+    except ValueError as e:
+      if "has no parent" not in str(e):
+        raise
+    else:
+      assert False, "expected ValueError when detaching a detached operation"
+
+    print(m1)
+    # CHECK-NOT: func private @foo