[mlir][CFGToSCF] Add interface changes for downstream projects

[zero9178]

Summary

This is a follow-up to https://reviews.llvm.org/D156889

Downstream projects may have more complicated ops than the control flow ops upstream and therefore need a more powerful interface to support the lifting process. Use cases include the propagation of (inherent) metadata that was previously on the control flow ops and now needs to be lifted to structured control flow ops.
Since the lifting process is inherently non-local in respect to the function-body, we require stronger guarantees from the interface.

This patch therefore makes two changes to the interface:

• Passes the terminator that is being replaced to createStructuredBranchRegionTerminatorOp
• Adds as precondition to createCFGSwitchOp that its predecessors are already correctly established

Asserts have been added to verify these were it makes sense and to correctly state intent. I have not added tests purely because testing preconditions like these is not really feasible (and incredibly specific).

[zero9178]

rebase

[definelicht]

LGTM

mlir/lib/Transforms/Utils/CFGToSCF.cpp

514

574

mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h

public:
   /// Creates an `scf.if` op if `controlFlowCondOp` is a `cf.cond_br` op or
   /// an `scf.index_switch` if `controlFlowCondOp` is a `cf.switch`.
   /// Returns failure otherwise.
   FailureOr<Operation *> createStructuredBranchRegionOp(
       OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
       MutableArrayRef<Region> regions) override;
 
   /// Creates an `scf.yield` op returning the given results.
