[mlir] Add a ControlFlowSink pass.

[Mogball]

Summary

Control-Flow Sink moves operations whose only uses are in conditionally-executed regions into those regions so that paths in which their results are not needed do not perform unnecessary computation.

Depends on D115087

[mehdi_amini]

mlir/include/mlir/Transforms/Passes.td

315

Isn't this the opposite of LICM?

[Mogball]

mlir/include/mlir/Transforms/Passes.td

315

Basically, yeah.

[jpienaar]

mlir/include/mlir/Transforms/Passes.td

315

At least wrt loops (this is peer of https://llvm.org/docs/Passes.html#sink-code-sinking)

[Mogball]

Fixing the build and updating the pass description.

[Mogball]

Corrected the diff.

[mehdi_amini]

I'm not sure about the interest of such a pass, without some profitability aspects.

I guess this could be a test pass, where the core of the transformation is exposed as a utility. The utility could have some filtering / callback to support the decision process and users could then implement the pass for their own needs.

[Mogball]

In D115088#3174610, @mehdi_amini wrote:

I'm not sure about the interest of such a pass, without some profitability aspects.

I guess this could be a test pass, where the core of the transformation is exposed as a utility. The utility could have some filtering / callback to support the decision process and users could then implement the pass for their own needs.

I'm not against this approach, as I have a specific use-case in mind (TF) that will probably require the API to be modified, but I have to yet decide what exactly the callbacks will be. However, the pass as-is should be a uniform gain, because it only considers non-memory-effecting ops.

[mehdi_amini]

I guess the "single conditionally-executed region" may be enough of a criteria here to be conservative and get started. I suspect we'll still need to evolve this to be more configurable, so I'd rather have the split between the utility and the pass up-front (even if the utility doesn't expose options immediately)

[Mogball]

Simplify the implementation.

[rriddle]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

94–95

Is a magic number necessary? Why not use Optional to represent unknown? Seems safer and less prone to error.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

174

Can we not use something like $_op here?

176

nit: Spell out auto here.

[rriddle]

mlir/lib/Transforms/ControlFlowSink.cpp

39–42

Can you separate the pass and utilty? similarly to how it's done in https://reviews.llvm.org/D114945

45–47

What does the failure represent here? Also, why are you stipulating only r-reference. This is a local utility class, you don't need to add that.

108–112

Seems like this is essentially just checking if all users are nested within region? Seems like that could be checked without the need of DominanceInfo.

118

119

123–124

141

If you want a stack, prefer using SmallVector instead.

142–144

148

You can merge the depth first comment into the comment before the loop.

154

Why the FailureOr? This always succeeds.

156–157

159

This should already be guaranteed by the infra, this situation can't happen in valid IR.

189

Do you need a LogicalResult return? This always succeeds.

[Mogball]

mlir/lib/Transforms/ControlFlowSink.cpp

45–47

Adding && made more sense when I had data structures stored in the struct (which I've since removed by simplifying the logic) that made the utility struct "single-use". I can drop it now.

108–112

What happens if the user is in a block unreachable from the entry block? Then moving the op to the entry block would result in invalid IR.

141

I used a deque after seeing this

154

I thought maybe I'd add LogicalResult in case the pass might fail in the future. But now that I think about it, I don't think it will. I'll drop them.

159

Ack

[bondhugula]

mlir/include/mlir/Transforms/ControlFlowSink.h

13

It is out of line to have a method with this name exposed on mlir:: -- the name isn't descriptive enough for this namespace as well even if this needs to be exposed.

mlir/lib/Transforms/ControlFlowSink.cpp

199–209

All code being called from here can just go into a utility in lib/Transforms/Utils/Utils.cpp?

LogicalResult mlir::controlFlowSink(RegionBranchOpInterface branchOp, DominanceInfo &domInfo) {
   ...
}

You can return the number of things that have been sunk if desired instead of LogicalResult. I'm also not sure that a failure of this utility should be the cause for a pass failure. Is the IR going to be invalid? If so, sinkRegions and/or sinkConditionalCode should be documenting that which they don't appear to.

[Mogball]

Review comments and separating utility from pass

[Mogball]

@rriddle @bondhugula pls

[bondhugula]

In D115088#3247232, @Mogball wrote:

@rriddle @bondhugula pls

I still see several comments that haven't been marked "Done".

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

91

"defined" or "known"?

109

A comment here that a None is returned for an unknown upper bound?

mlir/include/mlir/Transforms/ControlFlowSink.h

13

@Mogball Did you missing updating the revision?

[Mogball]

thanks!

mlir/include/mlir/Transforms/ControlFlowSink.h

13

Ah yeah I did.

[Mogball]

Update the diff

[bondhugula]

mlir/lib/Transforms/ControlFlowSink.cpp

16

This file is missing its primary include of Passes.h. In fact, mlir/Transforms/PassDetail.h is layered completely wrong -- it shouldn't be including Passes.h.

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

110–120

Why do you need a deque here? You are only pushing and popping at the end.

[bondhugula]

mlir/lib/Transforms/ControlFlowSink.cpp

16

Looks like mlir::FusionMode is needed for PassDetail.h and Passes.h was included there. That can be cleaned up as part of a separate refactoring.

[Mogball]

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

110–120

I saw this:

https://github.com/llvm/llvm-project/blob/81f9dc8eee3b9319d911583a01b0060b5bce6018/mlir/lib/Transforms/CSE.cpp#L211

[mehdi_amini]

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

110–120

This is referring to a micro-benchmark posted on the mailing list in 2012, do you have more recent data?

Right now when I run the benchmark on my Linux box I get:

Stack:
accum = 636
time = 4.486544 seconds


Deque:
accum = 636
time = 5.598040 seconds


Deque iter insert:
accum = 636
time = 10.433442 seconds


Vector:
accum = 636
time = 2.632639 seconds


List:
accum = 636
time = 31.212231 seconds

(clang++ -O3)

At the time they reported:

List: 92.473490 seconds
Stack: 15.575143 seconds
Deque: 9.129154 seconds
Deque iter insert: 7.384741 seconds
Vector: 13.387473 seconds

[rriddle]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

171

Might be good to add a helper static InvocationBounds InvocationBounds::getUnknown() method that makes it easier to do this.

mlir/include/mlir/Transforms/Utils.h

21

Can you forward declare the classes you need instead of adding an include?

mlir/lib/Transforms/ControlFlowSink.cpp

60

Is the mlir:: here necessary?

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

124–128

mlir/test/Transforms/control-flow-sink.mlir

7

Can you use proper file check variables for all of these SSA captures? Ideally the tests should only be testing the minimum necessary.

https://mlir.llvm.org/getting_started/TestingGuide/

(Same for the other tests)

[Mogball]

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

110–120

I do not. It may be worth changing the one in CSE to a vector too then!

[Mogball]

review comments

[Mogball]

actually update the diff

[jpienaar]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

161

Does this need an update?

mlir/include/mlir/Transforms/Passes.td

319

similar but opposite? :-)

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

94

The mix of Operation& and Operation* I find distracting, I'd just use Operation* throughout as that's the convention elsewhere

148

The name of the pass is more general than what it does, while the comments all refer to this being simple. Is the intention to generalize the pass later (and so the name) and these are just documenting the current state? Or is the intention to have a simple sinking pass and there be others?

150

Could we use llvm::zip?

[Mogball]

review comments

[Mogball]

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

148

It's meant to be generalized later such that the default cost model is the "simple" case (i.e. no duplicating ops). I'll make that clear in the documentation.

[rriddle]

Did a scan and looks good from my point of view as something that we can land and iterate on.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

171

Can we use it here ;)

[bondhugula]

LGTM.

Nit: "work queue" -> "work list" at many places.

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

44

Why is this returning an rvalue ref for a POD? Haven't seen this pattern in the tree.

[jpienaar]

Nice, thanks for refactoring!

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

90

is always known and at least 0 ?

Or is this trying to say it may be unknown but in that case 0 could be returned so there'll always be a lower bound even if not exact? (I think the latter)

mlir/include/mlir/Transforms/Passes.td

313

I'd split this into two and remove simple here now. First say what the pass is doing and then say it's restrictions (which currently is a very conservative cost model)

mlir/lib/Transforms/ControlFlowSink.cpp

26

Some of these now feel like properties of the helper methods rather than this pass.

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

10

utilities

97

And this is reason that the dominance isn't updated? But the dominance information cannot be marked as retained as this is only true locally (e.g., within cf op) [just checking my understanding what preserved means here]

110–120

Create a GitHub tracking issue for it :-)

