Skip to content

Commit 0d30968

Browse files
ergawyergawy
committed
[flang][fir] Basic PFT to MLIR lowering for do concurrent locality specifiers
Extends support for `fir.do_concurrent` locality specifiers to the PFT to MLIR level. This adds code-gen for generating the newly added `fir.local` ops and referencing these ops from `fir.do_concurrent.loop` ops that have locality specifiers attached to them. This reuses the `DataSharingProcessor` component and generalizes it a bit more to allow for handling `omp.private` ops and `fir.local` ops as well.
1 parent a4acb56 commit 0d30968

File tree

7 files changed

+209
-40
lines changed

7 files changed

+209
-40
lines changed

‎flang/include/flang/Lower/AbstractConverter.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,10 @@ class AbstractConverter {
348348
virtual Fortran::lower::SymbolBox
349349
lookupOneLevelUpSymbol(const Fortran::semantics::Symbol &sym) = 0;
350350

351+
/// Find the symbol in the inner-most level of the local map or return null.
352+
virtual Fortran::lower::SymbolBox
353+
shallowLookupSymbol(const Fortran::semantics::Symbol &sym) = 0;
354+
351355
/// Return the mlir::SymbolTable associated to the ModuleOp.
352356
/// Look-ups are faster using it than using module.lookup<>,
353357
/// but the module op should be queried in case of failure

‎flang/include/flang/Optimizer/Dialect/FIROps.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,10 @@ class CoordinateIndicesAdaptor {
147147
mlir::ValueRange values;
148148
};
149149

150+
struct LocalitySpecifierOperands {
151+
llvm::SmallVector<::mlir::Value> privateVars;
152+
llvm::SmallVector<::mlir::Attribute> privateSyms;
153+
};
150154
} // namespace fir
151155

152156
#endif // FORTRAN_OPTIMIZER_DIALECT_FIROPS_H

‎flang/include/flang/Optimizer/Dialect/FIROps.td‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3600,6 +3600,21 @@ def fir_LocalitySpecifierOp : fir_Op<"local", [IsolatedFromAbove]> {
36003600
];
36013601

