Skip to content

Commit b6893b6

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 5855225 commit b6893b6

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();
@@ -174,7 +184,8 @@ void DataSharingProcessor::cloneSymbol(const semantics::Symbol *sym) {
174184

175185
void DataSharingProcessor::copyFirstPrivateSymbol(
176186
const semantics::Symbol *sym, mlir::OpBuilder::InsertPoint *copyAssignIP) {
177-
if (sym->test(semantics::Symbol::Flag::OmpFirstPrivate))
187+
if (sym->test(semantics::Symbol::Flag::OmpFirstPrivate) ||
188+
sym->test(semantics::Symbol::Flag::LocalityLocalInit))
178189
converter.copyHostAssociateVar(*sym, copyAssignIP);
179190
}
180191

@@ -487,9 +498,9 @@ void DataSharingProcessor::privatize(mlir::omp::PrivateClauseOps *clauseOps) {
487498
if (const auto *commonDet =
488499
sym->detailsIf<semantics::CommonBlockDetails>()) {
489500
for (const auto &mem : commonDet->objects())
490-
doPrivatize(&*mem, clauseOps);
501+
privatizeSymbol(&*mem, clauseOps);
491502
} else
492-
doPrivatize(sym, clauseOps);
503+
privatizeSymbol(sym, clauseOps);
493504
}
494505
}
495506

@@ -506,22 +517,30 @@ void DataSharingProcessor::copyLastPrivatize(mlir::Operation *op) {
506517
}
507518
}
508519

509-
void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
510-
mlir::omp::PrivateClauseOps *clauseOps) {
520+
template <typename OpType, typename OperandsStructType>
521+
void DataSharingProcessor::privatizeSymbol(
522+
const semantics::Symbol *symToPrivatize, OperandsStructType *clauseOps) {
511523
if (!useDelayedPrivatization) {
512-
cloneSymbol(sym);
513-
copyFirstPrivateSymbol(sym);
524+
cloneSymbol(symToPrivatize);
525+
copyFirstPrivateSymbol(symToPrivatize);
514526
return;
515527
}
516528

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

522539
mlir::Location symLoc = hsb.getAddr().getLoc();
523540
std::string privatizerName = sym->name().ToString() + ".privatizer";
524-
bool isFirstPrivate = sym->test(semantics::Symbol::Flag::OmpFirstPrivate);
541+
bool isFirstPrivate =
542+
symToPrivatize->test(semantics::Symbol::Flag::OmpFirstPrivate) ||
543+
symToPrivatize->test(semantics::Symbol::Flag::LocalityLocalInit);
525544

526545
mlir::Value privVal = hsb.getAddr();
527546
mlir::Type allocType = privVal.getType();
@@ -555,24 +574,33 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
555574

556575
mlir::Type argType = privVal.getType();
557576

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

565584
if (auto existingPrivatizer =
566-
moduleOp.lookupSymbol<mlir::omp::PrivateClauseOp>(
567-
uniquePrivatizerName))
585+
moduleOp.lookupSymbol<OpType>(uniquePrivatizerName))
568586
return existingPrivatizer;
569587

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

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

618-
auto addSymbol = [&](unsigned argIdx, bool force = false) {
646+
auto addSymbol = [&](unsigned argIdx, const semantics::Symbol *symToMap,
647+
bool force = false) {
619648
symExV.match(
620649
[&](const fir::MutableBoxValue &box) {
621650
symTable.addSymbol(
622-
*sym, fir::substBase(box, copyRegion.getArgument(argIdx)),
623-
force);
651+
*symToMap,
652+
fir::substBase(box, copyRegion.getArgument(argIdx)), force);
624653
},
625654
[&](const auto &box) {
626-
symTable.addSymbol(*sym, copyRegion.getArgument(argIdx), force);
655+
symTable.addSymbol(*symToMap, copyRegion.getArgument(argIdx),
656+
force);
627657
});
628658
};
629659

630-
addSymbol(0, true);
660+
addSymbol(0, sym, true);
631661
lower::SymMapScope innerScope(symTable);
632-
addSymbol(1);
662+
addSymbol(1, symToPrivatize);
633663

634664
auto ip = firOpBuilder.saveInsertionPoint();
635-
copyFirstPrivateSymbol(sym, &ip);
636-
637-
firOpBuilder.create<mlir::omp::YieldOp>(
638-
hsb.getAddr().getLoc(), symTable.shallowLookupSymbol(*sym).getAddr());
665+
copyFirstPrivateSymbol(symToPrivatize, &ip);
666+
667+
if constexpr (std::is_same_v<OpType, mlir::omp::PrivateClauseOp>) {
668+
firOpBuilder.create<mlir::omp::YieldOp>(
669+
hsb.getAddr().getLoc(),
670+
symTable.shallowLookupSymbol(*symToPrivatize).getAddr());
671+
} else {
672+
firOpBuilder.create<fir::YieldOp>(
673+
hsb.getAddr().getLoc(),
674+
symTable.shallowLookupSymbol(*symToPrivatize).getAddr());
675+
}
639676
}
640677

641678
return result;
@@ -646,9 +683,22 @@ void DataSharingProcessor::doPrivatize(const semantics::Symbol *sym,
646683
clauseOps->privateVars.push_back(privVal);
647684
}
648685

649-
symToPrivatizer[sym] = privatizerOp;
686+
if (symToPrivatize->HasLocalLocality())
687+
allPrivatizedSymbols.insert(symToPrivatize);
650688
}
651689

690+
template void
691+
DataSharingProcessor::privatizeSymbol<mlir::omp::PrivateClauseOp,
692+
mlir::omp::PrivateClauseOps>(
693+
const semantics::Symbol *symToPrivatize,
694+
mlir::omp::PrivateClauseOps *clauseOps);
695+
696+
template void
697+
DataSharingProcessor::privatizeSymbol<fir::LocalitySpecifierOp,
698+
fir::LocalitySpecifierOperands>(
699+
const semantics::Symbol *symToPrivatize,
700+
fir::LocalitySpecifierOperands *clauseOps);
701+
652702
} // namespace omp
653703
} // namespace lower
654704
} // 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)