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