36023602
let extraClassDeclaration = [{
3603+
mlir::BlockArgument getInitMoldArg() {
3604+
auto &region = getInitRegion();
3605+
return region.empty() ? nullptr : region.getArgument(0);
3606+
}
3607+
mlir::BlockArgument getInitPrivateArg() {
3608+
auto &region = getInitRegion();
3609+
return region.empty() ? nullptr : region.getArgument(1);
3610+
}
3611+
3612+
/// Returns true if the init region might read from the mold argument
3613+
bool initReadsFromMold() {
3614+
mlir::BlockArgument moldArg = getInitMoldArg();
3615+
return moldArg && !moldArg.use_empty();
3616+
}
3617+
36033618
/// Get the type for arguments to nested regions. This should
36043619
/// generally be either the same as getType() or some pointer
36053620
/// type (pointing to the type allocated by this op).

‎flang/lib/Lower/Bridge.cpp‎

Lines changed: 50 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212

1313
#include "flang/Lower/Bridge.h"
1414

15+
#include "OpenMP/DataSharingProcessor.h"
16+
#include "OpenMP/Utils.h"
1517
#include "flang/Lower/Allocatable.h"
1618
#include "flang/Lower/CallInterface.h"
1719
#include "flang/Lower/Coarray.h"
@@ -1144,6 +1146,14 @@ class FirConverter : public Fortran::lower::AbstractConverter {
11441146
return name;
11451147
}
11461148

1149+
/// Find the symbol in the inner-most level of the local map or return null.
1150+
Fortran::lower::SymbolBox
1151+
shallowLookupSymbol(const Fortran::semantics::Symbol &sym) override {
1152+
if (Fortran::lower::SymbolBox v = localSymbols.shallowLookupSymbol(sym))
1153+
return v;
1154+
return {};
1155+
}
1156+
11471157
private:
11481158
FirConverter() = delete;
11491159
FirConverter(const FirConverter &) = delete;
@@ -1218,14 +1228,6 @@ class FirConverter : public Fortran::lower::AbstractConverter {
12181228
return {};
12191229
}
12201230

1221-
/// Find the symbol in the inner-most level of the local map or return null.
1222-
Fortran::lower::SymbolBox
1223-
shallowLookupSymbol(const Fortran::semantics::Symbol &sym) {
1224-
if (Fortran::lower::SymbolBox v = localSymbols.shallowLookupSymbol(sym))
1225-
return v;
1226-
return {};
1227-
}
1228-
12291231
/// Find the symbol in one level up of symbol map such as for host-association
12301232
/// in OpenMP code or return null.
12311233
Fortran::lower::SymbolBox
@@ -2029,9 +2031,30 @@ class FirConverter : public Fortran::lower::AbstractConverter {
20292031
void handleLocalitySpecs(const IncrementLoopInfo &info) {
20302032
Fortran::semantics::SemanticsContext &semanticsContext =
20312033
bridge.getSemanticsContext();
2032-
for (const Fortran::semantics::Symbol *sym : info.localSymList)
2034+
Fortran::lower::omp::DataSharingProcessor dsp(
2035+
*this, semanticsContext, getEval(),
2036+
/*useDelayedPrivatization=*/true, localSymbols);
2037+
fir::LocalitySpecifierOperands privateClauseOps;
2038+
auto doConcurrentLoopOp =
2039+
mlir::dyn_cast_if_present<fir::DoConcurrentLoopOp>(info.loopOp);
2040+
bool useDelayedPriv =
2041+
enableDelayedPrivatizationStaging && doConcurrentLoopOp;
2042+
2043+
for (const Fortran::semantics::Symbol *sym : info.localSymList) {
2044+
if (useDelayedPriv) {
2045+
dsp.privatizeSymbol<fir::LocalitySpecifierOp>(sym, &privateClauseOps);
2046+
continue;
2047+
}
2048+
20332049
createHostAssociateVarClone(*sym, /*skipDefaultInit=*/false);
2050+
}
2051+
20342052
for (const Fortran::semantics::Symbol *sym : info.localInitSymList) {
2053+
if (useDelayedPriv) {
2054+
dsp.privatizeSymbol<fir::LocalitySpecifierOp>(sym, &privateClauseOps);
2055+
continue;
2056+
}
2057+
20352058
createHostAssociateVarClone(*sym, /*skipDefaultInit=*/true);
20362059
const auto *hostDetails =
20372060
sym->detailsIf<Fortran::semantics::HostAssocDetails>();
@@ -2050,6 +2073,24 @@ class FirConverter : public Fortran::lower::AbstractConverter {
20502073
sym->detailsIf<Fortran::semantics::HostAssocDetails>();
20512074
copySymbolBinding(hostDetails->symbol(), *sym);
20522075
}
2076+
2077+
if (useDelayedPriv) {
2078+
doConcurrentLoopOp.getLocalVarsMutable().assign(
2079+
privateClauseOps.privateVars);
2080+
doConcurrentLoopOp.setLocalSymsAttr(
2081+
builder->getArrayAttr(privateClauseOps.privateSyms));
2082+
2083+
for (auto [sym, privateVar] : llvm::zip_equal(
2084+
dsp.getAllSymbolsToPrivatize(), privateClauseOps.privateVars)) {
2085+
auto arg = doConcurrentLoopOp.getRegion().begin()->addArgument(
2086+
privateVar.getType(), doConcurrentLoopOp.getLoc());
2087+
bindSymbol(*sym, hlfir::translateToExtendedValue(
2088+
privateVar.getLoc(), *builder, hlfir::Entity{arg},
2089+
/*contiguousHint=*/true)
2090+
.first);
2091+
}
2092+
}
2093+
20532094
// Note that allocatable, types with ultimate components, and type
20542095
// requiring finalization are forbidden in LOCAL/LOCAL_INIT (F2023 C1130),
20552096
// so no clean-up needs to be generated for these entities.

‎flang/lib/Lower/OpenMP/DataSharingProcessor.cpp‎

Lines changed: 77 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "flang/Optimizer/Builder/BoxValue.h"
2121
#include "flang/Optimizer/Builder/HLFIRTools.h"
2222
#include "flang/Optimizer/Builder/Todo.h"
23+
#include "flang/Optimizer/Dialect/FIROps.h"
2324
#include "flang/Optimizer/HLFIR/HLFIRDialect.h"
2425
#include "flang/Optimizer/HLFIR/HLFIROps.h"
2526
#include "flang/Semantics/attr.h"
@@ -53,6 +54,15 @@ DataSharingProcessor::DataSharingProcessor(
5354
});
5455
}
5556

57+
DataSharingProcessor::DataSharingProcessor(lower::AbstractConverter &converter,
58+
semantics::SemanticsContext &semaCtx,
59+
lower::pft::Evaluation &eval,
60+
bool useDelayedPrivatization,
61+
lower::SymMap &symTable)
62+
: DataSharingProcessor(converter, semaCtx, {}, eval,
63+
/*shouldCollectPreDeterminedSymols=*/false,
64+
useDelayedPrivatization, symTable) {}
65+
5666
void DataSharingProcessor::processStep1(
5767
mlir::omp::PrivateClauseOps *clauseOps) {
5868
collectSymbolsForPrivatization();
@@ -172,7 +182,8 @@ void DataSharingProcessor::cloneSymbol(const semantics::Symbol *sym) {
172182

173183
void DataSharingProcessor::copyFirstPrivateSymbol(
174184
const semantics::Symbol *sym, mlir::OpBuilder::InsertPoint *copyAssignIP) {
175-
if (sym->test(semantics::Symbol::Flag::OmpFirstPrivate))
185+
if (sym->test(semantics::Symbol::Flag::OmpFirstPrivate) ||
186+
sym->test(semantics::Symbol::Flag::LocalityLocalInit))
176187
converter.copyHostAssociateVar(*sym, copyAssignIP);
177188
}
178189

@@ -485,9 +496,9 @@ void DataSharingProcessor::privatize(mlir::omp::PrivateClauseOps *clauseOps) {
485496
if (const auto *commonDet =
486497
sym->detailsIf<semantics::CommonBlockDetails>()) {
487498
for (const auto &mem : commonDet->objects())
488-
doPrivatize(&*mem, clauseOps);
499+
privatizeSymbol(&*mem, clauseOps);
489500
} else
490-
doPrivatize(sym, clauseOps);
501+
privatizeSymbol(sym, clauseOps);
491502
}
492503
}
493504

@@ -504,22 +515,30 @@ void DataSharingProcessor::copyLastPrivatize(mlir::Operation *op) {
504515
}
505516
}
506517

507-
void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
508-
mlir::omp::PrivateClauseOps *clauseOps) {
518+
template <typename OpType, typename OperandsStructType>
519+
void DataSharingProcessor::privatizeSymbol(
520+
const semantics::Symbol *symToPrivatize, OperandsStructType *clauseOps) {
509521
if (!useDelayedPrivatization) {
510-
cloneSymbol(sym);
511-
copyFirstPrivateSymbol(sym);
522+
cloneSymbol(symToPrivatize);
523+
copyFirstPrivateSymbol(symToPrivatize);
512524
return;
513525
}
514526

515-
lower::SymbolBox hsb = converter.lookupOneLevelUpSymbol(*sym);
527+
const semantics::Symbol *sym = symToPrivatize->HasLocalLocality()
528+
? &symToPrivatize->GetUltimate()
529+
: symToPrivatize;
530+
lower::SymbolBox hsb = symToPrivatize->HasLocalLocality()
531+
? converter.shallowLookupSymbol(*sym)
532+
: converter.lookupOneLevelUpSymbol(*sym);
516533
assert(hsb && "Host symbol box not found");
517534
hlfir::Entity entity{hsb.getAddr()};
518535
bool cannotHaveNonDefaultLowerBounds = !entity.mayHaveNonDefaultLowerBounds();
519536

520537
mlir::Location symLoc = hsb.getAddr().getLoc();
521538
std::string privatizerName = sym->name().ToString() + ".privatizer";
522-
bool isFirstPrivate = sym->test(semantics::Symbol::Flag::OmpFirstPrivate);
539+
bool isFirstPrivate =
540+
symToPrivatize->test(semantics::Symbol::Flag::OmpFirstPrivate) ||
541+
symToPrivatize->test(semantics::Symbol::Flag::LocalityLocalInit);
523542

524543
mlir::Value privVal = hsb.getAddr();
525544
mlir::Type allocType = privVal.getType();
@@ -553,24 +572,33 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
553572

554573
mlir::Type argType = privVal.getType();
555574

556-
mlir::omp::PrivateClauseOp privatizerOp = [&]() {
575+
OpType privatizerOp = [&]() {
557576
auto moduleOp = firOpBuilder.getModule();
558577
auto uniquePrivatizerName = fir::getTypeAsString(
559578
allocType, converter.getKindMap(),
560579
converter.mangleName(*sym) +
561580
(isFirstPrivate ? "_firstprivate" : "_private"));
562581

563582
if (auto existingPrivatizer =
564-
moduleOp.lookupSymbol<mlir::omp::PrivateClauseOp>(
565-
uniquePrivatizerName))
583+
moduleOp.lookupSymbol<OpType>(uniquePrivatizerName))
566584
return existingPrivatizer;
567585

568586
mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
569587
firOpBuilder.setInsertionPointToStart(moduleOp.getBody());
570-
auto result = firOpBuilder.create<mlir::omp::PrivateClauseOp>(
571-
symLoc, uniquePrivatizerName, allocType,
572-
isFirstPrivate ? mlir::omp::DataSharingClauseType::FirstPrivate
573-
: mlir::omp::DataSharingClauseType::Private);
588+
OpType result;
589+
590+
if constexpr (std::is_same_v<OpType, mlir::omp::PrivateClauseOp>) {
591+
result = firOpBuilder.create<OpType>(
592+
symLoc, uniquePrivatizerName, allocType,
593+
isFirstPrivate ? mlir::omp::DataSharingClauseType::FirstPrivate
594+
: mlir::omp::DataSharingClauseType::Private);
595+
} else {
596+
result = firOpBuilder.create<OpType>(
597+
symLoc, uniquePrivatizerName, allocType,
598+
isFirstPrivate ? fir::LocalitySpecifierType::LocalInit
599+
: fir::LocalitySpecifierType::Local);
600+
}
601+
574602
fir::ExtendedValue symExV = converter.getSymbolExtendedValue(*sym);
575603
lower::SymMapScope outerScope(symTable);
576604

@@ -613,27 +641,36 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
613641
&copyRegion, /*insertPt=*/{}, {argType, argType}, {symLoc, symLoc});
614642
firOpBuilder.setInsertionPointToEnd(copyEntryBlock);
615643

616-
auto addSymbol = [&](unsigned argIdx, bool force = false) {
644+
auto addSymbol = [&](unsigned argIdx, const semantics::Symbol *symToMap,
645+
bool force = false) {
617646
symExV.match(
618647
[&](const fir::MutableBoxValue &box) {
619648
symTable.addSymbol(
620-
*sym, fir::substBase(box, copyRegion.getArgument(argIdx)),
621-
force);
649+
*symToMap,
650+
fir::substBase(box, copyRegion.getArgument(argIdx)), force);
622651
},
623652
[&](const auto &box) {
624-
symTable.addSymbol(*sym, copyRegion.getArgument(argIdx), force);
653+
symTable.addSymbol(*symToMap, copyRegion.getArgument(argIdx),
654+
force);
625655
});
626656
};
627657

628-
addSymbol(0, true);
658+
addSymbol(0, sym, true);
629659
lower::SymMapScope innerScope(symTable);
630-
addSymbol(1);
660+
addSymbol(1, symToPrivatize);
631661

632662
auto ip = firOpBuilder.saveInsertionPoint();
633-
copyFirstPrivateSymbol(sym, &ip);
634-
635-
firOpBuilder.create<mlir::omp::YieldOp>(
636-
hsb.getAddr().getLoc(), symTable.shallowLookupSymbol(*sym).getAddr());
663+
copyFirstPrivateSymbol(symToPrivatize, &ip);
664+
665+
if constexpr (std::is_same_v<OpType, mlir::omp::PrivateClauseOp>) {
666+
firOpBuilder.create<mlir::omp::YieldOp>(
667+
hsb.getAddr().getLoc(),
668+
symTable.shallowLookupSymbol(*symToPrivatize).getAddr());
669+
} else {
670+
firOpBuilder.create<fir::YieldOp>(
671+
hsb.getAddr().getLoc(),
672+
symTable.shallowLookupSymbol(*symToPrivatize).getAddr());
673+
}
637674
}
638675

639676
return result;
@@ -644,9 +681,22 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
644681
clauseOps->privateVars.push_back(privVal);
645682
}
646683

647-
symToPrivatizer[sym] = privatizerOp;
684+
if (symToPrivatize->HasLocalLocality())
685+
allPrivatizedSymbols.insert(symToPrivatize);
648686
}
649687

688+
template void
689+
DataSharingProcessor::privatizeSymbol<mlir::omp::PrivateClauseOp,
690+
mlir::omp::PrivateClauseOps>(
691+
const semantics::Symbol *symToPrivatize,
692+
mlir::omp::PrivateClauseOps *clauseOps);
693+
694+
template void
695+
DataSharingProcessor::privatizeSymbol<fir::LocalitySpecifierOp,
696+
fir::LocalitySpecifierOperands>(
697+
const semantics::Symbol *symToPrivatize,
698+
fir::LocalitySpecifierOperands *clauseOps);
699+
650700
} // namespace omp
651701
} // namespace lower
652702
} // namespace Fortran

