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/Analysis/LoopAnalysis.h" 38 #include "mlir/Dialect/Affine/IR/AffineOps.h" 39 #include "mlir/Dialect/Affine/Passes.h" 40 #include "mlir/IR/AffineExpr.h" 41 #include "mlir/IR/AffineMap.h" 42 #include "mlir/IR/BlockAndValueMapping.h" 43 #include "mlir/IR/Builders.h" 44 #include "mlir/Transforms/LoopUtils.h" 45 #include "llvm/ADT/DenseMap.h" 46 #include "llvm/Support/CommandLine.h" 47 48 using namespace mlir; 49 50 #define DEBUG_TYPE "affine-loop-unroll-jam" 51 52 namespace { 53 /// Loop unroll jam pass. Currently, this just unroll jams the first 54 /// outer loop in a Function. 55 struct LoopUnrollAndJam : public AffineLoopUnrollAndJamBase<LoopUnrollAndJam> { 56 explicit LoopUnrollAndJam(Optional<unsigned> unrollJamFactor = None) { 57 if (unrollJamFactor) 58 this->unrollJamFactor = *unrollJamFactor; 59 } 60 61 void runOnFunction() override; 62 }; 63 } // end anonymous namespace 64 65 std::unique_ptr<OperationPass<FuncOp>> 66 mlir::createLoopUnrollAndJamPass(int unrollJamFactor) { 67 return std::make_unique<LoopUnrollAndJam>( 68 unrollJamFactor == -1 ? None : Optional<unsigned>(unrollJamFactor)); 69 } 70 71 void LoopUnrollAndJam::runOnFunction() { 72 // Currently, just the outermost loop from the first loop nest is 73 // unroll-and-jammed by this pass. However, runOnAffineForOp can be called on 74 // any for operation. 75 auto &entryBlock = getFunction().front(); 76 if (auto forOp = dyn_cast<AffineForOp>(entryBlock.front())) 77 loopUnrollJamByFactor(forOp, unrollJamFactor); 78 } 79