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