xref: /llvm-project/mlir/lib/Dialect/Affine/Transforms/AffineParallelize.cpp (revision 4c48f016effde67d500fc95290096aec9f3bdb70)
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   void runOnOperation() override;
46 };
47 
48 /// Descriptor of a potentially parallelizable loop.
49 struct ParallelizationCandidate {
ParallelizationCandidate__anonaf1e49550111::ParallelizationCandidate50   ParallelizationCandidate(AffineForOp l, SmallVector<LoopReduction> &&r)
51       : loop(l), reductions(std::move(r)) {}
52 
53   /// The potentially parallelizable loop.
54   AffineForOp loop;
55   /// Desciprtors of reductions that can be parallelized in the loop.
56   SmallVector<LoopReduction> reductions;
57 };
58 } // namespace
59 
runOnOperation()60 void AffineParallelize::runOnOperation() {
61   func::FuncOp f = getOperation();
62 
63   // The walker proceeds in pre-order to process the outer loops first
64   // and control the number of outer parallel loops.
65   std::vector<ParallelizationCandidate> parallelizableLoops;
66   f.walk<WalkOrder::PreOrder>([&](AffineForOp loop) {
67     SmallVector<LoopReduction> reductions;
68     if (isLoopParallel(loop, parallelReductions ? &reductions : nullptr))
69       parallelizableLoops.emplace_back(loop, std::move(reductions));
70   });
71 
72   for (const ParallelizationCandidate &candidate : parallelizableLoops) {
73     unsigned numParentParallelOps = 0;
74     AffineForOp loop = candidate.loop;
75     for (Operation *op = loop->getParentOp();
76          op != nullptr && !op->hasTrait<OpTrait::AffineScope>();
77          op = op->getParentOp()) {
78       if (isa<AffineParallelOp>(op))
79         ++numParentParallelOps;
80     }
81 
82     if (numParentParallelOps < maxNested) {
83       if (failed(affineParallelize(loop, candidate.reductions))) {
84         LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] failed to parallelize\n"
85                                 << loop);
86       }
87     } else {
88       LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] too many nested loops\n"
89                               << loop);
90     }
91   }
92 }
93 
94 std::unique_ptr<OperationPass<func::FuncOp>>
createAffineParallelizePass()95 mlir::affine::createAffineParallelizePass() {
96   return std::make_unique<AffineParallelize>();
97 }
98