141

(void) to make it clear whether matches or not doesn't matter here

142

Nit: new line above (i was reading this as one clump and wondering why bounds were being called inside the loop :-))

mlir/test/Transforms/control-flow-sink.mlir

2

Is split-input-file adding anything here? We mostly only use it for error case checking, it returns a less useful file location upon failure.

[Mogball]

rebase

[Mogball]

Ah wait I totally forgot about the other comments -- guess I'm sending an NFC .-.

[Mogball]

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

90

The lower bound is known to be at least zero (because you can't negatively execute a region...)

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

97

Just locally in the subgraph. I'll update the comment.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 LogicalResult verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
                                             Optional<OperandRange> operands);
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
-// A constant value to represent unknown number of region invocations.
+/// A constant value to represent unknown number of region invocations.
 extern const int64_t kUnknownNumRegionInvocations;
 
 namespace detail {
 /// Verify that types match along control flow edges described the given op.
 LogicalResult verifyTypesAlongControlFlowEdges(Operation *op);
 } //  namespace detail
 
 /// This class represents a successor of a region. A region successor can either
public:
   /// the current region.
   ValueRange getSuccessorInputs() const { return inputs; }
 
 private:
   Region *region{nullptr};
   ValueRange inputs;
 };
 
+/// This class represents upper and lower bounds on the number of times a region
+/// of a `RegionBranchOpInterface` can be invoked. The lower bound is at least
+/// zero, but the upper bound may not be defined.
+class InvocationBounds {
+public:
+  /// Create invocation bounds. The lower bound must be at least 0 and only the
+  /// upper bound can be unknown.
+  InvocationBounds(int64_t lb, Optional<int64_t> ub) : lower(lb), upper(ub) {
+    assert(lb >= 0 && "lower bound must be at least 0");
+  }
+
+  /// Return the lower bound.
+  int64_t getLowerBound() const { return lower; }
+
+  /// Return the upper bound.
+  Optional<int64_t> getUpperBound() const { return upper; }
+
+private:
+  /// The minimum number of times the successor region will be invoked.
+  int64_t lower;
+  /// The maximum number of times the successor region will be invoked.
+  Optional<int64_t> upper;
+};
+
 /// Return `true` if `a` and `b` are in mutually exclusive regions as per
 /// RegionBranchOpInterface.
 bool insideMutuallyExclusiveRegions(Operation *a, Operation *b);
 
 //===----------------------------------------------------------------------===//
 // RegionBranchTerminatorOpInterface
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let description = [{
     branching behavior between held regions, i.e. this interface allows for
     expressing control flow information for region holding operations.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Returns the operands of this operation used as the entry arguments when
-        entering the region at `index`, which was specified as a successor of this
-        operation by `getSuccessorRegions`. These operands should correspond 1-1
-        with the successor inputs specified in `getSuccessorRegions`.
+        entering the region at `index`, which was specified as a successor of
+        this operation by `getSuccessorRegions`. These operands should
+        correspond 1-1 with the successor inputs specified in
+        `getSuccessorRegions`.
       }],
       "::mlir::OperandRange", "getSuccessorEntryOperands",
       (ins "unsigned":$index), [{}], /*defaultImplementation=*/[{
         auto operandEnd = this->getOperation()->operand_end();
         return ::mlir::OperandRange(operandEnd, operandEnd);
       }]
     >,
     InterfaceMethod<[{
         Returns the viable successors of a region at `index`, or the possible
         successors when branching from the parent op if `index` is None. These
         are the regions that may be selected during the flow of control. If
         `index` is None, `operands` is a set of optional attributes that
         either correspond to a constant value for each operand of this
         operation, or null if that operand is not a constant. If `index` is
         valid, `operands` corresponds to the entry values of the region at
         `index`. Only a region, i.e. a valid `index`, may use the parent
         operation as a successor. This method allows for describing which
         regions may be executed when entering an operation, and which regions
         are executed after having executed another region of the parent op. The
         successor region must be non-empty.
       }],
       "void", "getSuccessorRegions",