‎flang/lib/Lower/OpenMP/DataSharingProcessor.h‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,6 @@ class DataSharingProcessor {
7777
llvm::SetVector<const semantics::Symbol *> preDeterminedSymbols;
7878
llvm::SetVector<const semantics::Symbol *> allPrivatizedSymbols;
7979

80-
llvm::DenseMap<const semantics::Symbol *, mlir::omp::PrivateClauseOp>
81-
symToPrivatizer;
8280
lower::AbstractConverter &converter;
8381
semantics::SemanticsContext &semaCtx;
8482
fir::FirOpBuilder &firOpBuilder;
@@ -105,8 +103,6 @@ class DataSharingProcessor {
105103
void collectImplicitSymbols();
106104
void collectPreDeterminedSymbols();
107105
void privatize(mlir::omp::PrivateClauseOps *clauseOps);
108-
void doPrivatize(const semantics::Symbol *sym,
109-
mlir::omp::PrivateClauseOps *clauseOps);
110106
void copyLastPrivatize(mlir::Operation *op);
111107
void insertLastPrivateCompare(mlir::Operation *op);
112108
void cloneSymbol(const semantics::Symbol *sym);
@@ -125,6 +121,11 @@ class DataSharingProcessor {
125121
bool shouldCollectPreDeterminedSymbols,
126122
bool useDelayedPrivatization, lower::SymMap &symTable);
127123

124+
DataSharingProcessor(lower::AbstractConverter &converter,
125+
semantics::SemanticsContext &semaCtx,
126+
lower::pft::Evaluation &eval,
127+
bool useDelayedPrivatization, lower::SymMap &symTable);
128+
128129
// Privatisation is split into two steps.
129130
// Step1 performs cloning of all privatisation clauses and copying for
130131
// firstprivates. Step1 is performed at the place where process/processStep1
@@ -151,6 +152,11 @@ class DataSharingProcessor {
151152
? allPrivatizedSymbols.getArrayRef()
152153
: llvm::ArrayRef<const semantics::Symbol *>();
153154
}
155+
156+
template <typename OpType = mlir::omp::PrivateClauseOp,
157+
typename OperandsStructType = mlir::omp::PrivateClauseOps>
158+
void privatizeSymbol(const semantics::Symbol *symToPrivatize,
159+
OperandsStructType *clauseOps);
154160
};
155161

156162
} // namespace omp

0 commit comments

Comments
 (0)