86 lines · cpp
1//===- ParallelForToNestedFors.cpp - scf.parallel to nested scf.for ops --===//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// Transforms SCF.ParallelOp to nested scf.for ops.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/SCF/IR/SCF.h"14#include "mlir/Dialect/SCF/Transforms/Passes.h"15#include "mlir/Dialect/SCF/Transforms/Transforms.h"16#include "mlir/IR/PatternMatch.h"17#include "llvm/Support/Debug.h"18 19namespace mlir {20#define GEN_PASS_DEF_SCFPARALLELFORTONESTEDFORS21#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"22} // namespace mlir23 24#define DEBUG_TYPE "parallel-for-to-nested-fors"25using namespace mlir;26 27FailureOr<scf::LoopNest>28mlir::scf::parallelForToNestedFors(RewriterBase &rewriter,29 scf::ParallelOp parallelOp) {30 31 if (!parallelOp.getResults().empty())32 return rewriter.notifyMatchFailure(33 parallelOp, "Currently scf.parallel to scf.for conversion doesn't "34 "support scf.parallel with results.");35 36 rewriter.setInsertionPoint(parallelOp);37 38 Location loc = parallelOp.getLoc();39 SmallVector<Value> lowerBounds = parallelOp.getLowerBound();40 SmallVector<Value> upperBounds = parallelOp.getUpperBound();41 SmallVector<Value> steps = parallelOp.getStep();42 43 assert(lowerBounds.size() == upperBounds.size() &&44 lowerBounds.size() == steps.size() &&45 "Mismatched parallel loop bounds");46 47 scf::LoopNest loopNest =48 scf::buildLoopNest(rewriter, loc, lowerBounds, upperBounds, steps);49 50 SmallVector<Value> newInductionVars = llvm::map_to_vector(51 loopNest.loops, [](scf::ForOp forOp) { return forOp.getInductionVar(); });52 Block *linearizedBody = loopNest.loops.back().getBody();53 Block *parallelBody = parallelOp.getBody();54 rewriter.eraseOp(parallelBody->getTerminator());55 rewriter.inlineBlockBefore(parallelBody, linearizedBody->getTerminator(),56 newInductionVars);57 rewriter.eraseOp(parallelOp);58 return loopNest;59}60 61namespace {62struct ParallelForToNestedFors final63 : public impl::SCFParallelForToNestedForsBase<ParallelForToNestedFors> {64 void runOnOperation() override {65 Operation *parentOp = getOperation();66 IRRewriter rewriter(parentOp->getContext());67 68 parentOp->walk(69 [&](scf::ParallelOp parallelOp) {70 if (failed(scf::parallelForToNestedFors(rewriter, parallelOp))) {71 LLVM_DEBUG(72 llvm::dbgs()73 << "Failed to convert scf.parallel to nested scf.for ops for:\n"74 << parallelOp << "\n");75 return WalkResult::advance();76 }77 return WalkResult::advance();78 });79 }80};81} // namespace82 83std::unique_ptr<Pass> mlir::createParallelForToNestedForsPass() {84 return std::make_unique<ParallelForToNestedFors>();85}86