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 <iomanip>
23 #include <sstream>
24 
25 #include "mlir/AffineOps/AffineOps.h"
26 #include "mlir/Analysis/AffineAnalysis.h"
27 #include "mlir/Analysis/AffineStructures.h"
28 #include "mlir/Analysis/LoopAnalysis.h"
29 #include "mlir/Analysis/SliceAnalysis.h"
30 #include "mlir/Analysis/Utils.h"
31 #include "mlir/IR/AffineExpr.h"
32 #include "mlir/IR/AffineMap.h"
33 #include "mlir/IR/Builders.h"
34 #include "mlir/Pass/Pass.h"
35 #include "mlir/StandardOps/Ops.h"
36 #include "mlir/Transforms/LoopUtils.h"
37 #include "mlir/Transforms/Passes.h"
38 #include "mlir/Transforms/Utils.h"
39 #include "llvm/ADT/DenseMap.h"
40 #include "llvm/ADT/DenseSet.h"
41 #include "llvm/ADT/SetVector.h"
42 #include "llvm/Support/CommandLine.h"
43 #include "llvm/Support/Debug.h"
44 #include "llvm/Support/raw_ostream.h"
45 
46 #define DEBUG_TYPE "licm"
47 
48 using llvm::SetVector;
49 
50 using namespace mlir;
51 
52 namespace {
53 
54 /// Loop invariant code motion (LICM) pass.
55 /// TODO(asabne) : The pass is missing zero-trip tests.
56 /// TODO(asabne) : Check for the presence of side effects before hoisting.
57 struct LoopInvariantCodeMotion : public FunctionPass<LoopInvariantCodeMotion> {
58   void runOnFunction() override;
59   void runOnAffineForOp(AffineForOp forOp);
60   std::vector<AffineForOp> forOps;
61 };
62 } // end anonymous namespace
63 
64 FunctionPassBase *mlir::createLoopInvariantCodeMotionPass() {
65   return new LoopInvariantCodeMotion();
66 }
67 
68 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
69   auto *loopBody = forOp.getBody();
70 
71   // This is the place where hoisted instructions would reside.
72   FuncBuilder b(forOp.getOperation());
73 
74   // This vector is used to place loop invariant operations.
75   SmallVector<Operation *, 8> opsToMove;
76 
77   SetVector<Operation *> loopDefinedOps;
78   // Generate forward slice which contains ops that fall under the transitive
79   // definition closure following the loop induction variable.
80   getForwardSlice(forOp, &loopDefinedOps);
81 
82   LLVM_DEBUG(for (auto i
83                   : loopDefinedOps) {
84     (i->print(llvm::dbgs() << "\nLoop-dependent op\n"));
85   });
86 
87   for (auto &op : *loopBody) {
88     // If the operation is loop invariant, insert it into opsToMove.
89     if (!op.isa<AffineForOp>() && !op.isa<AffineTerminatorOp>() &&
90         loopDefinedOps.count(&op) != 1) {
91       LLVM_DEBUG(op.print(llvm::dbgs() << "\nLICM'ing op\n"));
92       opsToMove.push_back(&op);
93     }
94   }
95 
96   // For all instructions that we found to be invariant, place them sequentially
97   // right before the for loop.
98   for (auto *op : opsToMove) {
99     op->moveBefore(forOp);
100   }
101 
102   LLVM_DEBUG(forOp.getOperation()->print(llvm::dbgs() << "\nModified loop\n"));
103 
104   // If the for loop body has a single operation (the terminator), erase it.
105   if (forOp.getBody()->getOperations().size() == 1) {
106     assert(forOp.getBody()->getOperations().front().isa<AffineTerminatorOp>());
107     forOp.erase();
108   }
109 }
110 
111 void LoopInvariantCodeMotion::runOnFunction() {
112   forOps.clear();
113 
114   // Gather all loops in a function, and order them in innermost-loop-first
115   // order. This way, we first LICM from the inner loop, and place the ops in
116   // the outer loop, which in turn can be further LICM'ed. This saves iterating
117   // on the inner loop operations while LICMing through the outer loop.
118   getFunction().walk<AffineForOp>(
119       [&](AffineForOp forOp) { forOps.push_back(forOp); });
120   // We gather loops first, and then go over them later because we don't want to
121   // mess the iterators up.
122   for (auto op : forOps) {
123     LLVM_DEBUG(op.getOperation()->print(llvm::dbgs() << "\nOriginal loop\n"));
124     runOnAffineForOp(op);
125   }
126 }
127 
128 static PassRegistration<LoopInvariantCodeMotion>
129     pass("loop-invariant-code-motion",
130          "Hoist loop invariant instructions outside of the loop");
131