[mlir][cf] Add ControlFlow to SCF lifting pass

Structured control flow ops have proven very useful for many transformations doing analysis on conditional flow and loops. Doing these transformations on CFGs requires repeated analysis of the IR possibly leading to more complicated or less capable implementations. With structured control flow, a lot of the information is already present in the structure.

This patch therefore adds a transformation making it possible to lift arbitrary control flow graphs to structured control flow operations. The algorithm used is outlined in https://dl.acm.org/doi/10.1145/2693261. The complexity in implementing the algorithm was mostly spent correctly handling block arguments in MLIR (the paper only addresses the control flow graph part of it).

Note that the transformation has been implemented fully generically and does not depend on any dialect. An interface implemented by the caller is used to construct any operation necessary for the transformation, making it possible to create an interface implementation purpose fit for ones IR.

For the purpose of testing and due to likely being a very common scenario, this patch adds an interface implementation lifting the control flow dialect to the SCF dialect.
Note the use of the word "lifting". Unlike other conversion passes, this pass is not 100% guaranteed to convert all ControlFlow ops.
Only if the input region being transformed contains a single kind of return-like operations is it guaranteed to replace all control flow ops. If that is not the case, exactly one control flow op will remain branching to regions terminating with a given return-like operation (e.g. one region terminates with `llvm.return` the other with `llvm.unreachable`).

Differential Revision: https://reviews.llvm.org/D156889

Author: zero9178

llvm/include/llvm/ADT/STLExtras.h

