1 //===- ParallelLoopTiling.cpp - Tiles scf.parallel ---------------===// 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 tiling on parallel loops. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "PassDetail.h" 14 #include "mlir/Dialect/Affine/IR/AffineOps.h" 15 #include "mlir/Dialect/SCF/Passes.h" 16 #include "mlir/Dialect/SCF/SCF.h" 17 #include "mlir/Dialect/SCF/Transforms.h" 18 #include "mlir/Dialect/SCF/Utils.h" 19 #include "mlir/Dialect/StandardOps/IR/Ops.h" 20 21 using namespace mlir; 22 using namespace mlir::scf; 23 24 /// Tile a parallel loop of the form 25 /// scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3) 26 /// step (%arg4, %arg5) 27 /// 28 /// into 29 /// scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3) 30 /// step (%arg4*tileSize[0], 31 /// %arg5*tileSize[1]) 32 /// scf.parallel (%j0, %j1) = (0, 0) to (min(%arg4*tileSize[0], %arg2-%i0) 33 /// min(%arg5*tileSize[1], %arg3-%i1)) 34 /// step (%arg4, %arg5) 35 /// 36 /// where the uses of %i0 and %i1 in the loop body are replaced by 37 /// %i0 + j0 and %i1 + %j1. 38 // 39 /// The old loop is replaced with the new one. 40 std::pair<ParallelOp, ParallelOp> 41 mlir::scf::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) { 42 OpBuilder b(op); 43 auto zero = b.create<ConstantIndexOp>(op.getLoc(), 0); 44 SmallVector<Value, 2> tileSizeConstants; 45 tileSizeConstants.reserve(op.upperBound().size()); 46 for (size_t i = 0, end = op.upperBound().size(); i != end; ++i) { 47 if (i < tileSizes.size()) 48 tileSizeConstants.push_back( 49 b.create<ConstantIndexOp>(op.getLoc(), tileSizes[i])); 50 else 51 // Just pick 1 for the remaining dimensions. 52 tileSizeConstants.push_back(b.create<ConstantIndexOp>(op.getLoc(), 1)); 53 } 54 55 // Create the outer loop with adjusted steps. 56 SmallVector<Value, 2> newSteps; 57 newSteps.reserve(op.step().size()); 58 for (auto step : llvm::zip(op.step(), tileSizeConstants)) { 59 newSteps.push_back( 60 b.create<MulIOp>(op.getLoc(), std::get<0>(step), std::get<1>(step))); 61 } 62 auto outerLoop = b.create<ParallelOp>(op.getLoc(), op.lowerBound(), 63 op.upperBound(), newSteps); 64 b.setInsertionPointToStart(outerLoop.getBody()); 65 66 // Compute min(size, dim - offset) to avoid out-of-bounds accesses. 67 // FIXME: Instead of using min, we want to replicate the tail. This would give 68 // the inner loop constant bounds for easy vectorization. 69 auto minMap = AffineMap::get( 70 /*dimCount=*/3, /*symbolCount=*/0, 71 {getAffineDimExpr(/*position=*/0, b.getContext()), 72 getAffineDimExpr(/*position=*/1, b.getContext()) - 73 getAffineDimExpr(/*position=*/2, b.getContext())}, 74 b.getContext()); 75 76 // Create the inner loop with adjusted bounds. 77 SmallVector<Value, 2> newBounds; 78 newBounds.reserve(op.upperBound().size()); 79 for (auto dim : llvm::zip(outerLoop.lowerBound(), outerLoop.upperBound(), 80 outerLoop.step(), outerLoop.getInductionVars(), 81 op.step(), tileSizeConstants)) { 82 Value lowerBound, upperBound, newStep, iv, step, tileSizeConstant; 83 std::tie(lowerBound, upperBound, newStep, iv, step, tileSizeConstant) = dim; 84 // Collect the statically known loop bounds 85 auto lowerBoundConstant = 86 dyn_cast_or_null<ConstantIndexOp>(lowerBound.getDefiningOp()); 87 auto upperBoundConstant = 88 dyn_cast_or_null<ConstantIndexOp>(upperBound.getDefiningOp()); 89 auto stepConstant = dyn_cast_or_null<ConstantIndexOp>(step.getDefiningOp()); 90 auto tileSize = 91 cast<ConstantIndexOp>(tileSizeConstant.getDefiningOp()).getValue(); 92 // If the loop bounds and the loop step are constant and if the number of 93 // loop iterations is an integer multiple of the tile size, we use a static 94 // bound for the inner loop. 95 if (lowerBoundConstant && upperBoundConstant && stepConstant) { 96 auto numIterations = llvm::divideCeil(upperBoundConstant.getValue() - 97 lowerBoundConstant.getValue(), 98 stepConstant.getValue()); 99 if (numIterations % tileSize == 0) { 100 newBounds.push_back(newStep); 101 continue; 102 } 103 } 104 // Otherwise, we dynamically compute the bound for 105 // each iteration of the outer loop. 106 newBounds.push_back( 107 b.create<AffineMinOp>(op.getLoc(), b.getIndexType(), minMap, 108 ValueRange{newStep, upperBound, iv})); 109 } 110 auto innerLoop = b.create<ParallelOp>( 111 op.getLoc(), SmallVector<Value, 2>(newBounds.size(), zero), newBounds, 112 op.step()); 113 114 // Steal the body of the old parallel loop and erase it. 115 innerLoop.region().takeBody(op.region()); 116 117 // Insert computation for new index vectors and replace uses. 118 b.setInsertionPointToStart(innerLoop.getBody()); 119 for (auto ivs : 120 llvm::zip(innerLoop.getInductionVars(), outerLoop.getInductionVars())) { 121 Value inner_index = std::get<0>(ivs); 122 AddIOp newIndex = 123 b.create<AddIOp>(op.getLoc(), std::get<0>(ivs), std::get<1>(ivs)); 124 inner_index.replaceAllUsesExcept(newIndex, newIndex); 125 } 126 127 op.erase(); 128 return std::make_pair(outerLoop, innerLoop); 129 } 130 131 namespace { 132 struct ParallelLoopTiling 133 : public SCFParallelLoopTilingBase<ParallelLoopTiling> { 134 ParallelLoopTiling() = default; 135 explicit ParallelLoopTiling(ArrayRef<int64_t> tileSizes) { 136 this->tileSizes = tileSizes; 137 } 138 139 void runOnFunction() override { 140 SmallVector<ParallelOp, 2> innermostPloops; 141 getInnermostParallelLoops(getFunction().getOperation(), innermostPloops); 142 for (ParallelOp ploop : innermostPloops) { 143 // FIXME: Add reduction support. 144 if (ploop.getNumReductions() == 0) 145 tileParallelLoop(ploop, tileSizes); 146 } 147 } 148 }; 149 } // namespace 150 151 std::unique_ptr<Pass> 152 mlir::createParallelLoopTilingPass(ArrayRef<int64_t> tileSizes) { 153 return std::make_unique<ParallelLoopTiling>(tileSizes); 154 } 155