mlir/lib/Transforms/Utils/ControlFlowSinkUtils.cpp

   1 //===- ControlFlowSinkUtils.cpp - Code to perform control-flow sinking ----===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements utilities for control-flow sinking. Control-flow
  10 // sinking moves operations whose only uses are in conditionally-executed blocks
  11 // into those blocks so that they aren't executed on paths where their results
  12 // are not needed.
  13 //
  14 // Control-flow sinking is not implemented on BranchOpInterface because
  15 // sinking ops into the successors of branch operations may move ops into loops.
  16 // It is idiomatic MLIR to perform optimizations at IR levels that readily
  17 // provide the necessary information.
  18 //
  19 //===----------------------------------------------------------------------===//
  20
  21 #include "mlir/Transforms/ControlFlowSinkUtils.h"
  22 #include "mlir/IR/Dominance.h"
  23 #include "mlir/IR/Matchers.h"
  24 #include "mlir/Interfaces/ControlFlowInterfaces.h"
  25 #include <vector>
  26
  27 #define DEBUG_TYPE "cf-sink"
  28
  29 using namespace mlir;
  30
  31 namespace {
  32 /// A helper struct for control-flow sinking.
  33 class Sinker {
  34 public:
  35   /// Create an operation sinker with given dominance info.
  36   Sinker(function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,
  37          function_ref<void(Operation *, Region *)> moveIntoRegion,
  38          DominanceInfo &domInfo)
  39       : shouldMoveIntoRegion(shouldMoveIntoRegion),
  40         moveIntoRegion(moveIntoRegion), domInfo(domInfo) {}
  41
  42   /// Given a list of regions, find operations to sink and sink them. Return the
  43   /// number of operations sunk.
  44   size_t sinkRegions(RegionRange regions);
  45
  46 private:
  47   /// Given a region and an op which dominates the region, returns true if all
  48   /// users of the given op are dominated by the entry block of the region, and
  49   /// thus the operation can be sunk into the region.
  50   bool allUsersDominatedBy(Operation *op, Region *region);
  51
  52   /// Given a region and a top-level op (an op whose parent region is the given
  53   /// region), determine whether the defining ops of the op's operands can be
  54   /// sunk into the region.
  55   ///
  56   /// Add moved ops to the work queue.
  57   void tryToSinkPredecessors(Operation *user, Region *region,
  58                              std::vector<Operation *> &stack);
  59
  60   /// Iterate over all the ops in a region and try to sink their predecessors.
  61   /// Recurse on subgraphs using a work queue.
  62   void sinkRegion(Region *region);
  63
  64   /// The callback to determine whether an op should be moved in to a region.
  65   function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion;
  66   /// The calback to move an operation into the region.
  67   function_ref<void(Operation *, Region *)> moveIntoRegion;
  68   /// Dominance info to determine op user dominance with respect to regions.
  69   DominanceInfo &domInfo;
  70   /// The number of operations sunk.
  71   size_t numSunk = 0;
  72 };
  73 } // end anonymous namespace
  74
  75 bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {
  76   assert(region->findAncestorOpInRegion(*op) == nullptr &&
  77          "expected op to be defined outside the region");
  78   return llvm::all_of(op->getUsers(), [&](Operation *user) {
  79     // The user is dominated by the region if its containing block is dominated
  80     // by the region's entry block.
  81     return domInfo.dominates(&region->front(), user->getBlock());
  82   });
  83 }
  84
  85 void Sinker::tryToSinkPredecessors(Operation *user, Region *region,
  86                                    std::vector<Operation *> &stack) {
  87   LLVM_DEBUG(user->print(llvm::dbgs() << "\nContained op:\n"));
  88   for (Value value : user->getOperands()) {
  89     Operation *op = value.getDefiningOp();
  90     // Ignore block arguments and ops that are already inside the region.
  91     if (!op || op->getParentRegion() == region)
  92       continue;
  93     LLVM_DEBUG(op->print(llvm::dbgs() << "\nTry to sink:\n"));
  94
  95     // If the op's users are all in the region and it can be moved, then do so.
  96     if (allUsersDominatedBy(op, region) && shouldMoveIntoRegion(op, region)) {
  97       moveIntoRegion(op, region);
  98       ++numSunk;
  99       // Add the op to the work queue.
 100       stack.push_back(op);
 101     }
 102   }
 103 }
 104
 105 void Sinker::sinkRegion(Region *region) {
 106   // Initialize the work queue with all the ops in the region.
 107   std::vector<Operation *> stack;
 108   for (Operation &op : region->getOps())
 109     stack.push_back(&op);
 110
 111   // Process all the ops depth-first. This ensures that nodes of subgraphs are
 112   // sunk in the correct order.
 113   while (!stack.empty()) {
 114     Operation *op = stack.back();
 115     stack.pop_back();
 116     tryToSinkPredecessors(op, region, stack);
 117   }
 118 }
 119
 120 size_t Sinker::sinkRegions(RegionRange regions) {
 121   for (Region *region : regions)
 122     if (!region->empty())
 123       sinkRegion(region);
 124   return numSunk;
 125 }
 126
 127 size_t mlir::controlFlowSink(
 128     RegionRange regions, DominanceInfo &domInfo,
 129     function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,
 130     function_ref<void(Operation *, Region *)> moveIntoRegion) {
 131   return Sinker(shouldMoveIntoRegion, moveIntoRegion, domInfo)
 132       .sinkRegions(regions);
 133 }
 134
 135 void mlir::getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,
 136                                           SmallVectorImpl<Region *> &regions) {
 137   // Collect constant operands.
 138   SmallVector<Attribute> operands(branch->getNumOperands(), Attribute());
 139   for (auto [idx, operand] : llvm::enumerate(branch->getOperands()))
 140     (void)matchPattern(operand, m_Constant(&operands[idx]));
 141
 142   // Get the invocation bounds.
 143   SmallVector<InvocationBounds> bounds;
 144   branch.getRegionInvocationBounds(operands, bounds);
 145
 146   // For a simple control-flow sink, only consider regions that are executed at
 147   // most once.
 148   for (auto it : llvm::zip(branch->getRegions(), bounds)) {
 149     const InvocationBounds &bound = std::get<1>(it);
 150     if (bound.getUpperBound() && *bound.getUpperBound() <= 1)
 151       regions.push_back(&std::get<0>(it));
 152   }
 153 }