@@ -1792,7 +1792,7 @@ OutputIt copy_if(R &&Range, OutputIt Out, UnaryPredicate P) {
 template <typename T, typename R, typename Predicate>
 T *find_singleton(R &&Range, Predicate P, bool AllowRepeats = false) {
   T *RC = nullptr;
-  for (auto *A : Range) {
+  for (auto &&A : Range) {
     if (T *PRC = P(A, AllowRepeats)) {
       if (RC) {
         if (!AllowRepeats || PRC != RC)

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

@@ -0,0 +1,26 @@
+//===- ControlFlowToSCF.h - ControlFlow to SCF -------------*- C++ ------*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Define conversions from the ControlFlow dialect to the SCF dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_CONTROLFLOWTOSCF_CONTROLFLOWTOSCF_H
+#define MLIR_CONVERSION_CONTROLFLOWTOSCF_CONTROLFLOWTOSCF_H
+
+#include <memory>
+
+namespace mlir {
+class Pass;
+
+#define GEN_PASS_DECL_LIFTCONTROLFLOWTOSCFPASS
+#include "mlir/Conversion/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_CONTROLFLOWTOSCF_CONTROLFLOWTOSCF_H

mlir/include/mlir/Conversion/Passes.h

@@ -22,6 +22,7 @@
 #include "mlir/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.h"
 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
+#include "mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h"
 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.h"
 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRVPass.h"
 #include "mlir/Conversion/ConvertToLLVM/ToLLVMPass.h"

mlir/include/mlir/Conversion/Passes.td

@@ -283,6 +283,31 @@ def ConvertControlFlowToLLVMPass : Pass<"convert-cf-to-llvm", "ModuleOp"> {
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// ControlFlowToSCF
+//===----------------------------------------------------------------------===//
+
+def LiftControlFlowToSCFPass : Pass<"lift-cf-to-scf"> {
+  let summary = "Lift ControlFlow dialect to SCF dialect";
+  let description = [{
+    Lifts ControlFlow operations to SCF dialect operations.
+
+    This pass is prefixed with "lift" instead of "convert" as it is not always
+    guaranteed to replace all ControlFlow ops.
+    If a region contains only a single kind of return-like operation, all
+    ControlFlow operations will be replaced successfully.
+    Otherwise a single ControlFlow switch branching to one block per return-like
+    operation kind remains.
+  }];
+
+  let dependentDialects = ["scf::SCFDialect",
+                           "arith::ArithDialect",
+                           "ub::UBDialect",
+                           // TODO: This is only necessary until we have a
+                           //       ub.unreachable op.
+                           "func::FuncDialect"];
+}
+
 //===----------------------------------------------------------------------===//
 // ControlFlowToSPIRV
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/Block.h

@@ -59,6 +59,10 @@ class Block : public IRObjectWithUseList<BlockOperand>,
   /// the specified block.
   void insertBefore(Block *block);
 
+  /// Insert this block (which must not already be in a region) right after
+  /// the specified block.
+  void insertAfter(Block *block);
+
   /// Unlink this block from its current region and insert it right before the
   /// specific block.
   void moveBefore(Block *block);

mlir/include/mlir/IR/ValueRange.h

@@ -167,6 +167,14 @@ class MutableOperandRange {
     return operator OperandRange()[index];
   }
 
+  OperandRange::iterator begin() const {
+    return static_cast<OperandRange>(*this).begin();
+  }
+
+  OperandRange::iterator end() const {
+    return static_cast<OperandRange>(*this).end();
+  }
+
 private:
   /// Update the length of this range to the one provided.
   void updateLength(unsigned newLength);

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

@@ -82,6 +82,11 @@ class SuccessorOperands {
   /// Get the range of operands that are simply forwarded to the successor.
   OperandRange getForwardedOperands() const { return forwardedOperands; }
 
+  /// Get the range of operands that are simply forwarded to the successor.
+  MutableOperandRange getMutableForwardedOperands() const {
+    return forwardedOperands;
+  }
+
   /// Get a slice of the operands forwarded to the successor. The given range
   /// must not contain any operands produced by the operation.
   MutableOperandRange slice(unsigned subStart, unsigned subLen) const {

mlir/include/mlir/Transforms/CFGToSCF.h

@@ -0,0 +1,154 @@
+//===- CFGToSCF.h - Control Flow Graph to Structured Control Flow *- C++ -*===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines a generic `transformCFGToSCF` function that can be
+// used to lift any dialect operations implementing control flow graph
+// operations to any dialect implementing structured control flow operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_CFGTOSCF_H
+#define MLIR_TRANSFORMS_CFGTOSCF_H
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/Operation.h"
+
+namespace mlir {
+
+/// Interface that should be implemented by any caller of `transformCFGToSCF`.
+/// The transformation requires the caller to 1) create switch-like control
+/// flow operations for intermediate transformations and 2) to create
+/// the desired structured control flow ops.
+class CFGToSCFInterface {
+public:
+  virtual ~CFGToSCFInterface() = default;
+
+  /// Creates a structured control flow operation branching to one of `regions`.
+  /// It replaces `controlFlowCondOp` and must have `resultTypes` as results.
+  /// `regions` contains the list of branch regions corresponding to each
+  /// successor of `controlFlowCondOp`. Their bodies must simply be taken and
+  /// left as is.
+  /// 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;
+
+  /// 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
+  /// the iteration values for the next iteration if continuing, or the result
+  /// of the loop if exiting.
+  /// `condition` is guaranteed to be of the same type as values returned by
+  /// `getCFGSwitchValue` with either 0 or 1 as value.
+  ///
+  /// `loopValuesInit` are the values used to initialize the iteration
+  /// values of the loop.
+  /// Returns failure if incapable of creating a loop op.
+  virtual FailureOr<Operation *> createStructuredDoWhileLoopOp(
+      OpBuilder &builder, Operation *replacedOp, ValueRange loopValuesInit,
+      Value condition, ValueRange loopValuesNextIter, Region &&loopBody) = 0;
+
+  /// Creates a constant operation with a result representing `value` that is
+  /// 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
+  /// unconditional branch.
+  virtual void createCFGSwitchOp(Location loc, OpBuilder &builder, Value flag,
+                                 ArrayRef<unsigned> caseValues,
+                                 BlockRange caseDestinations,
+                                 ArrayRef<ValueRange> caseArguments,
+                                 Block *defaultDest,
+                                 ValueRange defaultArgs) = 0;
+
+  /// Creates a constant operation returning an undefined instance of `type`.
+  /// This is required by the transformation as the lifting process might create
+  /// control-flow paths where an SSA-value is undefined.
+  virtual Value getUndefValue(Location loc, OpBuilder &builder, Type type) = 0;
+
+  /// Creates a return-like terminator indicating unreachable.
+  /// This is required when the transformation encounters a statically known
+  /// infinite loop. Since structured control flow ops are not terminators,
+  /// after lifting an infinite loop, a terminator has to be placed after to
+  /// possibly satisfy the terminator requirement of the region originally
+  /// passed to `transformCFGToSCF`.
+  ///
+  /// `region` is guaranteed to be the region originally passed to
+  /// `transformCFGToSCF` and the op is guaranteed to always be an op in a block
+  /// directly nested under `region` after the transformation.
+  ///
+  /// Returns failure if incapable of creating an unreachable terminator.
+  virtual FailureOr<Operation *>
+  createUnreachableTerminator(Location loc, OpBuilder &builder,
+                              Region &region) = 0;
+
+  /// Helper function to create an unconditional branch using
+  /// `createCFGSwitchOp`.
+  void createSingleDestinationBranch(Location loc, OpBuilder &builder,
+                                     Value dummyFlag, Block *destination,
+                                     ValueRange arguments) {
+    createCFGSwitchOp(loc, builder, dummyFlag, {}, {}, {}, destination,
+                      arguments);
+  }
+
+  /// Helper function to create a conditional branch using
+  /// `createCFGSwitchOp`.
+  void createConditionalBranch(Location loc, OpBuilder &builder,
+                               Value condition, Block *trueDest,
+                               ValueRange trueArgs, Block *falseDest,
+                               ValueRange falseArgs) {
+    createCFGSwitchOp(loc, builder, condition, {0}, {falseDest}, {falseArgs},
+                      trueDest, trueArgs);
+  }
+};
+
+/// Transformation lifting any dialect implementing control flow graph
+/// operations to a dialect implementing structured control flow operations.
+/// `region` is the region that should be transformed.
+/// The implementation of `interface` is responsible for the conversion of the
+/// control flow operations to the structured control flow operations.
+///
+/// If the region contains only a single kind of return-like operation, all
+/// control flow graph operations will be converted successfully.
+/// Otherwise a single control flow graph operation branching to one block
+/// per return-like operation kind remains.
+///
+/// The transformation currently requires that all control flow graph operations
+/// have no side effects, implement the BranchOpInterface and does not have any
+/// operation produced successor operands.
+/// Returns failure if any of the preconditions are violated or if any of the
+/// methods of `interface` failed. The IR is left in an unspecified state.
+///
+/// Otherwise, returns true or false if any changes to the IR have been made.
+FailureOr<bool> transformCFGToSCF(Region &region, CFGToSCFInterface &interface,
+                                  DominanceInfo &dominanceInfo);
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_CFGTOSCF_H

mlir/lib/Conversion/CMakeLists.txt

@@ -12,6 +12,7 @@ add_subdirectory(ComplexToLLVM)
 add_subdirectory(ComplexToSPIRV)
 add_subdirectory(ComplexToStandard)
 add_subdirectory(ControlFlowToLLVM)
+add_subdirectory(ControlFlowToSCF)
 add_subdirectory(ControlFlowToSPIRV)
 add_subdirectory(ConvertToLLVM)
 add_subdirectory(FuncToLLVM)

mlir/lib/Conversion/ControlFlowToSCF/CMakeLists.txt

@@ -0,0 +1,23 @@
+add_mlir_conversion_library(MLIRControlFlowToSCF
+  ControlFlowToSCF.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ControlFlowToSCF
+
+  DEPENDS
+  MLIRConversionPassIncGen
+  intrinsics_gen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRAnalysis
+  MLIRArithDialect
+  MLIRControlFlowDialect
+  MLIRFuncDialect
+  MLIRSCFDialect
+  MLIRUBDialect
+  MLIRPass
+  MLIRTransformUtils
+)