Skip to content

[flang][mlir][OpenMP] Add support for COPYPRIVATE #73128

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

[flang][mlir][OpenMP] Add support for COPYPRIVATE #73128

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
8ce1986
[flang][mlir][OpenMP] Add support for COPYPRIVATE
luporl Dec 20, 2023
3a5e062
Fix mlir/test/Dialect/OpenMP/invalid.mlir
luporl Dec 20, 2023
029c364
Use getTypeAsString() to build copy function name
luporl Dec 21, 2023
a45b09b
Add integration test
luporl Dec 21, 2023
d5b44de
Use ArrayRef in OMPIRBuilder's createSingle()
luporl Dec 21, 2023
0bc008a
clang-format
luporl Dec 21, 2023
f2f25b8
Add missing tests
luporl Jan 4, 2024
ef5a0b7
Add support for vla/assumed-shape arrays
luporl Jan 22, 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
3 changes: 3 additions & 0 deletions flang/include/flang/Lower/AbstractConverter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,9 @@ class AbstractConverter {
const Fortran::semantics::Symbol &sym,
mlir::OpBuilder::InsertPoint *copyAssignIP = nullptr) = 0;

virtual void copyVar(mlir::Location loc, mlir::Value dst,
mlir::Value src) = 0;

/// For a given symbol, check if it is present in the inner-most
/// level of the symbol map.
virtual bool isPresentShallowLookup(Fortran::semantics::Symbol &sym) = 0;
Expand Down
137 changes: 79 additions & 58 deletions flang/lib/Lower/Bridge.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -743,6 +743,11 @@ class FirConverter : public Fortran::lower::AbstractConverter {
});
}

void copyVar(mlir::Location loc, mlir::Value dst,
mlir::Value src) override final {
copyVarHLFIR(loc, dst, src);
}

void copyHostAssociateVar(
const Fortran::semantics::Symbol &sym,
mlir::OpBuilder::InsertPoint *copyAssignIP = nullptr) override final {
Expand Down Expand Up @@ -777,64 +782,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
rhs_sb = &hsb;
}

mlir::Location loc = genLocation(sym.name());

if (lowerToHighLevelFIR()) {
hlfir::Entity lhs{lhs_sb->getAddr()};
hlfir::Entity rhs{rhs_sb->getAddr()};
// Temporary_lhs is set to true in hlfir.assign below to avoid user
// assignment to be used and finalization to be called on the LHS.
// This may or may not be correct but mimics the current behaviour
// without HLFIR.
auto copyData = [&](hlfir::Entity l, hlfir::Entity r) {
// Dereference RHS and load it if trivial scalar.
r = hlfir::loadTrivialScalar(loc, *builder, r);
builder->create<hlfir::AssignOp>(
loc, r, l,
/*isWholeAllocatableAssignment=*/false,
/*keepLhsLengthInAllocatableAssignment=*/false,
/*temporary_lhs=*/true);
};
if (lhs.isAllocatable()) {
// Deep copy allocatable if it is allocated.
// Note that when allocated, the RHS is already allocated with the LHS
// shape for copy on entry in createHostAssociateVarClone.
// For lastprivate, this assumes that the RHS was not reallocated in
// the OpenMP region.
lhs = hlfir::derefPointersAndAllocatables(loc, *builder, lhs);
mlir::Value addr = hlfir::genVariableRawAddress(loc, *builder, lhs);
mlir::Value isAllocated = builder->genIsNotNullAddr(loc, addr);
builder->genIfThen(loc, isAllocated)
.genThen([&]() {
// Copy the DATA, not the descriptors.
copyData(lhs, rhs);
})
.end();
} else if (lhs.isPointer()) {
// Set LHS target to the target of RHS (do not copy the RHS
// target data into the LHS target storage).
auto loadVal = builder->create<fir::LoadOp>(loc, rhs);
builder->create<fir::StoreOp>(loc, loadVal, lhs);
} else {
// Non ALLOCATABLE/POINTER variable. Simple DATA copy.
copyData(lhs, rhs);
}
} else {
fir::ExtendedValue lhs = symBoxToExtendedValue(*lhs_sb);
fir::ExtendedValue rhs = symBoxToExtendedValue(*rhs_sb);
mlir::Type symType = genType(sym);
if (auto seqTy = symType.dyn_cast<fir::SequenceType>()) {
Fortran::lower::StatementContext stmtCtx;
Fortran::lower::createSomeArrayAssignment(*this, lhs, rhs, localSymbols,
stmtCtx);
stmtCtx.finalizeAndReset();
} else if (lhs.getBoxOf<fir::CharBoxValue>()) {
fir::factory::CharacterExprHelper{*builder, loc}.createAssign(lhs, rhs);
} else {
auto loadVal = builder->create<fir::LoadOp>(loc, fir::getBase(rhs));
builder->create<fir::StoreOp>(loc, loadVal, fir::getBase(lhs));
}
}
copyVar(sym, *lhs_sb, *rhs_sb);