-      (ins "::mlir::Optional<unsigned>":$index, "::mlir::ArrayRef<::mlir::Attribute>":$operands,
+      (ins "::mlir::Optional<unsigned>":$index,
+           "::mlir::ArrayRef<::mlir::Attribute>":$operands,
            "::mlir::SmallVectorImpl<::mlir::RegionSuccessor> &":$regions)
     >,
     InterfaceMethod<[{
         Populates countPerRegion with the number of times this operation will
         invoke the attached regions (assuming the regions yield normally, i.e.
         do not abort or invoke an infinite loop). If the number of region
         invocations is not known statically it will set the number of
         invocations to `kUnknownNumRegionInvocations`.
 
         `operands` is a set of optional attributes that either correspond to a
         constant values for each operand of this operation, or null if that
         operand is not a constant.
       }],
       "void", "getNumRegionInvocations",
       (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands,
            "::mlir::SmallVectorImpl<int64_t> &":$countPerRegion), [{}],
       /*defaultImplementation=*/[{
         unsigned numRegions = this->getOperation()->getNumRegions();
         assert(countPerRegion.empty());
         countPerRegion.resize(numRegions, kUnknownNumRegionInvocations);
       }]
-    >
+    >,
+    InterfaceMethod<[{
+        Populates `invocationBounds` with the minimum and maximum number of
+        times this operation will invoke the attached regions (assuming the
+        regions yield normally, i.e. do not abort or invoke an infinite loop).
+        The minimum number of invocations is at least 0. If the maximum number
+        of invocations cannot be statically determined, then it will be set to
+        `InvocationBounds::kUnknownUpperBound`.
+
+        `operands` is a set of optional attributes that either correspond to a
+        constant values for each operand of this operation, or null if that
+        operand is not a constant.
+      }],
+      "void", "getRegionInvocationBounds",
+      (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands,
+           "::llvm::SmallVectorImpl<::mlir::InvocationBounds> &"
+             :$invocationBounds),
+           [{}],
+      /*defaultImplementation=*/[{
+        ::llvm::SmallVector<int64_t> countPerRegion;
+        ::llvm::cast<ConcreteOp>(this->getOperation()).getNumRegionInvocations(
+            operands, countPerRegion);
+        for (int64_t count : countPerRegion) {
+          if (count == ::mlir::kUnknownNumRegionInvocations)
+            invocationBounds.emplace_back(0, llvm::None);
+          else
+            invocationBounds.emplace_back(count, count);
+        }
+      }]
+    >,
   ];
 
   let verify = [{
     static_assert(!ConcreteOp::template hasTrait<OpTrait::ZeroRegion>(),
                   "expected operation to have non-zero regions");
     return success();
   }];
 

mlir/include/mlir/Transforms/ControlFlowSink.h