-  LogicalResult
-  createStructuredBranchRegionTerminatorOp(Location loc, OpBuilder &builder,
-                                           Operation *branchRegionOp,
-                                           ValueRange results) override;
+  LogicalResult createStructuredBranchRegionTerminatorOp(
+      Location loc, OpBuilder &builder, Operation *branchRegionOp,
+      Operation *replacedControlFlowOp, ValueRange results) override;
 
   /// Creates an `scf.while` op. The loop body is made the before-region of the
   /// while op and terminated with an `scf.condition` op. The after-region does
   /// nothing but forward the iteration variables.
   FailureOr<Operation *>
   createStructuredDoWhileLoopOp(OpBuilder &builder, Operation *replacedOp,
                                 ValueRange loopVariablesInit, Value condition,
                                 ValueRange loopVariablesNextIter,

mlir/include/mlir/Transforms/CFGToSCF.h

public:
   /// Returns failure if incapable of converting the control flow graph
   /// operation.
   virtual FailureOr<Operation *> createStructuredBranchRegionOp(
       OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
       MutableArrayRef<Region> regions) = 0;
 
   /// Creates a return-like terminator for a branch region of the op returned
   /// by `createStructuredBranchRegionOp`. `branchRegionOp` is the operation
-  /// returned by `createStructuredBranchRegionOp` while `results` are the
-  /// values that should be returned by the branch region.
-  virtual LogicalResult
-  createStructuredBranchRegionTerminatorOp(Location loc, OpBuilder &builder,
-                                           Operation *branchRegionOp,
-                                           ValueRange results) = 0;
+  /// returned by `createStructuredBranchRegionOp`.
+  /// `replacedControlFlowOp` is the control flow op being replaced by the
+  /// terminator or nullptr if the terminator is not replacing any existing
+  /// control flow op. `results` are the values that should be returned by the
+  /// branch region.
+  virtual LogicalResult createStructuredBranchRegionTerminatorOp(
+      Location loc, OpBuilder &builder, Operation *branchRegionOp,
+      Operation *replacedControlFlowOp, ValueRange results) = 0;
 
   /// Creates a structured control flow operation representing a do-while loop.
   /// The do-while loop is expected to have the exact same result types as the
   /// types of the iteration values.
   /// `loopBody` is the body of the loop. The implementation of this
   /// function must create a suitable terminator op at the end of the last block
   /// in `loopBody` which continues the loop if `condition` is 1 and exits the
   /// loop if 0. `loopValuesNextIter` are the values that have to be passed as
public:
   /// suitable as flag for `createCFGSwitchOp`.
   virtual Value getCFGSwitchValue(Location loc, OpBuilder &builder,
                                   unsigned value) = 0;
 
   /// Creates a switch CFG branch operation branching to one of
   /// `caseDestinations` or `defaultDest`. This is used by the transformation
   /// for intermediate transformations before lifting to structured control
   /// flow. The switch op branches based on `flag` which is guaranteed to be of
-  /// the same type as values returned by `getCFGSwitchValue`. Note:
-  /// `caseValues` and other related ranges may be empty to represent an
+  /// the same type as values returned by `getCFGSwitchValue`. The insertion
+  /// block of the builder is guaranteed to have its predecessors already set
+  /// to create an equivalent CFG after this operation.
+  /// Note: `caseValues` and other related ranges may be empty to represent an
   /// unconditional branch.
   virtual void createCFGSwitchOp(Location loc, OpBuilder &builder, Value flag,
                                  ArrayRef<unsigned> caseValues,
                                  BlockRange caseDestinations,
                                  ArrayRef<ValueRange> caseArguments,
                                  Block *defaultDest,
                                  ValueRange defaultArgs) = 0;
 

mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp

ControlFlowToSCFTransformation::createStructuredBranchRegionOp(
   controlFlowCondOp->emitOpError(
       "Cannot convert unknown control flow op to structured control flow");
   return failure();
 }
 
 LogicalResult
 ControlFlowToSCFTransformation::createStructuredBranchRegionTerminatorOp(
     Location loc, OpBuilder &builder, Operation *branchRegionOp,
-    ValueRange results) {
+    Operation *replacedControlFlowOp, ValueRange results) {
   builder.create<scf::YieldOp>(loc, results);
   return success();
 }
 
 FailureOr<Operation *>
 ControlFlowToSCFTransformation::createStructuredDoWhileLoopOp(
     OpBuilder &builder, Operation *replacedOp, ValueRange loopVariablesInit,
     Value condition, ValueRange loopVariablesNextIter, Region &&loopBody) {

mlir/lib/Transforms/Utils/CFGToSCF.cpp

void redirectEdge(Edge edge, ValueRange extraArgs = {}) const {
 
     edge.setSuccessor(multiplexerBlock);
     successorOperands.assign(newSuccOperands);
   }
 
   /// Creates a switch op using `builder` which dispatches to the original
   /// successors of the edges passed to `create` minus the ones in `excluded`.
   /// The builder's insertion point has to be in a block dominated by the
-  /// multiplexer block.
+  /// multiplexer block. All edges to the multiplexer block must have already
+  /// been redirected using `redirectEdge`.
   void createSwitch(
       Location loc, OpBuilder &builder, CFGToSCFInterface &interface,
       const SmallPtrSetImpl<Block *> &excluded = SmallPtrSet<Block *, 1>{}) {
     // We create the switch by creating a case for all entries and then
     // splitting of the last entry as a default case.
 
     SmallVector<ValueRange> caseArguments;
     SmallVector<unsigned> caseValues;
void createSwitch(
     Value realDiscriminator = discriminator;
     if (!realDiscriminator || caseArguments.size() == 1)
       realDiscriminator = getSwitchValue(0);
 
     caseValues.pop_back();
     Block *defaultDest = caseDestinations.pop_back_val();
     ValueRange defaultArgs = caseArguments.pop_back_val();
 
+    assert(!builder.getInsertionBlock()->hasNoPredecessors() &&
+           "Edges need to be redirected prior to creating switch.");
     interface.createCFGSwitchOp(loc, builder, realDiscriminator, caseValues,
                                 caseDestinations, caseArguments, defaultDest,
                                 defaultArgs);
   }
 
 private:
   /// Newly created multiplexer block.
   Block *multiplexerBlock;

 createSingleEntryBlock(Location loc, ArrayRef<Edge> entryEdges,
                        function_ref<Value(unsigned)> getSwitchValue,
                        function_ref<Value(Type)> getUndefValue,
                        CFGToSCFInterface &interface) {
   auto result = EdgeMultiplexer::create(
       loc, llvm::map_to_vector(entryEdges, std::mem_fn(&Edge::getSuccessor)),
       getSwitchValue, getUndefValue);
 
-  auto builder = OpBuilder::atBlockBegin(result.getMultiplexerBlock());
-  result.createSwitch(loc, builder, interface);
-
+  // Redirect the edges prior to creating the switch op.
+  // We guarantee that predecessors are up to date.
   for (Edge edge : entryEdges)
     result.redirectEdge(edge);
 
+  auto builder = OpBuilder::atBlockBegin(result.getMultiplexerBlock());
+  result.createSwitch(loc, builder, interface);
+
   return result;
 }
 
 namespace {
 /// Special loop properties of a structured loop.
 /// A structured loop is a loop satisfying all of the following:
 /// * Has at most one entry, one exit and one back edge.
 /// * The back edge originates from the same block as the exit edge.
static FailureOr<StructuredLoopProperties> createSingleExitingLatch(
 
   // Create a separate exit block that comes right after the latch.
   auto *exitBlock = new Block;
   exitBlock->insertAfter(latchBlock);
 
   // Since this is a loop, all back edges point to the same loop header.
   Block *loopHeader = backEdges.front().getSuccessor();
 
+  // Redirect the edges prior to creating the switch op.
+  // We guarantee that predecessors are up to date.
+
+  // Redirecting back edges with `shouldRepeat` as 1.
+  for (Edge backEdge : backEdges)
+    multiplexer.redirectEdge(backEdge, /*extraArgs=*/getSwitchValue(1));
+
+  // Redirecting exits edges with `shouldRepeat` as 0.
+  for (Edge exitEdge : exitEdges)
+    multiplexer.redirectEdge(exitEdge, /*extraArgs=*/getSwitchValue(0));
+
   // Create the new only back edge to the loop header. Branch to the
   // exit block otherwise.
   Value shouldRepeat = latchBlock->getArguments().back();
   {
     auto builder = OpBuilder::atBlockBegin(latchBlock);
     interface.createConditionalBranch(
         builder.getUnknownLoc(), builder, shouldRepeat, loopHeader,
         latchBlock->getArguments().take_front(loopHeader->getNumArguments()),
if (!exitEdges.empty()) {
       if (failed(terminator))
         return failure();
       // Transform the just created transform operation in the case that an
       // occurrence of it existed in input IR.
       exitCombiner.combineExit(*terminator, getSwitchValue);
     }
   }
 
-  // Redirecting back edges with `shouldRepeat` as 1.
-  for (Edge backEdge : backEdges)
-    multiplexer.redirectEdge(backEdge, /*extraArgs=*/getSwitchValue(1));
-
-  // Redirecting exits edges with `shouldRepeat` as 0.
-  for (Edge exitEdge : exitEdges)
-    multiplexer.redirectEdge(exitEdge, /*extraArgs=*/getSwitchValue(0));
-
   return StructuredLoopProperties{latchBlock, /*condition=*/shouldRepeat,
                                   exitBlock};
 }
 
 /// Transforms a structured loop into a loop in reduce form.
 ///
 /// Reduce form is defined as a structured loop where:
 /// (0) No values defined within the loop body are used outside the loop body.
while (!scc.isAtEnd()) {
     llvm::SmallSetVector<Block *, 4> cycleBlockSet(scc->begin(), scc->end());
     ++scc;
 
     CycleEdges edges = calculateCycleEdges(cycleBlockSet);
     Block *loopHeader = edges.entryEdges.front().getSuccessor();
     // First turn the cycle into a loop by creating a single entry block if
     // needed.
     if (edges.entryEdges.size() > 1) {
+      SmallVector<Edge> edgesToEntryBlocks;
+      llvm::append_range(edgesToEntryBlocks, edges.entryEdges);
+      llvm::append_range(edgesToEntryBlocks, edges.backEdges);
+
       EdgeMultiplexer multiplexer = createSingleEntryBlock(
-          loopHeader->getTerminator()->getLoc(), edges.entryEdges,
+          loopHeader->getTerminator()->getLoc(), edgesToEntryBlocks,
           getSwitchValue, getUndefValue, interface);
 
-      for (Edge edge : edges.backEdges)
-        multiplexer.redirectEdge(edge);
-
       loopHeader = multiplexer.getMultiplexerBlock();
     }
     cycleBlockSet.insert(loopHeader);
 
     // Then turn it into a structured loop by creating a single latch.
     FailureOr<StructuredLoopProperties> loopProperties =
         createSingleExitingLatch(
             edges.backEdges.front().getFromBlock()->getTerminator()->getLoc(),
if (failed(result))
       return failure();
     structuredCondOp = *result;
     regionEntry->getTerminator()->erase();
   }
 
   for (auto &&[block, valueRange] : createdEmptyBlocks) {
     auto builder = OpBuilder::atBlockEnd(block);
     LogicalResult result = interface.createStructuredBranchRegionTerminatorOp(
-        structuredCondOp->getLoc(), builder, structuredCondOp, valueRange);
+        structuredCondOp->getLoc(), builder, structuredCondOp, nullptr,
+        valueRange);
     if (failed(result))
       return failure();
   }
 
   // Any leftover users of the continuation must be from unconditional branches
   // in a branch region. There can only be at most one per branch region as
   // all branch regions have been made single-entry single-exit above.
   // Replace them with the region terminator.
   for (Operation *user : llvm::make_early_inc_range(continuation->getUsers())) {
     assert(user->getNumSuccessors() == 1);
     auto builder = OpBuilder::atBlockTerminator(user->getBlock());
     LogicalResult result = interface.createStructuredBranchRegionTerminatorOp(
-        user->getLoc(), builder, structuredCondOp,
+        user->getLoc(), builder, structuredCondOp, user,
         static_cast<OperandRange>(
             getMutableSuccessorOperands(user->getBlock(), 0)));
     if (failed(result))
       return failure();
     user->erase();
   }
 
   for (auto &&[oldValue, newValue] :