60ee0560da
`replaceAllUsesWith` is not safe to use in a dialect conversion and will be deactivated soon (#154112). Fix commit fixes some API violations. Also some general improvements.
210 lines
8.0 KiB
C++
210 lines
8.0 KiB
C++
//===-- SimdOnly.cpp ------------------------------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "flang/Optimizer/Builder/FIRBuilder.h"
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
|
|
#include "mlir/IR/MLIRContext.h"
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "llvm/Support/Debug.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"
|
|
|
|
/// Rewrite and remove OpenMP operations left after the parse tree rewriting for
|
|
/// -fopenmp-simd is done. If possible, OpenMP constructs should be rewritten at
|
|
/// the parse tree stage. This pass is supposed to only handle complexities
|
|
/// around untangling composite simd constructs, and perform the necessary
|
|
/// cleanup.
|
|
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>() &&
|
|
(mlir::isa<mlir::omp::YieldOp>(op) ||
|
|
mlir::isa<mlir::omp::ScanOp>(op) ||
|
|
mlir::isa<mlir::omp::LoopNestOp>(op) ||
|
|
mlir::isa<mlir::omp::TerminatorOp>(op)))
|
|
return rewriter.notifyMatchFailure(op, "Op is part of a simd construct");
|
|
|
|
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");
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << "SimdOnlyPass matched OpenMP op:\n");
|
|
LLVM_DEBUG(op->dump());
|
|
|
|
auto eraseUnlessUsedBySimd = [&](mlir::Operation *ompOp,
|
|
mlir::StringAttr name) {
|
|
if (auto uses =
|
|
mlir::SymbolTable::getSymbolUses(name, op->getParentOp())) {
|
|
for (auto &use : *uses)
|
|
if (mlir::isa<mlir::omp::SimdOp>(use.getUser()))
|
|
return rewriter.notifyMatchFailure(op,
|
|
"Op used by a simd construct");
|
|
}
|
|
rewriter.eraseOp(ompOp);
|
|
return mlir::success();
|
|
};
|
|
|
|
if (auto ompOp = mlir::dyn_cast<mlir::omp::PrivateClauseOp>(op))
|
|
return eraseUnlessUsedBySimd(ompOp, ompOp.getSymNameAttr());
|
|
if (auto ompOp = mlir::dyn_cast<mlir::omp::DeclareReductionOp>(op))
|
|
return eraseUnlessUsedBySimd(ompOp, ompOp.getSymNameAttr());
|
|
|
|
// Might be left over from rewriting composite simd with target map
|
|
if (mlir::isa<mlir::omp::MapBoundsOp>(op)) {
|
|
rewriter.eraseOp(op);
|
|
return mlir::success();
|
|
}
|
|
if (auto mapInfoOp = mlir::dyn_cast<mlir::omp::MapInfoOp>(op)) {
|
|
rewriter.replaceOp(mapInfoOp, {mapInfoOp.getVarPtr()});
|
|
return mlir::success();
|
|
}
|
|
|
|
// Might be leftover after parse tree rewriting
|
|
if (auto threadPrivateOp = mlir::dyn_cast<mlir::omp::ThreadprivateOp>(op)) {
|
|
rewriter.replaceOp(threadPrivateOp, {threadPrivateOp.getSymAddr()});
|
|
return mlir::success();
|
|
}
|
|
|
|
fir::FirOpBuilder builder(rewriter, op);
|
|
mlir::Location loc = op->getLoc();
|
|
|
|
auto inlineSimpleOp = [&](mlir::Operation *ompOp) -> bool {
|
|
if (!ompOp)
|
|
return false;
|
|
|
|
assert("OpenMP operation has one region" && ompOp->getNumRegions() == 1);
|
|
|
|
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, ompOp, {});
|
|
} 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);
|
|
mlir::cf::BranchOp::create(builder, 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);
|
|
mlir::cf::BranchOp::create(builder, loc, newBlock);
|
|
rewriter.eraseOp(terminatorOp);
|
|
}
|
|
}
|
|
|
|
rewriter.inlineRegionBefore(ompOp->getRegion(0), newBlock);
|
|
}
|
|
|
|
rewriter.eraseOp(op);
|
|
return true;
|
|
};
|
|
|
|
// Remove ops that will be surrounding simd once a composite simd construct
|
|
// goes through the codegen stage. All of the other ones should have alredy
|
|
// been removed in the parse tree rewriting stage.
|
|
if (inlineSimpleOp(mlir::dyn_cast<mlir::omp::TeamsOp>(op)) ||
|
|
inlineSimpleOp(mlir::dyn_cast<mlir::omp::ParallelOp>(op)) ||
|
|
inlineSimpleOp(mlir::dyn_cast<mlir::omp::TargetOp>(op)) ||
|
|
inlineSimpleOp(mlir::dyn_cast<mlir::omp::WsloopOp>(op)) ||
|
|
inlineSimpleOp(mlir::dyn_cast<mlir::omp::DistributeOp>(op)))
|
|
return mlir::success();
|
|
|
|
op->emitOpError("left unhandled after SimdOnly pass.");
|
|
return mlir::failure();
|
|
}
|
|
};
|
|
|
|
class SimdOnlyPass : public flangomp::impl::SimdOnlyPassBase<SimdOnlyPass> {
|
|
|
|
public:
|
|
SimdOnlyPass() = default;
|
|
|
|
void runOnOperation() override {
|
|
mlir::ModuleOp module = getOperation();
|
|
|
|
mlir::MLIRContext *context = &getContext();
|
|
mlir::RewritePatternSet patterns(context);
|
|
patterns.insert<SimdOnlyConversionPattern>(context);
|
|
|
|
mlir::GreedyRewriteConfig config;
|
|
// Prevent the pattern driver from merging blocks.
|
|
config.setRegionSimplificationLevel(
|
|
mlir::GreedySimplifyRegionLevel::Disabled);
|
|
|
|
if (mlir::failed(
|
|
mlir::applyPatternsGreedily(module, std::move(patterns), config))) {
|
|
mlir::emitError(module.getLoc(), "Error in SimdOnly conversion pass");
|
|
signalPassFailure();
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace
|