+#ifndef MLIR_TRANSFORMS_CONTROLFLOWSINK_H_
+#define MLIR_TRANSFORMS_CONTROLFLOWSINK_H_
+
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+
+namespace mlir {
+class DominanceInfo;
+
+/// Return a list of regions belonging to a `RegionBranchOpInterface` that are
+/// candidates for control-flow sink. For basic control-flow sink, the region
+/// must be known to be invoked at most once.
+LogicalResult getCandidateRegions(RegionBranchOpInterface branch,
+                                  SmallVectorImpl<Region *> &regionsToSink);
+
+/// Given a list of regions that belong to a `RegionBranchOpInterface`, try to
+/// sink in operations from outside the region whose only uses are in that
+/// region.
+LogicalResult sinkConditionalCode(ArrayRef<Region *> regionsToSink,
+                                  DominanceInfo &domInfo);
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_CONTROLFLOWSINK_H_

mlir/include/mlir/Transforms/Passes.h

 /// them with `RewritePatternSet::addWithLabel`. Debug names may be empty, but
 /// patterns created with `RewritePattern::create` have their default debug name
 /// set to their type name.
 std::unique_ptr<Pass>
 createCanonicalizerPass(const GreedyRewriteConfig &config,
                         ArrayRef<std::string> disabledPatterns = llvm::None,
                         ArrayRef<std::string> enabledPatterns = llvm::None);
 
+/// Creates a pass to perform control-flow sinking.
+std::unique_ptr<Pass> createControlFlowSinkPass();
+
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
 /// Creates a loop fusion pass which fuses loops according to type of fusion
 /// specified in `fusionMode`. Buffers of size less than or equal to
 /// `localBufSizeThreshold` are promoted to memory space `fastMemorySpace`.
 std::unique_ptr<OperationPass<FuncOp>>
 createLoopFusionPass(unsigned fastMemorySpace = 0,

mlir/include/mlir/Transforms/Passes.td

Option<"enableRegionSimplification", "region-simplify", "bool",
            /*default=*/"true",
            "Seed the worklist in general top-down order">,
     Option<"maxIterations", "max-iterations", "int64_t",
            /*default=*/"10",
            "Seed the worklist in general top-down order">
   ] # RewritePassUtils.options;
 }
 
