================
@@ -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);
+ auto *conditionalBlock = loopBlocks.first;
+ auto *firstBlock =
+ conditionalBlock->getNextNode(); // Start of the loop body
+ auto *lastBlock = loopBlocks.second; // Incrementing induction
variables
+
+ // If the induction variable is used within the loop and was originally
+ // not of type index, then we need to add a convert to the original
type
+ // and replace its uses inside the loop body.
+ if (oldIndexType != indexType) {
+ indVarArg = conditionalBlock->getArgument(0);
+ builder.setInsertionPointToStart(firstBlock);
+ auto convertedIndVar =
+ builder.createConvert(loc, oldIndexType, indVarArg);
+ rewriter.replaceUsesWithIf(
+ indVarArg, convertedIndVar, [&](auto &use) -> bool {
+ return use.getOwner() != convertedIndVar.getDefiningOp() &&
+ use.getOwner()->getBlock() != lastBlock;
+ });
+ }
+
+ // There might be an unused convert and an unused argument to the
block.
+ // If so, remove them.
+ if (lastBlock->front().getUses().empty())
+ lastBlock->front().erase();
+ for (auto arg : lastBlock->getArguments()) {
+ if (arg.getUses().empty())
+ lastBlock->eraseArgument(arg.getArgNumber());
+ }
+
+ // Any loop blocks that end in omp.yield should just branch to
+ // lastBlock.
+ for (auto *loopBlock = conditionalBlock; loopBlock != lastBlock;
+ loopBlock = loopBlock->getNextNode()) {
+ if (auto yieldOp =
+ mlir::dyn_cast<mlir::omp::YieldOp>(loopBlock->back())) {
+ builder.setInsertionPointToEnd(loopBlock);
+ builder.create<mlir::cf::BranchOp>(loc, lastBlock);
+ rewriter.eraseOp(yieldOp);
----------------
tblah wrote:
Again, I think we need to handle yield ops with results here, or if you are
certain this can never happen with the current operations this runs on, could
you add an assertion to catch if this changes in the future.
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