if (copyAssignIP && copyAssignIP->isSet() &&
sym.test(Fortran::semantics::Symbol::Flag::OmpLastPrivate)) {
Expand Down Expand Up @@ -1092,6 +1040,79 @@ class FirConverter : public Fortran::lower::AbstractConverter {
return true;
}

void copyVar(const Fortran::semantics::Symbol &sym,
const Fortran::lower::SymbolBox &lhs_sb,
const Fortran::lower::SymbolBox &rhs_sb) {
mlir::Location loc = genLocation(sym.name());
if (lowerToHighLevelFIR())
copyVarHLFIR(loc, lhs_sb.getAddr(), rhs_sb.getAddr());
else
copyVarFIR(loc, sym, lhs_sb, rhs_sb);
}

void copyVarHLFIR(mlir::Location loc, mlir::Value dst, mlir::Value src) {
assert(lowerToHighLevelFIR());
hlfir::Entity lhs{dst};
hlfir::Entity rhs{src};
// Temporary_lhs is set to true in hlfir.assign below to avoid user
// assignment to be used and finalization to be called on the LHS.
// This may or may not be correct but mimics the current behaviour
// without HLFIR.
auto copyData = [&](hlfir::Entity l, hlfir::Entity r) {
// Dereference RHS and load it if trivial scalar.
r = hlfir::loadTrivialScalar(loc, *builder, r);
builder->create<hlfir::AssignOp>(
loc, r, l,
/*isWholeAllocatableAssignment=*/false,
/*keepLhsLengthInAllocatableAssignment=*/false,
/*temporary_lhs=*/true);
};
if (lhs.isAllocatable()) {
// Deep copy allocatable if it is allocated.
// Note that when allocated, the RHS is already allocated with the LHS
// shape for copy on entry in createHostAssociateVarClone.
// For lastprivate, this assumes that the RHS was not reallocated in
// the OpenMP region.
lhs = hlfir::derefPointersAndAllocatables(loc, *builder, lhs);
mlir::Value addr = hlfir::genVariableRawAddress(loc, *builder, lhs);
mlir::Value isAllocated = builder->genIsNotNullAddr(loc, addr);
builder->genIfThen(loc, isAllocated)
.genThen([&]() {
// Copy the DATA, not the descriptors.
copyData(lhs, rhs);
})
.end();
} else if (lhs.isPointer()) {
// Set LHS target to the target of RHS (do not copy the RHS
// target data into the LHS target storage).
auto loadVal = builder->create<fir::LoadOp>(loc, rhs);
builder->create<fir::StoreOp>(loc, loadVal, lhs);
} else {
// Non ALLOCATABLE/POINTER variable. Simple DATA copy.
copyData(lhs, rhs);
}
}

void copyVarFIR(mlir::Location loc, const Fortran::semantics::Symbol &sym,
const Fortran::lower::SymbolBox &lhs_sb,
const Fortran::lower::SymbolBox &rhs_sb) {
assert(!lowerToHighLevelFIR());
fir::ExtendedValue lhs = symBoxToExtendedValue(lhs_sb);
fir::ExtendedValue rhs = symBoxToExtendedValue(rhs_sb);
mlir::Type symType = genType(sym);
if (auto seqTy = symType.dyn_cast<fir::SequenceType>()) {
Fortran::lower::StatementContext stmtCtx;
Fortran::lower::createSomeArrayAssignment(*this, lhs, rhs, localSymbols,
stmtCtx);
stmtCtx.finalizeAndReset();
} else if (lhs.getBoxOf<fir::CharBoxValue>()) {
fir::factory::CharacterExprHelper{*builder, loc}.createAssign(lhs, rhs);
} else {
auto loadVal = builder->create<fir::LoadOp>(loc, fir::getBase(rhs));
builder->create<fir::StoreOp>(loc, loadVal, fir::getBase(lhs));
}
}

/// Map a block argument to a result or dummy symbol. This is not the
/// definitive mapping. The specification expression have not been lowered
/// yet. The final mapping will be done using this pre-mapping in
Expand Down
175 changes: 170 additions & 5 deletions flang/lib/Lower/OpenMP.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#include "flang/Optimizer/Builder/BoxValue.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Builder/Todo.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/HLFIR/HLFIROps.h"
#include "flang/Parser/dump-parse-tree.h"
#include "flang/Parser/parse-tree.h"
Expand Down Expand Up @@ -592,6 +593,10 @@ class ClauseProcessor {
processAllocate(llvm::SmallVectorImpl<mlir::Value> &allocatorOperands,
llvm::SmallVectorImpl<mlir::Value> &allocateOperands) const;
bool processCopyin() const;
bool processCopyPrivate(
mlir::Location currentLocation,
llvm::SmallVectorImpl<mlir::Value> &copyPrivateVars,
llvm::SmallVectorImpl<mlir::Attribute> &copyPrivateFuncs) const;
bool processDepend(llvm::SmallVectorImpl<mlir::Attribute> &dependTypeOperands,
llvm::SmallVectorImpl<mlir::Value> &dependOperands) const;
bool
Expand Down Expand Up @@ -1160,6 +1165,102 @@ class ReductionProcessor {
}
};

/// Class that extracts information from the specified type.
class TypeInfo {
public:
TypeInfo(mlir::Type ty) { typeScan(ty); }

// Returns the length of character types.
std::optional<fir::CharacterType::LenType> getCharLength() const {
return charLen;
}

// Returns the shape of array types.
const llvm::SmallVector<int64_t> &getShape() const { return shape; }

// Is the type inside a box?
bool isBox() const { return inBox; }

private:
void typeScan(mlir::Type type);

std::optional<fir::CharacterType::LenType> charLen;
llvm::SmallVector<int64_t> shape;
bool inBox = false;
};

void TypeInfo::typeScan(mlir::Type ty) {
if (auto sty = mlir::dyn_cast<fir::SequenceType>(ty)) {
assert(shape.empty() && !sty.getShape().empty());
shape = llvm::SmallVector<int64_t>(sty.getShape());
typeScan(sty.getEleTy());
} else if (auto bty = mlir::dyn_cast<fir::BoxType>(ty)) {
inBox = true;
typeScan(bty.getEleTy());
} else if (auto cty = mlir::dyn_cast<fir::CharacterType>(ty)) {
charLen = cty.getLen();
} else if (auto hty = mlir::dyn_cast<fir::HeapType>(ty)) {
typeScan(hty.getEleTy());
} else if (auto pty = mlir::dyn_cast<fir::PointerType>(ty)) {
typeScan(pty.getEleTy());
}
}

// Create a function that performs a copy between two variables, compatible
// with their types and attributes.
static mlir::func::FuncOp
createCopyFunc(mlir::Location loc, Fortran::lower::AbstractConverter &converter,
mlir::Type varType, fir::FortranVariableFlagsEnum varAttrs) {
fir::FirOpBuilder &builder = converter.getFirOpBuilder();
mlir::ModuleOp module = builder.getModule();
mlir::Type eleTy = mlir::cast<fir::ReferenceType>(varType).getEleTy();
TypeInfo typeInfo(eleTy);
std::string copyFuncName =
fir::getTypeAsString(eleTy, builder.getKindMap(), "_copy");

if (auto decl = module.lookupSymbol<mlir::func::FuncOp>(copyFuncName))
return decl;

// create function
mlir::OpBuilder::InsertionGuard guard(builder);
mlir::OpBuilder modBuilder(module.getBodyRegion());
llvm::SmallVector<mlir::Type> argsTy = {varType, varType};
auto funcType = mlir::FunctionType::get(builder.getContext(), argsTy, {});
mlir::func::FuncOp funcOp =
modBuilder.create<mlir::func::FuncOp>(loc, copyFuncName, funcType);
funcOp.setVisibility(mlir::SymbolTable::Visibility::Private);
builder.createBlock(&funcOp.getRegion(), funcOp.getRegion().end(), argsTy,
{loc, loc});
builder.setInsertionPointToStart(&funcOp.getRegion().back());
// generate body
fir::FortranVariableFlagsAttr attrs;
if (varAttrs != fir::FortranVariableFlagsEnum::None)
attrs = fir::FortranVariableFlagsAttr::get(builder.getContext(), varAttrs);
llvm::SmallVector<mlir::Value> typeparams;
if (typeInfo.getCharLength().has_value()) {
mlir::Value charLen = builder.createIntegerConstant(
loc, builder.getCharacterLengthType(), *typeInfo.getCharLength());
typeparams.push_back(charLen);
}
mlir::Value shape;
if (!typeInfo.isBox() && !typeInfo.getShape().empty()) {
llvm::SmallVector<mlir::Value> extents;
for (auto extent : typeInfo.getShape())
extents.push_back(
builder.createIntegerConstant(loc, builder.getIndexType(), extent));
shape = builder.create<fir::ShapeOp>(loc, extents);
}
auto declDst = builder.create<hlfir::DeclareOp>(loc, funcOp.getArgument(0),
copyFuncName + "_dst", shape,
typeparams, attrs);
auto declSrc = builder.create<hlfir::DeclareOp>(loc, funcOp.getArgument(1),
copyFuncName + "_src", shape,
typeparams, attrs);
converter.copyVar(loc, declDst.getBase(), declSrc.getBase());
builder.create<mlir::func::ReturnOp>(loc);
return funcOp;
}

static mlir::omp::ScheduleModifier
translateScheduleModifier(const Fortran::parser::OmpScheduleModifierType &m) {
switch (m.v) {
Expand Down Expand Up @@ -1740,6 +1841,62 @@ bool ClauseProcessor::processCopyin() const {
return hasCopyin;
}

bool ClauseProcessor::processCopyPrivate(
mlir::Location currentLocation,
llvm::SmallVectorImpl<mlir::Value> &copyPrivateVars,
llvm::SmallVectorImpl<mlir::Attribute> &copyPrivateFuncs) const {
auto addCopyPrivateVar = [&](Fortran::semantics::Symbol *sym) {
mlir::Value symVal = converter.getSymbolAddress(*sym);
auto declOp = symVal.getDefiningOp<hlfir::DeclareOp>();
if (!declOp)
fir::emitFatalError(currentLocation,
"COPYPRIVATE is supported only in HLFIR mode");
symVal = declOp.getBase();
mlir::Type symType = symVal.getType();
fir::FortranVariableFlagsEnum attrs =
declOp.getFortranAttrs().has_value()
? *declOp.getFortranAttrs()
: fir::FortranVariableFlagsEnum::None;
mlir::Value cpVar = symVal;

// CopyPrivate variables must be passed by reference. However, in the case
// of assumed shapes/vla the type is not a !fir.ref, but a !fir.box.
// In these cases to retrieve the appropriate !fir.ref<!fir.box<...>> to
// access the data we need we must perform an alloca and then store to it
// and retrieve the data from the new alloca.
if (mlir::isa<fir::BaseBoxType>(symType)) {
fir::FirOpBuilder &builder = converter.getFirOpBuilder();
auto alloca = builder.create<fir::AllocaOp>(currentLocation, symType);
builder.create<fir::StoreOp>(currentLocation, symVal, alloca);
cpVar = alloca;
}

copyPrivateVars.push_back(cpVar);
mlir::func::FuncOp funcOp =
createCopyFunc(currentLocation, converter, cpVar.getType(), attrs);
copyPrivateFuncs.push_back(mlir::SymbolRefAttr::get(funcOp));
};

bool hasCopyPrivate = findRepeatableClause<ClauseTy::Copyprivate>(
[&](const ClauseTy::Copyprivate *copyPrivateClause,
const Fortran::parser::CharBlock &) {
const Fortran::parser::OmpObjectList &ompObjectList =
copyPrivateClause->v;
for (const Fortran::parser::OmpObject &ompObject : ompObjectList.v) {
Fortran::semantics::Symbol *sym = getOmpObjectSymbol(ompObject);
if (const auto *commonDetails =
sym->detailsIf<Fortran::semantics::CommonBlockDetails>()) {
for (const auto &mem : commonDetails->objects())
addCopyPrivateVar(&*mem);
break;
}
addCopyPrivateVar(sym);
}
});

return hasCopyPrivate;
}

bool ClauseProcessor::processDepend(
llvm::SmallVectorImpl<mlir::Attribute> &dependTypeOperands,
llvm::SmallVectorImpl<mlir::Value> &dependOperands) const {
Expand Down Expand Up @@ -2481,19 +2638,26 @@ genSingleOp(Fortran::lower::AbstractConverter &converter,
const Fortran::parser::OmpClauseList &beginClauseList,
const Fortran::parser::OmpClauseList &endClauseList) {
llvm::SmallVector<mlir::Value> allocateOperands, allocatorOperands;
llvm::SmallVector<mlir::Value> copyPrivateVars;
llvm::SmallVector<mlir::Attribute> copyPrivateFuncs;
mlir::UnitAttr nowaitAttr;

ClauseProcessor cp(converter, beginClauseList);
cp.processAllocate(allocatorOperands, allocateOperands);
cp.processTODO<Fortran::parser::OmpClause::Copyprivate>(
currentLocation, llvm::omp::Directive::OMPD_single);

ClauseProcessor(converter, endClauseList).processNowait(nowaitAttr);
ClauseProcessor ecp(converter, endClauseList);
ecp.processNowait(nowaitAttr);
ecp.processCopyPrivate(currentLocation, copyPrivateVars, copyPrivateFuncs);

return genOpWithBody<mlir::omp::SingleOp>(
converter, eval, genNested, currentLocation,
/*outerCombined=*/false, &beginClauseList, allocateOperands,
allocatorOperands, nowaitAttr);
allocatorOperands, copyPrivateVars,
copyPrivateFuncs.empty()
? nullptr
: mlir::ArrayAttr::get(converter.getFirOpBuilder().getContext(),
copyPrivateFuncs),
nowaitAttr);
}

static mlir::omp::TaskOp
Expand Down Expand Up @@ -3367,7 +3531,8 @@ genOMP(Fortran::lower::AbstractConverter &converter,

for (const auto &clause : endClauseList.v) {
mlir::Location clauseLocation = converter.genLocation(clause.source);
if (!std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u))
if (!std::get_if<Fortran::parser::OmpClause::Nowait>(&clause.u) &&
!std::get_if<Fortran::parser::OmpClause::Copyprivate>(&clause.u))
TODO(clauseLocation, "OpenMP Block construct clause");
}

Expand Down
Loading