+def ControlFlowSink : Pass<"control-flow-sink"> {
+  let summary = "Sink operations into conditional blocks";
+  let description = [{
+    This pass implements a simple control-flow sink on operations that implement
+    `RegionBranchOpInterface` by moving dominating operations whose only uses
+    are in a single conditionally-executed region into that region so that
+    executions paths where their results are not needed do not perform
+    unnecessary computations.
+
+    This is similar to loop-invariant code motion, which hoists operations out
+    of regions executed more than once.
+
+    It is recommended to run canonicalization first to remove unreachable
+    blocks: ops in unreachable blocks may prevent other operations from being
+    sunk as they may contain uses of their results
+  }];
+  let constructor = "::mlir::createControlFlowSinkPass()";
+  let statistics = [
+    Statistic<"numSunk", "num-sunk", "Number of operations sunk">,
+  ];
+}
+
 def CSE : Pass<"cse"> {
   let summary = "Eliminate common sub-expressions";
   let description = [{
     This pass implements a generalized algorithm for common sub-expression
     elimination. This pass relies on information provided by the
     `Memory SideEffect` interface to identify when it is safe to eliminate
     operations. See [Common subexpression elimination](https://en.wikipedia.org/wiki/Common_subexpression_elimination)
     for more general details on this optimization.

mlir/include/mlir/Transforms/Utils.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_UTILS_H
 #define MLIR_TRANSFORMS_UTILS_H
 
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Dominance.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 
 class AffineApplyOp;
 class AffineForOp;
 class Location;

 /// In the following cases, `sliceOps` remains empty:
 ///   1. If none of opInst's operands were the result of an affine.apply
 ///      (i.e., there was no affine computation slice to create).
 ///   2. If all the affine.apply op's supplying operands to this opInst did not
 ///      have any uses other than those in this opInst.
 void createAffineComputationSlice(Operation *opInst,
                                   SmallVectorImpl<AffineApplyOp> *sliceOps);
 
+/// Given a list of regions, perform control flow sinking on them. For each
+/// region, control-flow sinking moves operations that dominate the region but
+/// whose only users are in the region into the regions so that they aren't
+/// executed on paths where their results are not needed.
+///
+/// Example:
+///
+/// ```mlir
+/// %0 = arith.addi %arg0, %arg1
+/// scf.if %cond {
+///   scf.yield %0
+/// } else {
+///   scf.yield %arg2
+/// }
+/// ```
+///
+/// After control-flow sink:
+///
+/// ```mlir
+/// scf.if %cond {
+///   %0 = arith.addi %arg0, %arg1
+///   scf.yield %0
+/// } else {
+///   scf.yield %arg2
+/// }
+/// ```
+///
+/// Users must supply a callback `shouldMoveIntoRegion` that determines whether
+/// the given operation that only has users in the given operation should be
+/// moved into that region.
+///
+/// Returns the number of operations sunk.
+size_t
+controlFlowSink(ArrayRef<Region *> regions, DominanceInfo &domInfo,
+                function_ref<bool(Operation &, Region &)> shouldMoveIntoRegion);
+
+/// Populates `regions` with regions of the provided region branch op that are
+/// executed at most once at that are reachable given the current operands of
+/// the op. These regions can be passed to `controlFlowSink` to perform sinking
+/// on the regions of the operation.
+void getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,
+                                    SmallVectorImpl<Region *> &regions);
+
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_UTILS_H

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferOptimizations.cpp
   BufferResultsToOutParams.cpp
   BufferUtils.cpp
   Canonicalizer.cpp
+  ControlFlowSink.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
   NormalizeMemRefs.cpp
   OpStats.cpp

mlir/lib/Transforms/ControlFlowSink.cpp

+//===- ControlFlowSink.cpp - Code to perform control-flow sinking ---------===//
+//
+// 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 file implements a basic control-flow sink pass. Control-flow sinking
+// moves operations whose only uses are in conditionally-executed blocks in to
+// those blocks so that they aren't executed on paths where their results are
+// not needed.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/ControlFlowSink.h"
+#include "PassDetail.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Transforms/Utils.h"
+
+using namespace mlir;
+
+namespace {
+/// A basic control-flow sink pass. This pass analyzes the regions of operations
+/// that implement `RegionBranchOpInterface` that are reachable and executed at
+/// most once and sinks candidate operations that are side-effect free.
+struct ControlFlowSink : public ControlFlowSinkBase<ControlFlowSink> {
+  void runOnOperation() override;
+};
+} // end anonymous namespace
+
+/// Returns true if the given operation is side-effect free as are all of its
+/// nested operations.
+static bool isSideEffectFree(Operation &op) {
+  if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+    // If the op has side-effects, it cannot be moved.
+    if (!memInterface.hasNoEffect())
+      return false;
+    // If the op does not have recursive side effects, then it can be moved.
+    if (!op.hasTrait<OpTrait::HasRecursiveSideEffects>())
+      return true;
+  } else if (!op.hasTrait<OpTrait::HasRecursiveSideEffects>()) {
+    // Otherwise, if the op does not implement the memory effect interface and
+    // it does not have recursive side effects, then it cannot be known that the
+    // op is moveable.
+    return false;
+  }
+
+  // Recurse into the regions and ensure that all nested ops can also be moved.
+  for (auto &region : op.getRegions())
+    if (!llvm::all_of(region.getOps(), isSideEffectFree))
+      return false;
+  return true;
+}
+
+void ControlFlowSink::runOnOperation() {
+  auto &domInfo = getAnalysis<DominanceInfo>();
+  getOperation()->walk([&](RegionBranchOpInterface branch) {
+    SmallVector<Region *> regionsToSink;
+    mlir::getSinglyExecutedRegionsToSink(branch, regionsToSink);
+    numSunk = mlir::controlFlowSink(
+        regionsToSink, domInfo,
+        [](Operation &op, Region &) { return isSideEffectFree(op); });
+  });
+}
+
+std::unique_ptr<Pass> mlir::createControlFlowSinkPass() {
+  return std::make_unique<ControlFlowSink>();
+}

mlir/lib/Transforms/Utils/CMakeLists.txt

 add_mlir_library(MLIRTransformUtils
+  ControlFlowSinkUtils.cpp
   DialectConversion.cpp
   FoldUtils.cpp
   GreedyPatternRewriteDriver.cpp
   InliningUtils.cpp
   LoopFusionUtils.cpp
   LoopUtils.cpp
   RegionUtils.cpp
   Utils.cpp

mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

+//===- ControlFlowSinkUtils.cpp - Code to perform control-flow sinking ----===//
+//
+// 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 file implements utilityies for control-flow sinking. Control-flow
+// sinking moves operations whose only uses are in conditionally-executed blocks
+// into those blocks so that they aren't executed on paths where their results
+// are not needed.
+//
+// Control-flow sinking is not implemented on BranchOpInterface because
+// sinking ops into the successors of branch operations may move ops into loops.
+// It is idiomatic MLIR to perform optimizations at IR levels that readily
+// provide the necessary information.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Transforms/Utils.h"
+#include <deque>
+
+#define DEBUG_TYPE "cf-sink"
+
+using namespace mlir;
+
+namespace {
+/// A helper struct for control-flow sinking.
+class Sinker {
+public:
+  /// Create an operation sinker with given dominance info.
+  Sinker(function_ref<bool(Operation &, Region &)> shouldMoveIntoRegion,
+         DominanceInfo &domInfo)
+      : shouldMoveIntoRegion(shouldMoveIntoRegion), domInfo(domInfo),
+        numSunk(0) {}
+
+  /// Given a list of regions, find operations to sink and sink them. Return the
+  /// number of operations sunk.
+  size_t sinkRegions(ArrayRef<Region *> regions) &&;
+
+private:
+  /// Given a region and an op which dominates the region, returns true if all
+  /// users of the given op are dominated by the entry block of the region, and
+  /// thus the operation can be sunk into the region.
+  bool allUsersDominatedBy(Operation &op, Region &region);
+
+  /// Given a region and a top-level op (an op whose parent region is the given
+  /// region), determine whether the defining ops of the op's operands can be
+  /// sunk into the region.
+  ///
+  /// Add moved ops to the work queue.
+  void tryToSinkPredecessors(Operation &user, Region &region,
+                             std::deque<Operation *> &stack);
+
+  /// Iterate over all the ops in a region and try to sink their predecessors.
+  /// Recurse on subgraphs using a work queue.
+  void sinkRegion(Region &region);
+
+  /// The callback to determine whether an op should be moved in to a region.
+  function_ref<bool(Operation &, Region &)> shouldMoveIntoRegion;
+  /// Dominance info to determine op user dominance with respect to regions.
+  DominanceInfo &domInfo;
+  /// The number of operations sunk.
+  size_t numSunk;
+};
+} // end anonymous namespace
+
+bool Sinker::allUsersDominatedBy(Operation &op, Region &region) {
+  assert(region.findAncestorOpInRegion(op) == nullptr &&
+         "expected op to be defined outside the region");
+  return llvm::all_of(op.getUsers(), [&](Operation *user) {
+    // The user is dominated by the region if its containing block is dominated
+    // by the region's entry block.
+    return domInfo.dominates(&region.front(), user->getBlock());
+  });
+}
+
+void Sinker::tryToSinkPredecessors(Operation &user, Region &region,
+                                   std::deque<Operation *> &stack) {
+  LLVM_DEBUG(user.print(llvm::dbgs() << "\nContained op:\n"));
+  for (Value value : user.getOperands()) {
+    Operation *op = value.getDefiningOp();
+    // Ignore block arguments and ops that are already inside the region.
+    if (!op || op->getParentRegion() == &region)
+      continue;
+    LLVM_DEBUG(op->print(llvm::dbgs() << "\nTry to sink:\n"));
+
+    // If the op's users are all in the region and it can be moved, then do so.
+    if (allUsersDominatedBy(*op, region) && shouldMoveIntoRegion(*op, region)) {
+      // Move the op into the region's entry block. If the op is part of a
+      // subgraph, dependee ops would have been moved first, so inserting before
+      // the start of the block will ensure dominance is preserved. Ops can only
+      // be safely moved into the entry block as the region's other blocks may
+      // for a loop.
+      op->moveBefore(&region.front(), region.front().begin());
+      ++numSunk;
+      // Add the op to the work queue.
+      stack.push_back(op);
+    }
+  }
+}
+
+void Sinker::sinkRegion(Region &region) {
+  // Initialize the work queue with all the ops in the region.
+  std::deque<Operation *> stack;
+  for (auto &op : region.getOps())
+    stack.push_back(&op);
+
+  // Process all the ops depth-first. This ensures that nodes of subgraphs are
+  // sunk in the correct order.
+  while (!stack.empty()) {
+    Operation *op = stack.back();
+    stack.pop_back();
+    tryToSinkPredecessors(*op, region, stack);
+  }
+}
+
+size_t Sinker::sinkRegions(ArrayRef<Region *> regions) && {
+  for (auto *region : regions) {
+    if (region->empty())
+      continue;
+    sinkRegion(*region);
+  }
+  return numSunk;
+}
+
+size_t mlir::controlFlowSink(
+    ArrayRef<Region *> regions, DominanceInfo &domInfo,
+    function_ref<bool(Operation &, Region &)> shouldMoveIntoRegion) {
+  return Sinker(shouldMoveIntoRegion, domInfo).sinkRegions(regions);
+}
+
+void mlir::getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,
+                                          SmallVectorImpl<Region *> &regions) {
+  // Collect constant operands.
+  SmallVector<Attribute> operands(branch->getNumOperands(), Attribute());
+  for (auto &it : llvm::enumerate(branch->getOperands()))
+    matchPattern(it.value(), m_Constant(&operands[it.index()]));
+  // Get the invocation bounds.
+  SmallVector<InvocationBounds> bounds;
+  branch.getRegionInvocationBounds(operands, bounds);
+
+  // For a simple control-flow sink, only consider regions that are executed at
+  // most once.
+  for (Region &region : branch->getRegions()) {
+    InvocationBounds &bound = bounds[region.getRegionNumber()];
+    if (bound.getUpperBound() && *bound.getUpperBound() <= 1)
+      regions.push_back(&region);
+  }
+}

mlir/test/Transforms/control-flow-sink.mlir

+// RUN: mlir-opt -split-input-file -control-flow-sink %s | FileCheck %s
+
+// Test that operations can be sunk.
+
+// CHECK-LABEL: @test_simple_sink
+// CHECK: %0 = arith.subi %arg2, %arg1 : i32
+// CHECK: %1 = test.region_if %arg0: i1 -> i32 then {
+// CHECK:   %2 = arith.subi %arg1, %arg2 : i32
+// CHECK:   test.region_if_yield %2 : i32
+// CHECK: } else {
+// CHECK:   %2 = arith.addi %arg1, %arg1 : i32
+// CHECK:   %3 = arith.addi %0, %2 : i32
+// CHECK:   test.region_if_yield %3 : i32
+// CHECK: } join {
+// CHECK:   %2 = arith.addi %arg2, %arg2 : i32
+// CHECK:   %3 = arith.addi %2, %0 : i32
+// CHECK:   test.region_if_yield %3 : i32
+// CHECK: }
+// CHECK: return %1 : i32
+func @test_simple_sink(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
+  %0 = arith.subi %arg0, %arg1 : i32
+  %1 = arith.subi %arg1, %arg0 : i32
+  %2 = arith.addi %arg0, %arg0 : i32
+  %3 = arith.addi %arg1, %arg1 : i32
+  %result = test.region_if %cond : i1 -> i32 then {
+    test.region_if_yield %0 : i32
+  } else {
+    %v = arith.addi %1, %2 : i32
+    test.region_if_yield %v : i32
+  } join {
+    %v = arith.addi %3, %1 : i32
+    test.region_if_yield %v : i32
+  }
+  return %result : i32
+}
+
+// -----
+
+// Test that a region op can be sunk.
+
+// CHECK-LABEL: @test_region_sink
+// CHECK: %0 = test.region_if %arg0: i1 -> i32 then {
+// CHECK:   %1 = test.region_if %arg0: i1 -> i32 then {
+// CHECK:     test.region_if_yield %arg1 : i32
+// CHECK:   } else {
+// CHECK:     %2 = arith.subi %arg1, %arg2 : i32
+// CHECK:     test.region_if_yield %2 : i32
+// CHECK:   } join {
+// CHECK:     test.region_if_yield %arg2 : i32
+// CHECK:   }
+// CHECK:   test.region_if_yield %1 : i32
+// CHECK: } else {
+// CHECK:   test.region_if_yield %arg1 : i32
+// CHECK: } join {
+// CHECK:   test.region_if_yield %arg2 : i32
+// CHECK: }
+// CHECK: return %0 : i32
+func @test_region_sink(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
+  %0 = arith.subi %arg0, %arg1 : i32
+  %result0 = test.region_if %cond : i1 -> i32 then {
+    test.region_if_yield %arg0 : i32
+  } else {
+    test.region_if_yield %0 : i32
+  } join {
+    test.region_if_yield %arg1 : i32
+  }
+  %result1 = test.region_if %cond : i1 -> i32 then {
+    test.region_if_yield %result0 : i32
+  } else {
+    test.region_if_yield %arg0 : i32
+  } join {
+    test.region_if_yield %arg1 : i32
+  }
+  return %result1 : i32
+}
+
+// -----
+
+// Test that an entire subgraph can be sunk.
+
+// CHECK-LABEL: @test_subgraph_sink
+// CHECK: %0 = test.region_if %arg0: i1 -> i32 then {
+// CHECK:   %1 = arith.subi %arg1, %arg2 : i32
+// CHECK:   %2 = arith.addi %arg1, %arg2 : i32
+// CHECK:   %3 = arith.subi %arg2, %arg1 : i32
+// CHECK:   %4 = arith.muli %3, %3 : i32
+// CHECK:   %5 = arith.muli %2, %1 : i32
+// CHECK:   %6 = arith.addi %5, %4 : i32
+// CHECK:   test.region_if_yield %6 : i32
+// CHECK: } else {
+// CHECK:   test.region_if_yield %arg1 : i32
+// CHECK: } join {
+// CHECK:   test.region_if_yield %arg2 : i32
+// CHECK: }
+// CHECK: return %0 : i32
+func @test_subgraph_sink(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
+  %0 = arith.addi %arg0, %arg1 : i32
+  %1 = arith.subi %arg0, %arg1 : i32
+  %2 = arith.subi %arg1, %arg0 : i32
+  %3 = arith.muli %0, %1 : i32
+  %4 = arith.muli %2, %2 : i32
+  %5 = arith.addi %3, %4 : i32
+  %result = test.region_if %cond : i1 -> i32 then {
+    test.region_if_yield %5 : i32
+  } else {
+    test.region_if_yield %arg0 : i32
+  } join {
+    test.region_if_yield %arg1 : i32
+  }
+  return %result : i32
+}
+
+// -----
+
+// Test that ops can be sunk into regions with multiple blocks.
+
+// CHECK-LABEL: @test_multiblock_region_sink
+// CHECK: %0 = arith.addi %arg1, %arg2 : i32
+// CHECK: %1 = "test.any_cond"() ( {
+// CHECK:   %3 = arith.addi %0, %arg2 : i32
+// CHECK:   %4 = arith.addi %3, %arg1 : i32
+// CHECK:   br ^bb1(%4 : i32)
+// CHECK: ^bb1(%5: i32):  // pred: ^bb0
+// CHECK:   %6 = arith.addi %5, %4 : i32
+// CHECK:   "test.yield"(%6) : (i32) -> ()
+// CHECK: }) : () -> i32
+// CHECK: %2 = arith.addi %0, %1 : i32
+// CHECK: return %2 : i32
+func @test_multiblock_region_sink(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
+  // Value %0 is used in a nested region but also at this level
+  %0 = arith.addi %arg0, %arg1 : i32
+  %1 = arith.addi %0, %arg1 : i32
+  %2 = arith.addi %1, %arg0 : i32
+  %3 = "test.any_cond"() ({
+  ^bb0:
+    // Value %2 is used in both blocks.
+    br ^bb1(%2 : i32)
+  ^bb1(%arg2 : i32):
+    %4 = arith.addi %arg2, %2 : i32
+    "test.yield"(%4) : (i32) -> ()
+  }) : () -> i32
+  %result = arith.addi %0, %3 : i32
+  return %result : i32
+}
+
+// -----
+
+// Test that ops can be sunk recursively into nested regions.
+
+// CHECK-LABEL: @test_nested_region_sink
+// CHECK: %0 = test.region_if %arg0: i1 -> i32 then {
+// CHECK:   %1 = test.region_if %arg0: i1 -> i32 then {
+// CHECK:     %2 = arith.addi %arg1, %arg1 : i32
+// CHECK:     test.region_if_yield %2 : i32
+// CHECK:   } else {
+// CHECK:     test.region_if_yield %arg1 : i32
+// CHECK:   } join {
+// CHECK:     test.region_if_yield %arg1 : i32
+// CHECK:   }
+// CHECK:   test.region_if_yield %1 : i32
+// CHECK: } else {
+// CHECK:   test.region_if_yield %arg1 : i32
+// CHECK: } join {
+// CHECK:   test.region_if_yield %arg1 : i32
+// CHECK: }
+// CHECK: return %0 : i32
+func @test_nested_region_sink(%cond : i1, %arg0 : i32) -> i32 {
+  %0 = arith.addi %arg0, %arg0 : i32
+  %1 = test.region_if %cond  : i1 -> i32 then {
+    %2 = test.region_if %cond : i1 -> i32 then {
+      test.region_if_yield %0 : i32
+    } else {
+      test.region_if_yield %arg0 : i32
+    } join {
+      test.region_if_yield %arg0 : i32
+    }
+    test.region_if_yield %2 : i32
+  } else {
+    test.region_if_yield %arg0 : i32
+  } join {
+    test.region_if_yield %arg0 : i32
+  }
+  return %1 : i32
+}
+
+// -----
+
+// Test that ops are only moved into the entry block, even when their only uses
+// are further along.
+
+// CHECK-LABEL: @test_not_sunk_deeply
+// CHECK: %0 = "test.any_cond"() ( {
+// CHECK:   %1 = arith.addi %arg0, %arg0 : i32
+// CHECK:   br ^bb1
+// CHECK: ^bb1:  // pred: ^bb0
+// CHECK:   "test.yield"(%1) : (i32) -> ()
+// CHECK: }) : () -> i32
+// CHECK: return %0 : i32
+func @test_not_sunk_deeply(%arg0 : i32) -> i32 {
+  %0 = arith.addi %arg0, %arg0 : i32
+  %1 = "test.any_cond"() ({
+  ^bb0:
+    br ^bb1
+  ^bb1:
+    "test.yield"(%0) : (i32) -> ()
+  }) : () -> i32
+  return %1 : i32
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (index.hasValue()) {
     return;
   }
 
   // The then and else regions are the entry regions of this op.
   regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));
   regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));
 }
 
