1 //===- LoopInvariantCodeMotion.cpp - Code to perform loop fusion-----------===//
2 //
3 // Copyright 2019 The MLIR Authors.
4 //
5 // Licensed under the Apache License, Version 2.0 (the "License");
6 // you may not use this file except in compliance with the License.
7 // You may obtain a copy of the License at
8 //
9 //   http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing, software
12 // distributed under the License is distributed on an "AS IS" BASIS,
13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 // See the License for the specific language governing permissions and
15 // limitations under the License.
16 // =============================================================================
17 //
18 // This file implements loop invariant code motion.
19 //
20 //===----------------------------------------------------------------------===//
21 
22 #include "mlir/Transforms/Passes.h"
23 
24 #include "mlir/IR/Builders.h"
25 #include "mlir/IR/Function.h"
26 #include "mlir/Pass/Pass.h"
27 #include "mlir/Transforms/LoopLikeInterface.h"
28 #include "mlir/Transforms/SideEffectsInterface.h"
29 #include "llvm/ADT/SmallPtrSet.h"
30 #include "llvm/Support/CommandLine.h"
31 #include "llvm/Support/Debug.h"
32 
33 #define DEBUG_TYPE "licm"
34 
35 using namespace mlir;
36 
37 namespace {
38 
39 using SideEffecting = SideEffectsInterface::SideEffecting;
40 
41 /// Loop invariant code motion (LICM) pass.
42 struct LoopInvariantCodeMotion : public OperationPass<LoopInvariantCodeMotion> {
43 public:
44   void runOnOperation() override;
45 };
46 
47 // Checks whether the given op can be hoisted by checking that
48 // - the op and any of its contained operations do not depend on SSA values
49 //   defined inside of the loop (by means of calling definedOutside).
50 // - the op has no side-effects. If sideEffecting is Never, sideeffects of this
51 //   op and its nested ops are ignored.
52 static bool canBeHoisted(Operation *op,
53                          llvm::function_ref<bool(Value *)> definedOutside,
54                          SideEffecting sideEffecting,
55                          SideEffectsInterface &interface) {
56   // Check that dependencies are defined outside of loop.
57   if (!llvm::all_of(op->getOperands(), definedOutside))
58     return false;
59   // Check whether this op is side-effect free. If we already know that there
60   // can be no side-effects because the surrounding op has claimed so, we can
61   // (and have to) skip this step.
62   auto thisOpIsSideEffecting = sideEffecting;
63   if (thisOpIsSideEffecting != SideEffecting::Never) {
64     thisOpIsSideEffecting = interface.isSideEffecting(op);
65     // If the op always has sideeffects, we cannot hoist.
66     if (thisOpIsSideEffecting == SideEffecting::Always)
67       return false;
68   }
69   // Recurse into the regions for this op and check whether the contained ops
70   // can be hoisted.
71   for (auto &region : op->getRegions()) {
72     for (auto &block : region.getBlocks()) {
73       for (auto &innerOp : block) {
74         if (innerOp.isKnownTerminator())
75           continue;
76         if (!canBeHoisted(&innerOp, definedOutside, thisOpIsSideEffecting,
77                           interface))
78           return false;
79       }
80     }
81   }
82   return true;
83 }
84 
85 static LogicalResult moveLoopInvariantCode(LoopLikeOpInterface looplike,
86                                            SideEffectsInterface &interface) {
87   auto &loopBody = looplike.getLoopBody();
88 
89   // We use two collections here as we need to preserve the order for insertion
90   // and this is easiest.
91   SmallPtrSet<Operation *, 8> willBeMovedSet;
92   SmallVector<Operation *, 8> opsToMove;
93 
94   // Helper to check whether an operation is loop invariant wrt. SSA properties.
95   auto isDefinedOutsideOfBody = [&](Value *value) {
96     auto definingOp = value->getDefiningOp();
97     return (definingOp && !!willBeMovedSet.count(definingOp)) ||
98            looplike.isDefinedOutsideOfLoop(value);
99   };
100 
101   // Do not use walk here, as we do not want to go into nested regions and hoist
102   // operations from there. These regions might have semantics unknown to this
103   // rewriting. If the nested regions are loops, they will have been processed.
104   for (auto &block : loopBody) {
105     for (auto &op : block.without_terminator()) {
106       if (canBeHoisted(&op, isDefinedOutsideOfBody,
107                        mlir::SideEffectsDialectInterface::Recursive,
108                        interface)) {
109         opsToMove.push_back(&op);
110         willBeMovedSet.insert(&op);
111       }
112     }
113   }
114 
115   // For all instructions that we found to be invariant, move outside of the
116   // loop.
117   auto result = looplike.moveOutOfLoop(opsToMove);
118   LLVM_DEBUG(looplike.print(llvm::dbgs() << "Modified loop\n"));
119   return result;
120 }
121 
122 } // end anonymous namespace
123 
124 void LoopInvariantCodeMotion::runOnOperation() {
125   SideEffectsInterface interface(&getContext());
126   // Walk through all loops in a function in innermost-loop-first order. This
127   // way, we first LICM from the inner loop, and place the ops in
128   // the outer loop, which in turn can be further LICM'ed.
129   getOperation()->walk([&](Operation *op) {
130     if (auto looplike = dyn_cast<LoopLikeOpInterface>(op)) {
131       LLVM_DEBUG(op->print(llvm::dbgs() << "\nOriginal loop\n"));
132       if (failed(moveLoopInvariantCode(looplike, interface)))
133         signalPassFailure();
134     }
135   });
136 }
137 
138 // Include the generated code for the loop-like interface here, as it otherwise
139 // has no compilation unit. This works as loop-invariant code motion is the
140 // only user of that interface.
141 #include "mlir/Transforms/LoopLikeInterface.cpp.inc"
142 
143 std::unique_ptr<Pass> mlir::createLoopInvariantCodeMotionPass() {
144   return std::make_unique<LoopInvariantCodeMotion>();
145 }
146 
147 static PassRegistration<LoopInvariantCodeMotion>
148     pass("loop-invariant-code-motion",
149          "Hoist loop invariant instructions outside of the loop");
150