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