[MLIR] Add block detach func to CAPI and use it in Python bindings

[jdd]

Summary

Adds mlirBlockDetach to the CAPI to remove a block from its parent
region. Use it in the Python bindings to implement
Block.append_to(region).

mlir/include/mlir-c/IR.h

 /// ownership to the caller.
 MLIR_CAPI_EXPORTED MlirBlock mlirBlockCreate(intptr_t nArgs,
                                              MlirType const *args,
                                              MlirLocation const *locs);
 
 /// Takes a block owned by the caller and destroys it.
 MLIR_CAPI_EXPORTED void mlirBlockDestroy(MlirBlock block);
 
+/// Detach a block from the owning region and assume ownership.
+MLIR_CAPI_EXPORTED void mlirBlockDetach(MlirBlock block);
+
 /// Checks whether a block is null.
 static inline bool mlirBlockIsNull(MlirBlock block) { return !block.ptr; }
 
 /// Checks whether two blocks handles point to the same block. This does not
 /// perform deep comparison.
 MLIR_CAPI_EXPORTED bool mlirBlockEqual(MlirBlock block, MlirBlock other);
 
 /// Returns the closest surrounding operation that contains this block.

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyBlock>(m, "Block", py::module_local())
                                               argLocs.data());
             mlirRegionInsertOwnedBlock(parent, 0, block);
             return PyBlock(parent.getParentOperation(), block);
           },
           py::arg("parent"), py::arg("arg_types") = py::list(),
           "Creates and returns a new Block at the beginning of the given "
           "region (with given argument types).")
       .def(
+          "append_to",
+          [](PyBlock &self, PyRegion &region) {
+            MlirBlock b = self.get();
+            if (!mlirRegionIsNull(mlirBlockGetParentRegion(b)))
+              mlirBlockDetach(b);
+            mlirRegionAppendOwnedBlock(region.get(), b);
+          },
+          "Append this block to a region, transferring ownership if necessary")
+      .def(
           "create_before",
           [](PyBlock &self, py::args pyArgTypes) {
             self.checkValid();
             llvm::SmallVector<MlirType, 4> argTypes;
             llvm::SmallVector<MlirLocation, 4> argLocs;
             argTypes.reserve(pyArgTypes.size());
             argLocs.reserve(pyArgTypes.size());
             for (auto &pyArg : pyArgTypes) {

mlir/lib/CAPI/IR/IR.cpp

void mlirBlockInsertOwnedOperationBefore(MlirBlock block,
   assert(unwrap(reference)->getBlock() == unwrap(block) &&
          "expected reference operation to belong to the block");
   unwrap(block)->getOperations().insert(Block::iterator(unwrap(reference)),
                                         unwrap(operation));
 }
 
 void mlirBlockDestroy(MlirBlock block) { delete unwrap(block); }
 
+void mlirBlockDetach(MlirBlock block) {
+  Block *b = unwrap(block);
+  b->getParent()->getBlocks().remove(b);
+}
+
 intptr_t mlirBlockGetNumArguments(MlirBlock block) {
   return static_cast<intptr_t>(unwrap(block)->getNumArguments());
 }
 
 MlirValue mlirBlockAddArgument(MlirBlock block, MlirType type,
                                MlirLocation loc) {
   return wrap(unwrap(block)->addArgument(unwrap(type), unwrap(loc)));
 }

mlir/test/CAPI/ir.c

static void buildWithInsertionsAndPrint(MlirContext ctx) {
   MlirRegion region = mlirOperationGetRegion(op, 0);
 
   // Use integer types of different bitwidth as block arguments in order to
   // differentiate blocks.
   MlirType i1 = mlirIntegerTypeGet(ctx, 1);
   MlirType i2 = mlirIntegerTypeGet(ctx, 2);
   MlirType i3 = mlirIntegerTypeGet(ctx, 3);
   MlirType i4 = mlirIntegerTypeGet(ctx, 4);
+  MlirType i5 = mlirIntegerTypeGet(ctx, 5);
   MlirBlock block1 = mlirBlockCreate(1, &i1, &loc);
   MlirBlock block2 = mlirBlockCreate(1, &i2, &loc);
   MlirBlock block3 = mlirBlockCreate(1, &i3, &loc);
   MlirBlock block4 = mlirBlockCreate(1, &i4, &loc);
+  MlirBlock block5 = mlirBlockCreate(1, &i5, &loc);
   // Insert blocks so as to obtain the 1-2-3-4 order,
   mlirRegionInsertOwnedBlockBefore(region, nullBlock, block3);
   mlirRegionInsertOwnedBlockBefore(region, block3, block2);
   mlirRegionInsertOwnedBlockAfter(region, nullBlock, block1);
   mlirRegionInsertOwnedBlockAfter(region, block3, block4);
+  mlirRegionInsertOwnedBlockBefore(region, block3, block5);
 
   MlirOperationState op1State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op1"), loc);
   MlirOperationState op2State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op2"), loc);
   MlirOperationState op3State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op3"), loc);
   MlirOperationState op4State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op4"), loc);
   MlirOperationState op5State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op5"), loc);
   MlirOperationState op6State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op6"), loc);
   MlirOperationState op7State =
       mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op7"), loc);
+  MlirOperationState op8State =
+      mlirOperationStateGet(mlirStringRefCreateFromCString("dummy.op8"), loc);
   MlirOperation op1 = mlirOperationCreate(&op1State);
   MlirOperation op2 = mlirOperationCreate(&op2State);
   MlirOperation op3 = mlirOperationCreate(&op3State);
   MlirOperation op4 = mlirOperationCreate(&op4State);
   MlirOperation op5 = mlirOperationCreate(&op5State);
   MlirOperation op6 = mlirOperationCreate(&op6State);
   MlirOperation op7 = mlirOperationCreate(&op7State);
+  MlirOperation op8 = mlirOperationCreate(&op8State);
 
   // Insert operations in the first block so as to obtain the 1-2-3-4 order.
   MlirOperation nullOperation = mlirBlockGetFirstOperation(block1);
   assert(mlirOperationIsNull(nullOperation));
   mlirBlockInsertOwnedOperationBefore(block1, nullOperation, op3);
   mlirBlockInsertOwnedOperationBefore(block1, op3, op2);
   mlirBlockInsertOwnedOperationAfter(block1, nullOperation, op1);
   mlirBlockInsertOwnedOperationAfter(block1, op3, op4);
 
   // Append operations to the rest of blocks to make them non-empty and thus
   // printable.
   mlirBlockAppendOwnedOperation(block2, op5);
   mlirBlockAppendOwnedOperation(block3, op6);
   mlirBlockAppendOwnedOperation(block4, op7);
+  mlirBlockAppendOwnedOperation(block5, op8);
+
+  // Remove block5.
+  mlirBlockDetach(block5);
+  mlirBlockDestroy(block5);
 
   mlirOperationDump(op);
   mlirOperationDestroy(op);
   mlirContextSetAllowUnregisteredDialects(ctx, false);
   // clang-format off
   // CHECK-LABEL:  "insertion.order.test"
   // CHECK:      ^{{.*}}(%{{.*}}: i1
   // CHECK:        "dummy.op1"
   // CHECK-NEXT:   "dummy.op2"
   // CHECK-NEXT:   "dummy.op3"
   // CHECK-NEXT:   "dummy.op4"
   // CHECK:      ^{{.*}}(%{{.*}}: i2
   // CHECK:        "dummy.op5"
+  // CHECK-NOT:  ^{{.*}}(%{{.*}}: i5
+  // CHECK-NOT:    "dummy.op8"
   // CHECK:      ^{{.*}}(%{{.*}}: i3
   // CHECK:        "dummy.op6"
   // CHECK:      ^{{.*}}(%{{.*}}: i4
   // CHECK:        "dummy.op7"
   // clang-format on
 }
 
 /// Creates operations with type inference and tests various failure modes.

mlir/test/python/ir/blocks.py

with InsertionPoint(module.body):
       f = func.FuncOp("test", ([f32], []))
       entry_block = Block.create_at_start(f.operation.regions[0], [f32])
       with InsertionPoint(entry_block):
         func.ReturnOp([])
 
     print(module)
     assert module.operation.verify()
     assert f.body.blocks[0] == entry_block
+
+
+# CHECK-LABEL: TEST: testBlockMove
+# CHECK:  %0 = "realop"() ({
+# CHECK:  ^bb0([[ARG0:%.+]]: f32):
+# CHECK:    "ret"([[ARG0]]) : (f32) -> ()
+# CHECK:  }) : () -> f32
+@run
+def testBlockMove():
+  with Context() as ctx, Location.unknown():
+    ctx.allow_unregistered_dialects = True
+    module = Module.create()
+    f32 = F32Type.get()
+    with InsertionPoint(module.body):
+      dummy = Operation.create("dummy", regions=1)
+      block = Block.create_at_start(dummy.operation.regions[0], [f32])
+      with InsertionPoint(block):
+        ret_op = Operation.create("ret", operands=[block.arguments[0]])
+      realop = Operation.create("realop",
+                                results=[r.type for r in ret_op.operands],
+                                regions=1)
+      block.append_to(realop.operation.regions[0])
+      dummy.operation.erase()
+    print(module)