1 | //===- AffineParallelize.cpp - Affineparallelize Pass---------------------===// |
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 a parallelizer for affine loop nests that is able to |
10 | // perform inner or outer loop parallelization. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #include "mlir/Dialect/Affine/Passes.h" |
15 | |
16 | #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h" |
17 | #include "mlir/Dialect/Affine/Analysis/AffineStructures.h" |
18 | #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h" |
19 | #include "mlir/Dialect/Affine/Analysis/Utils.h" |
20 | #include "mlir/Dialect/Affine/IR/AffineOps.h" |
21 | #include "mlir/Dialect/Affine/IR/AffineValueMap.h" |
22 | #include "mlir/Dialect/Affine/LoopUtils.h" |
23 | #include "mlir/Dialect/Affine/Passes.h.inc" |
24 | #include "mlir/Dialect/Affine/Utils.h" |
25 | #include "mlir/Dialect/Func/IR/FuncOps.h" |
26 | #include "llvm/Support/Debug.h" |
27 | #include <deque> |
28 | |
29 | namespace mlir { |
30 | namespace affine { |
31 | #define GEN_PASS_DEF_AFFINEPARALLELIZE |
32 | #include "mlir/Dialect/Affine/Passes.h.inc" |
33 | } // namespace affine |
34 | } // namespace mlir |
35 | |
36 | #define DEBUG_TYPE "affine-parallel" |
37 | |
38 | using namespace mlir; |
39 | using namespace mlir::affine; |
40 | |
41 | namespace { |
42 | /// Convert all parallel affine.for op into 1-D affine.parallel op. |
43 | struct AffineParallelize |
44 | : public affine::impl::AffineParallelizeBase<AffineParallelize> { |
45 | using AffineParallelizeBase<AffineParallelize>::AffineParallelizeBase; |
46 | |
47 | void runOnOperation() override; |
48 | }; |
49 | |
50 | /// Descriptor of a potentially parallelizable loop. |
51 | struct ParallelizationCandidate { |
52 | ParallelizationCandidate(AffineForOp l, SmallVector<LoopReduction> &&r) |
53 | : loop(l), reductions(std::move(r)) {} |
54 | |
55 | /// The potentially parallelizable loop. |
56 | AffineForOp loop; |
57 | /// Desciprtors of reductions that can be parallelized in the loop. |
58 | SmallVector<LoopReduction> reductions; |
59 | }; |
60 | } // namespace |
61 | |
62 | void AffineParallelize::runOnOperation() { |
63 | func::FuncOp f = getOperation(); |
64 | |
65 | // The walker proceeds in pre-order to process the outer loops first |
66 | // and control the number of outer parallel loops. |
67 | std::vector<ParallelizationCandidate> parallelizableLoops; |
68 | f.walk<WalkOrder::PreOrder>([&](AffineForOp loop) { |
69 | SmallVector<LoopReduction> reductions; |
70 | if (isLoopParallel(loop, parallelReductions ? &reductions : nullptr)) |
71 | parallelizableLoops.emplace_back(loop, std::move(reductions)); |
72 | }); |
73 | |
74 | for (const ParallelizationCandidate &candidate : parallelizableLoops) { |
75 | unsigned numParentParallelOps = 0; |
76 | AffineForOp loop = candidate.loop; |
77 | for (Operation *op = loop->getParentOp(); |
78 | op != nullptr && !op->hasTrait<OpTrait::AffineScope>(); |
79 | op = op->getParentOp()) { |
80 | if (isa<AffineParallelOp>(Val: op)) |
81 | ++numParentParallelOps; |
82 | } |
83 | |
84 | if (numParentParallelOps < maxNested) { |
85 | if (failed(affineParallelize(loop, candidate.reductions))) { |
86 | LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] failed to parallelize\n" |
87 | << loop); |
88 | } |
89 | } else { |
90 | LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] too many nested loops\n" |
91 | << loop); |
92 | } |
93 | } |
94 | } |
95 | |