74 lines · cpp
1//===- ForallToFor.cpp - scf.forall to scf.for loop conversion ------------===//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.ForallOp's into SCF.ForOp's.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/SCF/Transforms/Passes.h"14 15#include "mlir/Dialect/SCF/IR/SCF.h"16#include "mlir/Dialect/SCF/Transforms/Transforms.h"17#include "mlir/IR/PatternMatch.h"18 19namespace mlir {20#define GEN_PASS_DEF_SCFFORALLTOFORLOOP21#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"22} // namespace mlir23 24using namespace mlir;25using scf::LoopNest;26 27LogicalResult28mlir::scf::forallToForLoop(RewriterBase &rewriter, scf::ForallOp forallOp,29 SmallVectorImpl<Operation *> *results) {30 OpBuilder::InsertionGuard guard(rewriter);31 rewriter.setInsertionPoint(forallOp);32 33 Location loc = forallOp.getLoc();34 SmallVector<Value> lbs = forallOp.getLowerBound(rewriter);35 SmallVector<Value> ubs = forallOp.getUpperBound(rewriter);36 SmallVector<Value> steps = forallOp.getStep(rewriter);37 LoopNest loopNest = scf::buildLoopNest(rewriter, loc, lbs, ubs, steps);38 39 SmallVector<Value> ivs = llvm::map_to_vector(40 loopNest.loops, [](scf::ForOp loop) { return loop.getInductionVar(); });41 42 Block *innermostBlock = loopNest.loops.back().getBody();43 rewriter.eraseOp(forallOp.getBody()->getTerminator());44 rewriter.inlineBlockBefore(forallOp.getBody(), innermostBlock,45 innermostBlock->getTerminator()->getIterator(),46 ivs);47 rewriter.eraseOp(forallOp);48 49 if (results) {50 llvm::move(loopNest.loops, std::back_inserter(*results));51 }52 53 return success();54}55 56namespace {57struct ForallToForLoop : public impl::SCFForallToForLoopBase<ForallToForLoop> {58 void runOnOperation() override {59 Operation *parentOp = getOperation();60 IRRewriter rewriter(parentOp->getContext());61 62 parentOp->walk([&](scf::ForallOp forallOp) {63 if (failed(scf::forallToForLoop(rewriter, forallOp))) {64 return signalPassFailure();65 }66 });67 }68};69} // namespace70 71std::unique_ptr<Pass> mlir::createForallToForLoopPass() {72 return std::make_unique<ForallToForLoop>();73}74