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 "PassDetail.h"
15 #include "mlir/Analysis/AffineAnalysis.h"
16 #include "mlir/Analysis/AffineStructures.h"
17 #include "mlir/Analysis/LoopAnalysis.h"
18 #include "mlir/Analysis/Utils.h"
19 #include "mlir/Dialect/Affine/IR/AffineOps.h"
20 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
21 #include "mlir/Dialect/Affine/Passes.h"
22 #include "mlir/Dialect/Affine/Passes.h.inc"
23 #include "mlir/Dialect/Affine/Utils.h"
24 #include "mlir/Transforms/LoopUtils.h"
25 #include "llvm/Support/Debug.h"
26 #include <deque>
27 
28 #define DEBUG_TYPE "affine-parallel"
29 
30 using namespace mlir;
31 
32 namespace {
33 /// Convert all parallel affine.for op into 1-D affine.parallel op.
34 struct AffineParallelize : public AffineParallelizeBase<AffineParallelize> {
35   void runOnFunction() override;
36 };
37 
38 /// Descriptor of a potentially parallelizable loop.
39 struct ParallelizationCandidate {
40   ParallelizationCandidate(AffineForOp l, SmallVector<LoopReduction> &&r)
41       : loop(l), reductions(std::move(r)) {}
42 
43   /// The potentially parallelizable loop.
44   AffineForOp loop;
45   /// Desciprtors of reductions that can be parallelized in the loop.
46   SmallVector<LoopReduction> reductions;
47 };
48 } // namespace
49 
50 void AffineParallelize::runOnFunction() {
51   FuncOp f = getFunction();
52 
53   // The walker proceeds in post-order, but we need to process outer loops first
54   // to control the number of outer parallel loops, so push candidate loops to
55   // the front of a deque.
56   std::deque<ParallelizationCandidate> parallelizableLoops;
57   f.walk([&](AffineForOp loop) {
58     SmallVector<LoopReduction> reductions;
59     if (isLoopParallel(loop, parallelReductions ? &reductions : nullptr))
60       parallelizableLoops.emplace_back(loop, std::move(reductions));
61   });
62 
63   for (const ParallelizationCandidate &candidate : parallelizableLoops) {
64     unsigned numParentParallelOps = 0;
65     AffineForOp loop = candidate.loop;
66     for (Operation *op = loop->getParentOp();
67          op != nullptr && !op->hasTrait<OpTrait::AffineScope>();
68          op = op->getParentOp()) {
69       if (isa<AffineParallelOp>(op))
70         ++numParentParallelOps;
71     }
72 
73     if (numParentParallelOps < maxNested) {
74       if (failed(affineParallelize(loop, candidate.reductions))) {
75         LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] failed to parallelize\n"
76                                 << loop);
77       }
78     } else {
79       LLVM_DEBUG(llvm::dbgs() << "[" DEBUG_TYPE "] too many nested loops\n"
80                               << loop);
81     }
82   }
83 }
84 
85 std::unique_ptr<OperationPass<FuncOp>> mlir::createAffineParallelizePass() {
86   return std::make_unique<AffineParallelize>();
87 }
88