================
@@ -0,0 +1,360 @@
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Transforms/Utils.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include <llvm/Support/Debug.h>
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/IR/Operation.h>
+#include <mlir/IR/PatternMatch.h>
+#include <mlir/Support/LLVM.h>
+
+namespace flangomp {
+#define GEN_PASS_DEF_SIMDONLYPASS
+#include "flang/Optimizer/OpenMP/Passes.h.inc"
+} // namespace flangomp
+
+namespace {
+
+#define DEBUG_TYPE "omp-simd-only-pass"
+
+class SimdOnlyConversionPattern : public mlir::RewritePattern {
+public:
+ SimdOnlyConversionPattern(mlir::MLIRContext *ctx)
+ : mlir::RewritePattern(MatchAnyOpTypeTag{}, 1, ctx) {}
+
+ mlir::LogicalResult
+ matchAndRewrite(mlir::Operation *op,
+ mlir::PatternRewriter &rewriter) const override {
+ if (op->getDialect()->getNamespace() !=
+ mlir::omp::OpenMPDialect::getDialectNamespace())
+ return rewriter.notifyMatchFailure(op, "Not an OpenMP op");
+
+ if (auto simdOp = mlir::dyn_cast<mlir::omp::SimdOp>(op)) {
+ // Remove the composite attr given that the op will no longer be
composite
+ if (simdOp.isComposite()) {
+ simdOp.setComposite(false);
+ return mlir::success();
+ }
+
+ return rewriter.notifyMatchFailure(op, "Op is a plain SimdOp");
+ }
+
+ if (op->getParentOfType<mlir::omp::SimdOp>())
+ return rewriter.notifyMatchFailure(op, "Op is nested under a SimdOp");
+
+ if (!mlir::isa<mlir::func::FuncOp>(op->getParentOp()) &&
+ (mlir::isa<mlir::omp::TerminatorOp>(op) ||
+ mlir::isa<mlir::omp::YieldOp>(op)))
+ return rewriter.notifyMatchFailure(op,
+ "Non top-level yield or terminator");
+
+ // SectionOp overrides its BlockArgInterface based on the parent
SectionsOp.
+ // We need to make sure we only rewrite omp.sections once all omp.section
+ // ops inside it have been rewritten, otherwise the individual omp.section
+ // ops will not be able to access their argument values.
+ if (auto sectionsOp = mlir::dyn_cast<mlir::omp::SectionsOp>(op)) {
+ for (auto &opInSections : sectionsOp.getRegion().getOps())
+ if (mlir::isa<mlir::omp::SectionOp>(opInSections))
+ return rewriter.notifyMatchFailure(
+ op, "SectionsOp still contains individual sections");
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "SimdOnlyPass matched OpenMP op:\n");
+ LLVM_DEBUG(op->dump());
+
+ // Erase ops that don't need any special handling
+ if (mlir::isa<mlir::omp::BarrierOp>(op) ||
+ mlir::isa<mlir::omp::FlushOp>(op) ||
+ mlir::isa<mlir::omp::TaskyieldOp>(op) ||
+ mlir::isa<mlir::omp::MapBoundsOp>(op) ||
+ mlir::isa<mlir::omp::TargetEnterDataOp>(op) ||
+ mlir::isa<mlir::omp::TargetExitDataOp>(op) ||
+ mlir::isa<mlir::omp::TargetUpdateOp>(op) ||
+ mlir::isa<mlir::omp::OrderedOp>(op) ||
+ mlir::isa<mlir::omp::CancelOp>(op) ||
+ mlir::isa<mlir::omp::CancellationPointOp>(op) ||
+ mlir::isa<mlir::omp::ScanOp>(op) ||
+ mlir::isa<mlir::omp::TaskwaitOp>(op)) {
+ rewriter.eraseOp(op);
+ return mlir::success();
+ }
+
+ fir::FirOpBuilder builder(rewriter, op);
+ mlir::Location loc = op->getLoc();
+
+ auto inlineSimpleOp = [&](mlir::Operation *ompOp) -> bool {
+ if (!ompOp)
+ return false;
+
+ llvm::SmallVector<std::pair<mlir::Value, mlir::BlockArgument>>
+ blockArgsPairs;
+ if (auto iface =
+ mlir::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(op)) {
+ iface.getBlockArgsPairs(blockArgsPairs);
+ for (auto [value, argument] : blockArgsPairs)
+ rewriter.replaceAllUsesWith(argument, value);
+ }
+
+ if (ompOp->getRegion(0).getBlocks().size() == 1) {
+ auto &block = *ompOp->getRegion(0).getBlocks().begin();
+ // This block is about to be removed so any arguments should have been
+ // replaced by now.
+ block.eraseArguments(0, block.getNumArguments());
+ if (auto terminatorOp =
+ mlir::dyn_cast<mlir::omp::TerminatorOp>(block.back())) {
+ rewriter.eraseOp(terminatorOp);
+ }
+ rewriter.inlineBlockBefore(&block, op, {});
+ } else {
+ // When dealing with multi-block regions we need to fix up the control
+ // flow
+ auto *origBlock = ompOp->getBlock();
+ auto *newBlock = rewriter.splitBlock(origBlock, ompOp->getIterator());
+ auto *innerFrontBlock = &ompOp->getRegion(0).getBlocks().front();
+ builder.setInsertionPointToEnd(origBlock);
+ builder.create<mlir::cf::BranchOp>(loc, innerFrontBlock);
+ // We are no longer passing any arguments to the first block in the
+ // region, so this should be safe to erase.
+ innerFrontBlock->eraseArguments(0, innerFrontBlock->getNumArguments());
+
+ for (auto &innerBlock : ompOp->getRegion(0).getBlocks()) {
+ // Remove now-unused block arguments
+ for (auto arg : innerBlock.getArguments()) {
+ if (arg.getUses().empty())
+ innerBlock.eraseArgument(arg.getArgNumber());
+ }
+ if (auto terminatorOp =
+ mlir::dyn_cast<mlir::omp::TerminatorOp>(innerBlock.back())) {
+ builder.setInsertionPointToEnd(&innerBlock);
+ builder.create<mlir::cf::BranchOp>(loc, newBlock);
+ rewriter.eraseOp(terminatorOp);
+ }
+ }
+
+ rewriter.inlineRegionBefore(ompOp->getRegion(0), newBlock);
+ }
+
+ rewriter.eraseOp(op);
+ return true;
+ };
+
+ if (auto ompOp = mlir::dyn_cast<mlir::omp::LoopNestOp>(op)) {
+ mlir::Type indexType = builder.getIndexType();
+ mlir::Type oldIndexType = ompOp.getIVs().begin()->getType();
+ builder.setInsertionPoint(op);
+ auto one = builder.create<mlir::arith::ConstantIndexOp>(loc, 1);
+
+ // Generate the new loop nest
+ mlir::Block *nestBody = nullptr;
+ fir::DoLoopOp outerLoop = nullptr;
+ llvm::SmallVector<mlir::Value> loopIndArgs;
+ for (auto extent : ompOp.getLoopUpperBounds()) {
+ auto ub = builder.createConvert(loc, indexType, extent);
+ auto doLoop = builder.create<fir::DoLoopOp>(loc, one, ub, one, false);
+ nestBody = doLoop.getBody();
+ builder.setInsertionPointToStart(nestBody);
+ // Convert the indices to the type used inside the loop if needed
+ if (oldIndexType != indexType) {
+ auto convertedIndVar = builder.createConvert(
+ loc, oldIndexType, doLoop.getInductionVar());
+ loopIndArgs.push_back(convertedIndVar);
+ } else {
+ loopIndArgs.push_back(doLoop.getInductionVar());
+ }
+ if (!outerLoop)
+ outerLoop = doLoop;
+ }
+
+ // Move the omp loop body into the new loop body
+ if (ompOp->getRegion(0).getBlocks().size() == 1) {
+ auto &block = *ompOp->getRegion(0).getBlocks().begin();
+ rewriter.mergeBlocks(&block, nestBody, loopIndArgs);
+
+ // Find the new loop block terminator and move it before the end of the
+ // block
+ for (auto &loopBodyOp : nestBody->getOperations()) {
+ if (auto resultOp = mlir::dyn_cast<fir::ResultOp>(loopBodyOp)) {
+ rewriter.moveOpBefore(resultOp.getOperation(), &nestBody->back());
+ break;
+ }
+ }
+
+ // Remove omp.yield at the end of the loop body
+ if (auto yieldOp =
mlir::dyn_cast<mlir::omp::YieldOp>(nestBody->back()))
+ rewriter.eraseOp(yieldOp);
+ // DoLoopOp does not support multi-block regions, thus if we're dealing
+ // with multiple blocks we need to convert it into basic control-flow
+ // operations.
+ } else {
+ rewriter.inlineRegionBefore(ompOp->getRegion(0), nestBody);
+ auto indVarArg = outerLoop->getRegion(0).front().getArgument(0);
+ // fir::convertDoLoopToCFG expects the induction variable to be of type
+ // index while the OpenMP LoopNestOp can have indices of different
+ // types. We need to work around it.
+ if (indVarArg.getType() != indexType)
+ indVarArg.setType(indexType);
+
+ auto loopBlocks =
+ fir::convertDoLoopToCFG(outerLoop, rewriter, false, false);
----------------
mrkajetanp wrote:
There is one 4 lines above, just above the 'else', but I'll move it closer to
this for visibility.
https://github.com/llvm/llvm-project/pull/150269
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
