1 //===- LoopUnrollAndJam.cpp - Code to perform loop unroll and jam ---------===// 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 unroll and jam. Unroll and jam is a transformation 10 // that improves locality, in particular, register reuse, while also improving 11 // operation level parallelism. The example below shows what it does in nearly 12 // the general case. Loop unroll and jam currently works if the bounds of the 13 // loops inner to the loop being unroll-jammed do not depend on the latter. 14 // 15 // Before After unroll and jam of i by factor 2: 16 // 17 // for i, step = 2 18 // for i S1(i); 19 // S1; S2(i); 20 // S2; S1(i+1); 21 // for j S2(i+1); 22 // S3; for j 23 // S4; S3(i, j); 24 // S5; S4(i, j); 25 // S6; S3(i+1, j) 26 // S4(i+1, j) 27 // S5(i); 28 // S6(i); 29 // S5(i+1); 30 // S6(i+1); 31 // 32 // Note: 'if/else' blocks are not jammed. So, if there are loops inside if 33 // op's, bodies of those loops will not be jammed. 34 //===----------------------------------------------------------------------===// 35 36 #include "PassDetail.h" 37 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h" 38 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h" 39 #include "mlir/Dialect/Affine/IR/AffineOps.h" 40 #include "mlir/Dialect/Affine/LoopUtils.h" 41 #include "mlir/Dialect/Affine/Passes.h" 42 #include "mlir/IR/AffineExpr.h" 43 #include "mlir/IR/AffineMap.h" 44 #include "mlir/IR/BlockAndValueMapping.h" 45 #include "mlir/IR/Builders.h" 46 #include "llvm/ADT/DenseMap.h" 47 #include "llvm/Support/CommandLine.h" 48 49 using namespace mlir; 50 51 #define DEBUG_TYPE "affine-loop-unroll-jam" 52 53 namespace { 54 /// Loop unroll jam pass. Currently, this just unroll jams the first 55 /// outer loop in a Function. 56 struct LoopUnrollAndJam : public AffineLoopUnrollAndJamBase<LoopUnrollAndJam> { 57 explicit LoopUnrollAndJam(Optional<unsigned> unrollJamFactor = None) { 58 if (unrollJamFactor) 59 this->unrollJamFactor = *unrollJamFactor; 60 } 61 62 void runOnOperation() override; 63 }; 64 } // namespace 65 66 std::unique_ptr<OperationPass<func::FuncOp>> 67 mlir::createLoopUnrollAndJamPass(int unrollJamFactor) { 68 return std::make_unique<LoopUnrollAndJam>( 69 unrollJamFactor == -1 ? None : Optional<unsigned>(unrollJamFactor)); 70 } 71 72 void LoopUnrollAndJam::runOnOperation() { 73 if (getOperation().isExternal()) 74 return; 75 76 // Currently, just the outermost loop from the first loop nest is 77 // unroll-and-jammed by this pass. However, runOnAffineForOp can be called on 78 // any for operation. 79 auto &entryBlock = getOperation().front(); 80 if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front())) 81 (void)loopUnrollJamByFactor(forOp, unrollJamFactor); 82 } 83