eaa0d281b6
This patch is a required change for the device side IR to maintain apporpiate links for declare target variables to their global variables for later lowering. It is also a requirement to clone over map bounds and entry operations to maintain the correct information for later lowering of the IR. It simply tries to clone over the relevant information maintaining the appropriate links they would have maintained prior to the pass, rather than redirecting them to new function arguments which causes a loss of information in the case of Declare Target and map information. Depends on D158734 reviewers: TIFitis, razvanlupusoru Differential Revision: https://reviews.llvm.org/D158735
257 lines
10 KiB
C++
257 lines
10 KiB
C++
#include "flang/Optimizer/Dialect/FIRDialect.h"
|
|
#include "flang/Optimizer/Dialect/FIROps.h"
|
|
#include "flang/Optimizer/Dialect/FIRType.h"
|
|
#include "flang/Optimizer/Support/InternalNames.h"
|
|
#include "flang/Optimizer/Transforms/Passes.h"
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
|
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
|
|
#include "mlir/IR/BuiltinDialect.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/IRMapping.h"
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/IR/SymbolTable.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "mlir/Transforms/RegionUtils.h"
|
|
#include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
|
|
|
|
namespace fir {
|
|
#define GEN_PASS_DEF_OMPEARLYOUTLININGPASS
|
|
#include "flang/Optimizer/Transforms/Passes.h.inc"
|
|
} // namespace fir
|
|
|
|
namespace {
|
|
class OMPEarlyOutliningPass
|
|
: public fir::impl::OMPEarlyOutliningPassBase<OMPEarlyOutliningPass> {
|
|
|
|
std::string getOutlinedFnName(llvm::StringRef parentName, unsigned count) {
|
|
return std::string(parentName) + "_omp_outline_" + std::to_string(count);
|
|
}
|
|
|
|
// Given a value this function will iterate over an operators results
|
|
// and return the relevant index for the result the value corresponds to.
|
|
// There may be a simpler way to do this however.
|
|
unsigned getResultIndex(mlir::Value value, mlir::Operation *op) {
|
|
for (unsigned i = 0; i < op->getNumResults(); ++i) {
|
|
if (op->getResult(i) == value)
|
|
return i;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
bool isDeclareTargetOp(mlir::Operation *op) {
|
|
if (fir::AddrOfOp addressOfOp = mlir::dyn_cast<fir::AddrOfOp>(op))
|
|
if (fir::GlobalOp gOp = mlir::dyn_cast<fir::GlobalOp>(
|
|
addressOfOp->getParentOfType<mlir::ModuleOp>().lookupSymbol(
|
|
addressOfOp.getSymbol())))
|
|
if (auto declareTargetGlobal =
|
|
llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
|
|
gOp.getOperation()))
|
|
if (declareTargetGlobal.isDeclareTarget())
|
|
return true;
|
|
return false;
|
|
}
|
|
|
|
// Currently used for cloning arguments that are nested. Should be
|
|
// extendable where required, perhaps via operation
|
|
// specialisation/overloading, if something needs specialised handling.
|
|
// NOTE: Results in duplication of some values that would otherwise be
|
|
// a single SSA value shared between operations, this is tidied up on
|
|
// lowering to some extent.
|
|
mlir::Operation *
|
|
cloneArgAndChildren(mlir::OpBuilder &builder, mlir::Operation *op,
|
|
llvm::SetVector<mlir::Value> &inputs,
|
|
mlir::Block::BlockArgListType &newInputs) {
|
|
mlir::IRMapping valueMap;
|
|
for (auto opValue : op->getOperands()) {
|
|
if (opValue.getDefiningOp()) {
|
|
auto resIdx = getResultIndex(opValue, opValue.getDefiningOp());
|
|
valueMap.map(opValue,
|
|
cloneArgAndChildren(builder, opValue.getDefiningOp(),
|
|
inputs, newInputs)
|
|
->getResult(resIdx));
|
|
} else {
|
|
for (auto inArg : llvm::zip(inputs, newInputs)) {
|
|
if (opValue == std::get<0>(inArg))
|
|
valueMap.map(opValue, std::get<1>(inArg));
|
|
}
|
|
}
|
|
}
|
|
|
|
return builder.clone(*op, valueMap);
|
|
}
|
|
|
|
void cloneMapOpVariables(mlir::OpBuilder &builder, mlir::IRMapping &valueMap,
|
|
mlir::IRMapping &mapInfoMap,
|
|
llvm::SetVector<mlir::Value> &inputs,
|
|
mlir::Block::BlockArgListType &newInputs,
|
|
mlir::Value varPtr) {
|
|
if (fir::BoxAddrOp boxAddrOp =
|
|
mlir::dyn_cast_if_present<fir::BoxAddrOp>(varPtr.getDefiningOp())) {
|
|
mlir::Value newV =
|
|
cloneArgAndChildren(builder, boxAddrOp, inputs, newInputs)
|
|
->getResult(0);
|
|
mapInfoMap.map(varPtr, newV);
|
|
valueMap.map(boxAddrOp, newV);
|
|
return;
|
|
}
|
|
|
|
if (varPtr.getDefiningOp() && isDeclareTargetOp(varPtr.getDefiningOp())) {
|
|
fir::AddrOfOp addrOp =
|
|
mlir::dyn_cast<fir::AddrOfOp>(varPtr.getDefiningOp());
|
|
mlir::Value newV = builder.clone(*addrOp)->getResult(0);
|
|
mapInfoMap.map(varPtr, newV);
|
|
valueMap.map(addrOp, newV);
|
|
return;
|
|
}
|
|
|
|
for (auto inArg : llvm::zip(inputs, newInputs)) {
|
|
if (varPtr == std::get<0>(inArg))
|
|
mapInfoMap.map(varPtr, std::get<1>(inArg));
|
|
}
|
|
}
|
|
|
|
mlir::func::FuncOp outlineTargetOp(mlir::OpBuilder &builder,
|
|
mlir::omp::TargetOp &targetOp,
|
|
mlir::func::FuncOp &parentFunc,
|
|
unsigned count) {
|
|
// NOTE: once implicit captures are handled appropriately in the initial
|
|
// PFT lowering if it is possible, we can remove the usage of
|
|
// getUsedValuesDefinedAbove and instead just iterate over the target op's
|
|
// operands (or just the map arguments) and perhaps refactor this function
|
|
// a little.
|
|
// Collect inputs
|
|
llvm::SetVector<mlir::Value> inputs;
|
|
mlir::Region &targetRegion = targetOp.getRegion();
|
|
mlir::getUsedValuesDefinedAbove(targetRegion, inputs);
|
|
|
|
// filter out declareTarget and map entries which are specially handled
|
|
// at the moment, so we do not wish these to end up as function arguments
|
|
// which would just be more noise in the IR.
|
|
for (auto value : inputs)
|
|
if (value.getDefiningOp())
|
|
if (mlir::isa<mlir::omp::MapInfoOp>(value.getDefiningOp()) ||
|
|
isDeclareTargetOp(value.getDefiningOp()))
|
|
inputs.remove(value);
|
|
|
|
// Create new function and initialize
|
|
mlir::FunctionType funcType = builder.getFunctionType(
|
|
mlir::TypeRange(inputs.getArrayRef()), mlir::TypeRange());
|
|
std::string parentName(parentFunc.getName());
|
|
std::string funcName = getOutlinedFnName(parentName, count);
|
|
auto loc = targetOp.getLoc();
|
|
mlir::func::FuncOp newFunc =
|
|
mlir::func::FuncOp::create(loc, funcName, funcType);
|
|
mlir::Block *entryBlock = newFunc.addEntryBlock();
|
|
builder.setInsertionPointToStart(entryBlock);
|
|
mlir::Block::BlockArgListType newInputs = entryBlock->getArguments();
|
|
|
|
// Set the declare target information, the outlined function
|
|
// is always a host function.
|
|
if (auto parentDTOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
|
|
parentFunc.getOperation()))
|
|
if (auto newDTOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(
|
|
newFunc.getOperation()))
|
|
newDTOp.setDeclareTarget(mlir::omp::DeclareTargetDeviceType::host,
|
|
parentDTOp.getDeclareTargetCaptureClause());
|
|
|
|
// Set the early outlining interface parent name
|
|
if (auto earlyOutlineOp =
|
|
llvm::dyn_cast<mlir::omp::EarlyOutliningInterface>(
|
|
newFunc.getOperation()))
|
|
earlyOutlineOp.setParentName(parentName);
|
|
|
|
// The value map for the newly generated Target Operation, we must
|
|
// remap most of the input.
|
|
mlir::IRMapping valueMap;
|
|
|
|
// Special handling for map, declare target and regular map variables
|
|
// are handled slightly differently for the moment, declare target has
|
|
// its addressOfOp cloned over, whereas we skip it for the regular map
|
|
// variables. We need knowledge of which global is linked to the map
|
|
// operation for declare target, whereas we aren't bothered for the
|
|
// regular map variables for the moment. We could treat both the same,
|
|
// however, cloning across the minimum for the moment to avoid
|
|
// optimisations breaking segments of the lowering seems prudent as this
|
|
// was the original intent of the pass.
|
|
for (auto oper : targetOp.getOperation()->getOperands()) {
|
|
if (auto mapEntry =
|
|
mlir::dyn_cast<mlir::omp::MapInfoOp>(oper.getDefiningOp())) {
|
|
mlir::IRMapping mapInfoMap;
|
|
for (auto bound : mapEntry.getBounds()) {
|
|
if (auto mapEntryBound = mlir::dyn_cast<mlir::omp::DataBoundsOp>(
|
|
bound.getDefiningOp())) {
|
|
mapInfoMap.map(bound, cloneArgAndChildren(builder, mapEntryBound,
|
|
inputs, newInputs)
|
|
->getResult(0));
|
|
}
|
|
}
|
|
|
|
cloneMapOpVariables(builder, valueMap, mapInfoMap, inputs, newInputs,
|
|
mapEntry.getVarPtr());
|
|
|
|
if (mapEntry.getVarPtrPtr())
|
|
cloneMapOpVariables(builder, valueMap, mapInfoMap, inputs, newInputs,
|
|
mapEntry.getVarPtrPtr());
|
|
|
|
valueMap.map(
|
|
mapEntry,
|
|
builder.clone(*mapEntry.getOperation(), mapInfoMap)->getResult(0));
|
|
}
|
|
}
|
|
|
|
for (auto inArg : llvm::zip(inputs, newInputs))
|
|
valueMap.map(std::get<0>(inArg), std::get<1>(inArg));
|
|
|
|
// Clone the target op into the new function
|
|
builder.clone(*(targetOp.getOperation()), valueMap);
|
|
|
|
// Create return op
|
|
builder.create<mlir::func::ReturnOp>(loc);
|
|
|
|
return newFunc;
|
|
}
|
|
|
|
// Returns true if a target region was found int the function.
|
|
bool outlineTargetOps(mlir::OpBuilder &builder,
|
|
mlir::func::FuncOp &functionOp,
|
|
mlir::ModuleOp &moduleOp,
|
|
llvm::SmallVectorImpl<mlir::func::FuncOp> &newFuncs) {
|
|
unsigned count = 0;
|
|
for (auto TargetOp : functionOp.getOps<mlir::omp::TargetOp>()) {
|
|
mlir::func::FuncOp outlinedFunc =
|
|
outlineTargetOp(builder, TargetOp, functionOp, count);
|
|
newFuncs.push_back(outlinedFunc);
|
|
count++;
|
|
}
|
|
return count > 0;
|
|
}
|
|
|
|
void runOnOperation() override {
|
|
mlir::ModuleOp moduleOp = getOperation();
|
|
mlir::MLIRContext *context = &getContext();
|
|
mlir::OpBuilder builder(context);
|
|
llvm::SmallVector<mlir::func::FuncOp> newFuncs;
|
|
|
|
for (auto functionOp :
|
|
llvm::make_early_inc_range(moduleOp.getOps<mlir::func::FuncOp>())) {
|
|
bool outlined = outlineTargetOps(builder, functionOp, moduleOp, newFuncs);
|
|
if (outlined)
|
|
functionOp.erase();
|
|
}
|
|
|
|
for (auto newFunc : newFuncs)
|
|
moduleOp.push_back(newFunc);
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
namespace fir {
|
|
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
|
|
createOMPEarlyOutliningPass() {
|
|
return std::make_unique<OMPEarlyOutliningPass>();
|
|
}
|
|
} // namespace fir
|