93 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a parallelizer for affine loop nests that is able to10// 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/LoopUtils.h"22#include "mlir/Dialect/Affine/Passes.h.inc"23#include "mlir/Dialect/Affine/Utils.h"24#include "mlir/Dialect/Func/IR/FuncOps.h"25#include "llvm/Support/Debug.h"26 27namespace mlir {28namespace affine {29#define GEN_PASS_DEF_AFFINEPARALLELIZE30#include "mlir/Dialect/Affine/Passes.h.inc"31} // namespace affine32} // namespace mlir33 34#define DEBUG_TYPE "affine-parallel"35 36using namespace mlir;37using namespace mlir::affine;38 39namespace {40/// Convert all parallel affine.for op into 1-D affine.parallel op.41struct AffineParallelize42 : public affine::impl::AffineParallelizeBase<AffineParallelize> {43 using AffineParallelizeBase<AffineParallelize>::AffineParallelizeBase;44 45 void runOnOperation() override;46};47 48/// Descriptor of a potentially parallelizable loop.49struct ParallelizationCandidate {50 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} // namespace59 60void AffineParallelize::runOnOperation() {61 func::FuncOp f = getOperation();62 63 // The walker proceeds in pre-order to process the outer loops first64 // 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