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 #include "mlir/Transforms/RegionUtils.h"
20 #include "llvm/Support/CommandLine.h"
21 
22 using namespace mlir;
23 using namespace mlir::scf;
24 
25 /// Tile a parallel loop of the form
26 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
27 ///                                             step (%arg4, %arg5)
28 ///
29 /// into
30 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
31 ///                                             step (%arg4*tileSize[0],
32 ///                                                   %arg5*tileSize[1])
33 ///     scf.parallel (%j0, %j1) = (0, 0) to (min(tileSize[0], %arg2-%i0)
34 ///                                           min(tileSize[1], %arg3-%i1))
35 ///                                        step (%arg4, %arg5)
36 ///
37 /// where the uses of %i0 and %i1 in the loop body are replaced by
38 /// %i0 + j0 and %i1 + %j1.
39 //
40 /// The old loop is replaced with the new one.
41 void 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 bounds : llvm::zip(tileSizeConstants, outerLoop.upperBound(),
80                                outerLoop.getInductionVars())) {
81     newBounds.push_back(b.create<AffineMinOp>(
82         op.getLoc(), b.getIndexType(), minMap,
83         ValueRange{std::get<0>(bounds), std::get<1>(bounds),
84                    std::get<2>(bounds)}));
85   }
86   auto innerLoop = b.create<ParallelOp>(
87       op.getLoc(), SmallVector<Value, 2>(newBounds.size(), zero), newBounds,
88       op.step());
89 
90   // Steal the body of the old parallel loop and erase it.
91   innerLoop.region().takeBody(op.region());
92 
93   // Insert computation for new index vectors and replace uses.
94   b.setInsertionPointToStart(innerLoop.getBody());
95   for (auto ivs :
96        llvm::zip(innerLoop.getInductionVars(), outerLoop.getInductionVars())) {
97     Value inner_index = std::get<0>(ivs);
98     AddIOp newIndex =
99         b.create<AddIOp>(op.getLoc(), std::get<0>(ivs), std::get<1>(ivs));
100     inner_index.replaceAllUsesExcept(
101         newIndex, SmallPtrSet<Operation *, 1>{newIndex.getOperation()});
102   }
103 
104   op.erase();
105 }
106 
107 /// Get a list of most nested parallel loops. Assumes that ParallelOps are only
108 /// directly nested.
109 static bool getInnermostNestedLoops(Block *block,
110                                     SmallVectorImpl<ParallelOp> &loops) {
111   bool hasInnerLoop = false;
112   for (auto parallelOp : block->getOps<ParallelOp>()) {
113     hasInnerLoop = true;
114     if (!getInnermostNestedLoops(parallelOp.getBody(), loops))
115       loops.push_back(parallelOp);
116   }
117   return hasInnerLoop;
118 }
119 
120 namespace {
121 struct ParallelLoopTiling
122     : public SCFParallelLoopTilingBase<ParallelLoopTiling> {
123   ParallelLoopTiling() = default;
124   explicit ParallelLoopTiling(ArrayRef<int64_t> tileSizes) {
125     this->tileSizes = tileSizes;
126   }
127 
128   void runOnFunction() override {
129     SmallVector<ParallelOp, 2> mostNestedParallelOps;
130     for (Block &block : getFunction()) {
131       getInnermostNestedLoops(&block, mostNestedParallelOps);
132     }
133     for (ParallelOp pLoop : mostNestedParallelOps) {
134       tileParallelLoop(pLoop, tileSizes);
135     }
136   }
137 };
138 } // namespace
139 
140 std::unique_ptr<Pass>
141 mlir::createParallelLoopTilingPass(ArrayRef<int64_t> tileSizes) {
142   return std::make_unique<ParallelLoopTiling>(tileSizes);
143 }
144