+void RegionIfOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands,
+    SmallVectorImpl<InvocationBounds> &invocationBounds) {
+  // Each region is invoked at most once.
+  invocationBounds.assign(/*NumElts=*/3, /*Elt=*/{0, 1});
+}
+
+//===----------------------------------------------------------------------===//
+// AnyCondOp
+//===----------------------------------------------------------------------===//
+
+void AnyCondOp::getSuccessorRegions(Optional<unsigned> index,
+                                    ArrayRef<Attribute> operands,
+                                    SmallVectorImpl<RegionSuccessor> &regions) {
+  // The parent op branches into the only region, and the region branches back
+  // to the parent op.
+  if (index)
+    regions.emplace_back(&getRegion());
+  else
+    regions.emplace_back(getResults());
+}
+
+void AnyCondOp::getNumRegionInvocations(
+    ArrayRef<Attribute> operands, SmallVectorImpl<int64_t> &countPerRegion) {
+  countPerRegion.assign(/*NumElts=*/1, /*Elt=*/1);
+}
+
 //===----------------------------------------------------------------------===//
 // SingleNoTerminatorCustomAsmOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseSingleNoTerminatorCustomAsmOp(OpAsmParser &parser,
                                                       OperationState &state) {
   Region *body = state.addRegion();
   if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))

