Skip to content

[MLIR][Flang][OpenMP] Make omp.wsloop into a loop wrapper #88403

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 8 commits into from
Closed

[MLIR][Flang][OpenMP] Make omp.wsloop into a loop wrapper #88403

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
281121e
[MLIR][OpenMP] Add omp.loop_nest operation
skatrak Mar 29, 2024
2452bc7
[MLIR][OpenMP] Introduce the LoopWrapperInterface
skatrak Apr 1, 2024
e3c440c
Update op description according to review comments
skatrak Apr 2, 2024
904f274
Address review comments
skatrak Apr 2, 2024
99e8123
Merge branch 'users/skatrak/spr/loop-nest-01-mlir' into users/skatrak…
skatrak Apr 2, 2024
3bcb419
[MLIR][Flang][OpenMP] Make omp.simdloop into a loop wrapper
skatrak Apr 2, 2024
50e4f78
[MLIR][Flang][OpenMP] Make omp.wsloop into a loop wrapper
skatrak Apr 11, 2024
1dfae18
Merge branch 'main' into users/skatrak/spr/loop-nest-04-wsloop-mlir
skatrak Apr 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 27 additions & 21 deletions flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,8 +135,14 @@ void DataSharingProcessor::insertBarrier() {
}

void DataSharingProcessor::insertLastPrivateCompare(mlir::Operation *op) {
mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
mlir::omp::LoopNestOp loopOp;
if (auto wrapper = mlir::dyn_cast<mlir::omp::LoopWrapperInterface>(op))
loopOp = wrapper.isWrapper()
? mlir::cast<mlir::omp::LoopNestOp>(wrapper.getWrappedLoop())
: nullptr;

bool cmpCreated = false;
mlir::OpBuilder::InsertPoint localInsPt = firOpBuilder.saveInsertionPoint();
for (const omp::Clause &clause : clauses) {
if (clause.id != llvm::omp::OMPC_lastprivate)
continue;
Expand Down Expand Up @@ -215,33 +221,34 @@ void DataSharingProcessor::insertLastPrivateCompare(mlir::Operation *op) {
// Update the original variable just before exiting the worksharing
// loop. Conversion as follows:
//
// omp.wsloop {
// omp.wsloop { ...
// ... store
// store ===> %v = arith.addi %iv, %step
// omp.yield %cmp = %step < 0 ? %v < %ub : %v > %ub
// } fir.if %cmp {
// fir.store %v to %loopIV
// ^%lpv_update_blk:
// }
// omp.yield
// }
//
// omp.wsloop { omp.wsloop {
// omp.loop_nest { omp.loop_nest {
// ... ...
// store ===> store
// omp.yield %v = arith.addi %iv, %step
// } %cmp = %step < 0 ? %v < %ub : %v > %ub
// omp.terminator fir.if %cmp {
// } fir.store %v to %loopIV
// ^%lpv_update_blk:
// }
// omp.yield
// }
// omp.terminator
// }

// Only generate the compare once in presence of multiple LastPrivate
// clauses.
if (cmpCreated)
continue;
cmpCreated = true;

mlir::Location loc = op->getLoc();
mlir::Operation *lastOper = op->getRegion(0).back().getTerminator();
mlir::Location loc = loopOp.getLoc();
mlir::Operation *lastOper = loopOp.getRegion().back().getTerminator();
firOpBuilder.setInsertionPoint(lastOper);

mlir::Value iv = op->getRegion(0).front().getArguments()[0];
mlir::Value ub =
mlir::dyn_cast<mlir::omp::WsloopOp>(op).getUpperBound()[0];
mlir::Value step = mlir::dyn_cast<mlir::omp::WsloopOp>(op).getStep()[0];
mlir::Value iv = loopOp.getIVs()[0];
mlir::Value ub = loopOp.getUpperBound()[0];
mlir::Value step = loopOp.getStep()[0];

// v = iv + step
// cmp = step < 0 ? v < ub : v > ub
Expand All @@ -260,15 +267,14 @@ void DataSharingProcessor::insertLastPrivateCompare(mlir::Operation *op) {
auto ifOp = firOpBuilder.create<fir::IfOp>(loc, cmpOp, /*else*/ false);
firOpBuilder.setInsertionPointToStart(&ifOp.getThenRegion().front());
assert(loopIV && "loopIV was not set");
firOpBuilder.create<fir::StoreOp>(op->getLoc(), v, loopIV);
firOpBuilder.create<fir::StoreOp>(loopOp.getLoc(), v, loopIV);
lastPrivIP = firOpBuilder.saveInsertionPoint();
} else {
TODO(converter.getCurrentLocation(),
"lastprivate clause in constructs other than "
"simd/worksharing-loop");
}
}
firOpBuilder.restoreInsertionPoint(localInsPt);
}

void DataSharingProcessor::collectSymbols(
Expand Down
115 changes: 42 additions & 73 deletions flang/lib/Lower/OpenMP/OpenMP.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,9 @@ getDeclareTargetFunctionDevice(
static llvm::SmallVector<const Fortran::semantics::Symbol *>
genLoopVars(mlir::Operation *op, Fortran::lower::AbstractConverter &converter,
mlir::Location &loc,
llvm::ArrayRef<const Fortran::semantics::Symbol *> args) {
llvm::ArrayRef<const Fortran::semantics::Symbol *> args,
llvm::ArrayRef<const Fortran::semantics::Symbol *> wrapperSyms = {},
llvm::ArrayRef<mlir::BlockArgument> wrapperArgs = {}) {
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
auto &region = op->getRegion(0);

Expand All @@ -380,6 +382,14 @@ genLoopVars(mlir::Operation *op, Fortran::lower::AbstractConverter &converter,
llvm::SmallVector<mlir::Type> tiv(args.size(), loopVarType);
llvm::SmallVector<mlir::Location> locs(args.size(), loc);
firOpBuilder.createBlock(&region, {}, tiv, locs);

// Bind the entry block arguments of parent wrappers to the corresponding
// symbols. Do it here so that any hlfir.declare operations created as a
// result are inserted inside of the omp.loop_nest rather than the wrapper
// operations.
for (auto [arg, prv] : llvm::zip_equal(wrapperSyms, wrapperArgs))
converter.bindSymbol(*arg, prv);

// The argument is not currently in memory, so make a temporary for the
// argument, and store it there, then bind that location to the argument.
mlir::Operation *storeOp = nullptr;
Expand Down Expand Up @@ -410,58 +420,6 @@ static void genReductionVars(
}
}

static llvm::SmallVector<const Fortran::semantics::Symbol *>
genLoopAndReductionVars(
mlir::Operation *op, Fortran::lower::AbstractConverter &converter,
mlir::Location &loc,
llvm::ArrayRef<const Fortran::semantics::Symbol *> loopArgs,
llvm::ArrayRef<const Fortran::semantics::Symbol *> reductionArgs,
llvm::ArrayRef<mlir::Type> reductionTypes) {
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();

llvm::SmallVector<mlir::Type> blockArgTypes;
llvm::SmallVector<mlir::Location> blockArgLocs;
blockArgTypes.reserve(loopArgs.size() + reductionArgs.size());
blockArgLocs.reserve(blockArgTypes.size());
mlir::Block *entryBlock;

if (loopArgs.size()) {
std::size_t loopVarTypeSize = 0;
for (const Fortran::semantics::Symbol *arg : loopArgs)
loopVarTypeSize = std::max(loopVarTypeSize, arg->GetUltimate().size());
mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
std::fill_n(std::back_inserter(blockArgTypes), loopArgs.size(),
loopVarType);
std::fill_n(std::back_inserter(blockArgLocs), loopArgs.size(), loc);
}
if (reductionArgs.size()) {
llvm::copy(reductionTypes, std::back_inserter(blockArgTypes));
std::fill_n(std::back_inserter(blockArgLocs), reductionArgs.size(), loc);
}
entryBlock = firOpBuilder.createBlock(&op->getRegion(0), {}, blockArgTypes,
blockArgLocs);
// The argument is not currently in memory, so make a temporary for the
// argument, and store it there, then bind that location to the argument.
if (loopArgs.size()) {
mlir::Operation *storeOp = nullptr;
for (auto [argIndex, argSymbol] : llvm::enumerate(loopArgs)) {
mlir::Value indexVal =
fir::getBase(op->getRegion(0).front().getArgument(argIndex));
storeOp =
createAndSetPrivatizedLoopVar(converter, loc, indexVal, argSymbol);
}
firOpBuilder.setInsertionPointAfter(storeOp);
}
// Bind the reduction arguments to their block arguments
for (auto [arg, prv] : llvm::zip_equal(
reductionArgs,
llvm::drop_begin(entryBlock->getArguments(), loopArgs.size()))) {
converter.bindSymbol(*arg, prv);
}

return llvm::SmallVector<const Fortran::semantics::Symbol *>(loopArgs);
}

static void
markDeclareTarget(mlir::Operation *op,
Fortran::lower::AbstractConverter &converter,
Expand Down Expand Up @@ -1292,20 +1250,16 @@ static void genWsloopClauses(
Fortran::lower::AbstractConverter &converter,
Fortran::semantics::SemanticsContext &semaCtx,
Fortran::lower::StatementContext &stmtCtx,
Fortran::lower::pft::Evaluation &eval,
const Fortran::parser::OmpClauseList &beginClauses,
const Fortran::parser::OmpClauseList *endClauses, mlir::Location loc,
mlir::omp::WsloopClauseOps &clauseOps,
llvm::SmallVectorImpl<const Fortran::semantics::Symbol *> &iv,
llvm::SmallVectorImpl<mlir::Type> &reductionTypes,
llvm::SmallVectorImpl<const Fortran::semantics::Symbol *> &reductionSyms) {
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
ClauseProcessor bcp(converter, semaCtx, beginClauses);
bcp.processCollapse(loc, eval, clauseOps, iv);
bcp.processOrdered(clauseOps);
bcp.processReduction(loc, clauseOps, &reductionTypes, &reductionSyms);
bcp.processSchedule(stmtCtx, clauseOps);
clauseOps.loopInclusiveAttr = firOpBuilder.getUnitAttr();
// TODO Support delayed privatization.

if (ReductionProcessor::doReductionByRef(clauseOps.reductionVars))
Expand Down Expand Up @@ -1844,34 +1798,49 @@ genWsloopOp(Fortran::lower::AbstractConverter &converter,
Fortran::lower::pft::Evaluation &eval, mlir::Location loc,
const Fortran::parser::OmpClauseList &beginClauseList,
const Fortran::parser::OmpClauseList *endClauseList) {
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
DataSharingProcessor dsp(converter, semaCtx, beginClauseList, eval);
dsp.processStep1();

Fortran::lower::StatementContext stmtCtx;
mlir::omp::WsloopClauseOps clauseOps;
mlir::omp::LoopNestClauseOps loopClauseOps;
mlir::omp::WsloopClauseOps wsClauseOps;
llvm::SmallVector<const Fortran::semantics::Symbol *> iv;
llvm::SmallVector<mlir::Type> reductionTypes;
llvm::SmallVector<const Fortran::semantics::Symbol *> reductionSyms;
genWsloopClauses(converter, semaCtx, stmtCtx, eval, beginClauseList,
endClauseList, loc, clauseOps, iv, reductionTypes,
reductionSyms);
genLoopNestClauses(converter, semaCtx, eval, beginClauseList, loc,
loopClauseOps, iv);
genWsloopClauses(converter, semaCtx, stmtCtx, beginClauseList, endClauseList,
loc, wsClauseOps, reductionTypes, reductionSyms);

// Create omp.wsloop wrapper and populate entry block arguments with reduction
// variables.
auto wsloopOp = firOpBuilder.create<mlir::omp::WsloopOp>(loc, wsClauseOps);
llvm::SmallVector<mlir::Location> reductionLocs(reductionSyms.size(), loc);
mlir::Block *wsloopEntryBlock = firOpBuilder.createBlock(
&wsloopOp.getRegion(), {}, reductionTypes, reductionLocs);
firOpBuilder.setInsertionPoint(
Fortran::lower::genOpenMPTerminator(firOpBuilder, wsloopOp, loc));

// Create nested omp.loop_nest and fill body with loop contents.
auto loopOp = firOpBuilder.create<mlir::omp::LoopNestOp>(loc, loopClauseOps);

auto *nestedEval = getCollapsedLoopEval(
eval, Fortran::lower::getCollapseValue(beginClauseList));

auto ivCallback = [&](mlir::Operation *op) {
return genLoopAndReductionVars(op, converter, loc, iv, reductionSyms,
reductionTypes);
return genLoopVars(op, converter, loc, iv, reductionSyms,
wsloopEntryBlock->getArguments());
};

return genOpWithBody<mlir::omp::WsloopOp>(
OpWithBodyGenInfo(converter, semaCtx, loc, *nestedEval,
llvm::omp::Directive::OMPD_do)
.setClauses(&beginClauseList)
.setDataSharingProcessor(&dsp)
.setReductions(&reductionSyms, &reductionTypes)
.setGenRegionEntryCb(ivCallback),
clauseOps);
createBodyOfOp(*loopOp,
OpWithBodyGenInfo(converter, semaCtx, loc, *nestedEval,
llvm::omp::Directive::OMPD_do)
.setClauses(&beginClauseList)
.setDataSharingProcessor(&dsp)
.setReductions(&reductionSyms, &reductionTypes)
.setGenRegionEntryCb(ivCallback));
return wsloopOp;
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -2542,8 +2511,8 @@ static void genOMP(Fortran::lower::AbstractConverter &converter,
mlir::Operation *Fortran::lower::genOpenMPTerminator(fir::FirOpBuilder &builder,
mlir::Operation *op,
mlir::Location loc) {
if (mlir::isa<mlir::omp::WsloopOp, mlir::omp::DeclareReductionOp,
mlir::omp::AtomicUpdateOp, mlir::omp::LoopNestOp>(op))
if (mlir::isa<mlir::omp::AtomicUpdateOp, mlir::omp::DeclareReductionOp,
mlir::omp::LoopNestOp>(op))
return builder.create<mlir::omp::YieldOp>(loc);
return builder.create<mlir::omp::TerminatorOp>(loc);
}
Expand Down
Loading