210 lines · cpp
1//===-- SimdOnly.cpp ------------------------------------------------------===//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#include "flang/Optimizer/Builder/FIRBuilder.h"10#include "mlir/Dialect/Arith/IR/Arith.h"11#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"12#include "mlir/Dialect/Func/IR/FuncOps.h"13#include "mlir/Dialect/OpenMP/OpenMPDialect.h"14#include "mlir/IR/MLIRContext.h"15#include "mlir/IR/Operation.h"16#include "mlir/IR/PatternMatch.h"17#include "mlir/Pass/Pass.h"18#include "mlir/Support/LLVM.h"19#include "mlir/Transforms/GreedyPatternRewriteDriver.h"20#include "llvm/Support/Debug.h"21 22namespace flangomp {23#define GEN_PASS_DEF_SIMDONLYPASS24#include "flang/Optimizer/OpenMP/Passes.h.inc"25} // namespace flangomp26 27namespace {28 29#define DEBUG_TYPE "omp-simd-only-pass"30 31/// Rewrite and remove OpenMP operations left after the parse tree rewriting for32/// -fopenmp-simd is done. If possible, OpenMP constructs should be rewritten at33/// the parse tree stage. This pass is supposed to only handle complexities34/// around untangling composite simd constructs, and perform the necessary35/// cleanup.36class SimdOnlyConversionPattern : public mlir::RewritePattern {37public:38 SimdOnlyConversionPattern(mlir::MLIRContext *ctx)39 : mlir::RewritePattern(MatchAnyOpTypeTag{}, 1, ctx) {}40 41 mlir::LogicalResult42 matchAndRewrite(mlir::Operation *op,43 mlir::PatternRewriter &rewriter) const override {44 if (op->getDialect()->getNamespace() !=45 mlir::omp::OpenMPDialect::getDialectNamespace())46 return rewriter.notifyMatchFailure(op, "Not an OpenMP op");47 48 if (auto simdOp = mlir::dyn_cast<mlir::omp::SimdOp>(op)) {49 // Remove the composite attr given that the op will no longer be composite50 if (simdOp.isComposite()) {51 simdOp.setComposite(false);52 return mlir::success();53 }54 55 return rewriter.notifyMatchFailure(op, "Op is a plain SimdOp");56 }57 58 if (op->getParentOfType<mlir::omp::SimdOp>() &&59 (mlir::isa<mlir::omp::YieldOp>(op) ||60 mlir::isa<mlir::omp::ScanOp>(op) ||61 mlir::isa<mlir::omp::LoopNestOp>(op) ||62 mlir::isa<mlir::omp::TerminatorOp>(op)))63 return rewriter.notifyMatchFailure(op, "Op is part of a simd construct");64 65 if (!mlir::isa<mlir::func::FuncOp>(op->getParentOp()) &&66 (mlir::isa<mlir::omp::TerminatorOp>(op) ||67 mlir::isa<mlir::omp::YieldOp>(op)))68 return rewriter.notifyMatchFailure(op,69 "Non top-level yield or terminator");70 71 LLVM_DEBUG(llvm::dbgs() << "SimdOnlyPass matched OpenMP op:\n");72 LLVM_DEBUG(op->dump());73 74 auto eraseUnlessUsedBySimd = [&](mlir::Operation *ompOp,75 mlir::StringAttr name) {76 if (auto uses =77 mlir::SymbolTable::getSymbolUses(name, op->getParentOp())) {78 for (auto &use : *uses)79 if (mlir::isa<mlir::omp::SimdOp>(use.getUser()))80 return rewriter.notifyMatchFailure(op,81 "Op used by a simd construct");82 }83 rewriter.eraseOp(ompOp);84 return mlir::success();85 };86 87 if (auto ompOp = mlir::dyn_cast<mlir::omp::PrivateClauseOp>(op))88 return eraseUnlessUsedBySimd(ompOp, ompOp.getSymNameAttr());89 if (auto ompOp = mlir::dyn_cast<mlir::omp::DeclareReductionOp>(op))90 return eraseUnlessUsedBySimd(ompOp, ompOp.getSymNameAttr());91 92 // Might be left over from rewriting composite simd with target map93 if (mlir::isa<mlir::omp::MapBoundsOp>(op)) {94 rewriter.eraseOp(op);95 return mlir::success();96 }97 if (auto mapInfoOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(op)) {98 rewriter.replaceOp(mapInfoOp, {mapInfoOp.getVarPtr()});99 return mlir::success();100 }101 102 // Might be leftover after parse tree rewriting103 if (auto threadPrivateOp = mlir::dyn_cast<mlir::omp::ThreadprivateOp>(op)) {104 rewriter.replaceOp(threadPrivateOp, {threadPrivateOp.getSymAddr()});105 return mlir::success();106 }107 108 fir::FirOpBuilder builder(rewriter, op);109 mlir::Location loc = op->getLoc();110 111 auto inlineSimpleOp = [&](mlir::Operation *ompOp) -> bool {112 if (!ompOp)113 return false;114 115 assert("OpenMP operation has one region" && ompOp->getNumRegions() == 1);116 117 llvm::SmallVector<std::pair<mlir::Value, mlir::BlockArgument>>118 blockArgsPairs;119 if (auto iface =120 mlir::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(op)) {121 iface.getBlockArgsPairs(blockArgsPairs);122 for (auto [value, argument] : blockArgsPairs)123 rewriter.replaceAllUsesWith(argument, value);124 }125 126 if (ompOp->getRegion(0).getBlocks().size() == 1) {127 auto &block = *ompOp->getRegion(0).getBlocks().begin();128 // This block is about to be removed so any arguments should have been129 // replaced by now.130 block.eraseArguments(0, block.getNumArguments());131 if (auto terminatorOp =132 mlir::dyn_cast<mlir::omp::TerminatorOp>(block.back())) {133 rewriter.eraseOp(terminatorOp);134 }135 rewriter.inlineBlockBefore(&block, ompOp, {});136 } else {137 // When dealing with multi-block regions we need to fix up the control138 // flow139 auto *origBlock = ompOp->getBlock();140 auto *newBlock = rewriter.splitBlock(origBlock, ompOp->getIterator());141 auto *innerFrontBlock = &ompOp->getRegion(0).getBlocks().front();142 builder.setInsertionPointToEnd(origBlock);143 mlir::cf::BranchOp::create(builder, loc, innerFrontBlock);144 // We are no longer passing any arguments to the first block in the145 // region, so this should be safe to erase.146 innerFrontBlock->eraseArguments(0, innerFrontBlock->getNumArguments());147 148 for (auto &innerBlock : ompOp->getRegion(0).getBlocks()) {149 // Remove now-unused block arguments150 for (auto arg : innerBlock.getArguments()) {151 if (arg.getUses().empty())152 innerBlock.eraseArgument(arg.getArgNumber());153 }154 if (auto terminatorOp =155 mlir::dyn_cast<mlir::omp::TerminatorOp>(innerBlock.back())) {156 builder.setInsertionPointToEnd(&innerBlock);157 mlir::cf::BranchOp::create(builder, loc, newBlock);158 rewriter.eraseOp(terminatorOp);159 }160 }161 162 rewriter.inlineRegionBefore(ompOp->getRegion(0), newBlock);163 }164 165 rewriter.eraseOp(op);166 return true;167 };168 169 // Remove ops that will be surrounding simd once a composite simd construct170 // goes through the codegen stage. All of the other ones should have alredy171 // been removed in the parse tree rewriting stage.172 if (inlineSimpleOp(mlir::dyn_cast<mlir::omp::TeamsOp>(op)) ||173 inlineSimpleOp(mlir::dyn_cast<mlir::omp::ParallelOp>(op)) ||174 inlineSimpleOp(mlir::dyn_cast<mlir::omp::TargetOp>(op)) ||175 inlineSimpleOp(mlir::dyn_cast<mlir::omp::WsloopOp>(op)) ||176 inlineSimpleOp(mlir::dyn_cast<mlir::omp::DistributeOp>(op)))177 return mlir::success();178 179 op->emitOpError("left unhandled after SimdOnly pass.");180 return mlir::failure();181 }182};183 184class SimdOnlyPass : public flangomp::impl::SimdOnlyPassBase<SimdOnlyPass> {185 186public:187 SimdOnlyPass() = default;188 189 void runOnOperation() override {190 mlir::ModuleOp module = getOperation();191 192 mlir::MLIRContext *context = &getContext();193 mlir::RewritePatternSet patterns(context);194 patterns.insert<SimdOnlyConversionPattern>(context);195 196 mlir::GreedyRewriteConfig config;197 // Prevent the pattern driver from merging blocks.198 config.setRegionSimplificationLevel(199 mlir::GreedySimplifyRegionLevel::Disabled);200 201 if (mlir::failed(202 mlir::applyPatternsGreedily(module, std::move(patterns), config))) {203 mlir::emitError(module.getLoc(), "Error in SimdOnly conversion pass");204 signalPassFailure();205 }206 }207};208 209} // namespace210