Moving Operations from one Block into another Operation/Block

MLIR
9

Hi all, we (GrAI Matter Labs) have a very similar problem (and implementation) as well. First, we cluster related ops based on MLIR’s slice analysis. Something like:

llvm::SetVector<Operation *> slice;
mlir::getForwardSlice(states.getOperation(),
                      &slice,
                      [&states](Operation *op) -> bool {
                          return somePredicate(op);
                      });
slice.insert(states.getOperation());
llvm::SetVector<Operation *> backwardRoots(slice.begin(), slice.end());
for (auto *root: backwardRoots) {
    llvm::SetVector<Operation *> backwardSlice;
    ::mlir::getBackwardSlice(
            root, &backwardSlice, [&states, &root](Operation *op) -> bool {
                return someOtherPredicate(op);
            });
    slice.insert(backwardSlice.begin(), backwardSlice.end());
}

return topologicalSort(slice);

In our case the ops need not be contiguous in a block. Afterwards these get moved into some new op region:

template <typename RegionOpTy>
RegionOpTy Cluster::makeRegionOp(OpBuilder &builder, const Cluster& cluster, BlockAndValueMapping &map) {
    Region body;
    mlir::IRRewriter bodyBuilder(builder.getContext());
    prepareBuilderForFloatingRegion(bodyBuilder, body);

    for (auto op : cluster.ops) bodyBuilder.clone(*op, map);
    
    auto mappedValues = cluster.mappedValuesUsedOutside(map).takeVector();
    auto returnOp = bodyBuilder.create<SomeTerminator>(
        UnknownLoc::get(builder.getContext()), llvm::makeArrayRef(mappedValues));

    auto regionOp = builder.create<RegionOpTy>(UnknownLoc::get(builder.getContext()),
                                               returnOp.getOperandTypes());
    regionOp.body().takeBody(body);
    return regionOp;
}

and a bunch of updating ops in the original region to use result values of this new op.

If something like that is available builtin we would be very happy to use that as well.