mlir/test/lib/Dialect/Test/TestOps.td

def RegionIfYieldOp : TEST_Op<"region_if_yield",
       [NoSideEffect, ReturnLike, Terminator]> {
   let arguments = (ins Variadic<AnyType>:$results);
   let assemblyFormat = [{
     $results `:` type($results) attr-dict
   }];
 }
 
 def RegionIfOp : TEST_Op<"region_if",
-      [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                                 ["getRegionInvocationBounds"]>,
        SingleBlockImplicitTerminator<"RegionIfYieldOp">,
        RecursiveSideEffects]> {
   let description =[{
     Represents an abstract if-then-else-join pattern. In this context, the then
     and else regions jump to the join region, which finally returns to its
     parent op.
-    }];
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parseRegionIfOp(parser, result); }];
   let arguments = (ins Variadic<AnyType>);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$thenRegion,
                         AnyRegion:$elseRegion,
                         AnyRegion:$joinRegion);
   let extraClassDeclaration = [{
     ::mlir::Block::BlockArgListType getThenArgs() {
       return getBody(0)->getArguments();
     }
     ::mlir::Block::BlockArgListType getElseArgs() {
       return getBody(1)->getArguments();
     }
     ::mlir::Block::BlockArgListType getJoinArgs() {
       return getBody(2)->getArguments();
     }
     ::mlir::OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 }
 
+def AnyCondOp : TEST_Op<"any_cond",
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                                 ["getNumRegionInvocations"]>,
+       RecursiveSideEffects]> {
+  let results = (outs Variadic<AnyType>:$results);
+  let regions = (region AnyRegion:$region);
+}
+
 //===----------------------------------------------------------------------===//
 // Test TableGen generated build() methods
 //===----------------------------------------------------------------------===//
 
 def TableGenConstant : TEST_Op<"tblgen_constant"> {
   let results = (outs AnyType);
 }