xref: /llvm-project/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp (revision bb2f87af0ac9ad8cc0efcb523e108e715ee9ab93)
1 //===- LoopInvariantCodeMotion.cpp - Code to perform loop fusion-----------===//
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 loop invariant code motion.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "PassDetail.h"
14 #include "mlir/Transforms/Passes.h"
15 
16 #include "mlir/IR/Builders.h"
17 #include "mlir/IR/BuiltinOps.h"
18 #include "mlir/Interfaces/LoopLikeInterface.h"
19 #include "mlir/Interfaces/SideEffectInterfaces.h"
20 #include "mlir/Transforms/LoopUtils.h"
21 #include "llvm/ADT/SmallPtrSet.h"
22 #include "llvm/Support/CommandLine.h"
23 #include "llvm/Support/Debug.h"
24 
25 #define DEBUG_TYPE "licm"
26 
27 using namespace mlir;
28 
29 namespace {
30 /// Loop invariant code motion (LICM) pass.
31 struct LoopInvariantCodeMotion
32     : public LoopInvariantCodeMotionBase<LoopInvariantCodeMotion> {
33   void runOnOperation() override;
34 };
35 } // namespace
36 
37 // Checks whether the given op can be hoisted by checking that
38 // - the op and any of its contained operations do not depend on SSA values
39 //   defined inside of the loop (by means of calling definedOutside).
40 // - the op has no side-effects. If sideEffecting is Never, sideeffects of this
41 //   op and its nested ops are ignored.
42 static bool canBeHoisted(Operation *op,
43                          function_ref<bool(Value)> definedOutside) {
44   // Check that dependencies are defined outside of loop.
45   if (!llvm::all_of(op->getOperands(), definedOutside))
46     return false;
47   // Check whether this op is side-effect free. If we already know that there
48   // can be no side-effects because the surrounding op has claimed so, we can
49   // (and have to) skip this step.
50   if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
51     if (!memInterface.hasNoEffect())
52       return false;
53     // If the operation doesn't have side effects and it doesn't recursively
54     // have side effects, it can always be hoisted.
55     if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>())
56       return true;
57 
58     // Otherwise, if the operation doesn't provide the memory effect interface
59     // and it doesn't have recursive side effects we treat it conservatively as
60     // side-effecting.
61   } else if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
62     return false;
63   }
64 
65   // Recurse into the regions for this op and check whether the contained ops
66   // can be hoisted.
67   for (auto &region : op->getRegions()) {
68     for (auto &block : region) {
69       for (auto &innerOp : block)
70         if (!canBeHoisted(&innerOp, definedOutside))
71           return false;
72     }
73   }
74   return true;
75 }
76 
77 LogicalResult mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) {
78   auto &loopBody = looplike.getLoopBody();
79 
80   // We use two collections here as we need to preserve the order for insertion
81   // and this is easiest.
82   SmallPtrSet<Operation *, 8> willBeMovedSet;
83   SmallVector<Operation *, 8> opsToMove;
84 
85   // Helper to check whether an operation is loop invariant wrt. SSA properties.
86   auto isDefinedOutsideOfBody = [&](Value value) {
87     auto *definingOp = value.getDefiningOp();
88     return (definingOp && !!willBeMovedSet.count(definingOp)) ||
89            looplike.isDefinedOutsideOfLoop(value);
90   };
91 
92   // Do not use walk here, as we do not want to go into nested regions and hoist
93   // operations from there. These regions might have semantics unknown to this
94   // rewriting. If the nested regions are loops, they will have been processed.
95   for (auto &block : loopBody) {
96     for (auto &op : block.without_terminator()) {
97       if (canBeHoisted(&op, isDefinedOutsideOfBody)) {
98         opsToMove.push_back(&op);
99         willBeMovedSet.insert(&op);
100       }
101     }
102   }
103 
104   // For all instructions that we found to be invariant, move outside of the
105   // loop.
106   auto result = looplike.moveOutOfLoop(opsToMove);
107   LLVM_DEBUG(looplike.print(llvm::dbgs() << "\n\nModified loop:\n"));
108   return result;
109 }
110 
111 void LoopInvariantCodeMotion::runOnOperation() {
112   // Walk through all loops in a function in innermost-loop-first order. This
113   // way, we first LICM from the inner loop, and place the ops in
114   // the outer loop, which in turn can be further LICM'ed.
115   getOperation()->walk([&](LoopLikeOpInterface loopLike) {
116     LLVM_DEBUG(loopLike.print(llvm::dbgs() << "\nOriginal loop:\n"));
117     if (failed(moveLoopInvariantCode(loopLike)))
118       signalPassFailure();
119   });
120 }
121 
122 std::unique_ptr<Pass> mlir::createLoopInvariantCodeMotionPass() {
123   return std::make_unique<LoopInvariantCodeMotion>();
124 }
125