Skip to content

Commit 26b8be2

Browse files
ergawyergawy
authored
[flang][OpenMP][MLIR] Basic support for delayed privatization code-gen (#81833)
Adds basic support for emitting delayed privatizers from flang. So far, only types of symbols are supported (i.e. scalars), support for more complicated types will be added later. This also makes sure that reduction and delayed privatization work properly together by merging the body-gen callbacks for both in case both clauses are present on the parallel construct.
1 parent 1d61709 commit 26b8be2

13 files changed

+422
-28
lines changed

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

+6
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
#include "flang/Common/Fortran.h"
1717
#include "flang/Lower/LoweringOptions.h"
1818
#include "flang/Lower/PFTDefs.h"
19+
#include "flang/Lower/SymbolMap.h"
1920
#include "flang/Optimizer/Builder/BoxValue.h"
2021
#include "flang/Semantics/symbol.h"
2122
#include "mlir/IR/Builders.h"
@@ -299,6 +300,11 @@ class AbstractConverter {
299300
return loweringOptions;
300301
}
301302

303+
/// Find the symbol in one level up of symbol map such as for host-association
304+
/// in OpenMP code or return null.
305+
virtual Fortran::lower::SymbolBox
306+
lookupOneLevelUpSymbol(const Fortran::semantics::Symbol &sym) = 0;
307+
302308
private:
303309
/// Options controlling lowering behavior.
304310
const Fortran::lower::LoweringOptions &loweringOptions;

‎flang/lib/Lower/Bridge.cpp

+12-1
Original file line numberDiff line numberDiff line change
@@ -1000,6 +1000,17 @@ class FirConverter : public Fortran::lower::AbstractConverter {
10001000
if (sym.detailsIf<Fortran::semantics::CommonBlockDetails>())
10011001
return symMap->lookupSymbol(sym);
10021002

1003+
// For symbols to be privatized in OMP, the symbol is mapped to an
1004+
// instance of `SymbolBox::Intrinsic` (i.e. a direct mapping to an MLIR
1005+
// SSA value). This MLIR SSA value is the block argument to the
1006+
// `omp.private`'s `alloc` block. If this is the case, we return this
1007+
// `SymbolBox::Intrinsic` value.
1008+
if (Fortran::lower::SymbolBox v = symMap->lookupSymbol(sym))
1009+
return v.match(
1010+
[&](const Fortran::lower::SymbolBox::Intrinsic &)
1011+
-> Fortran::lower::SymbolBox { return v; },
1012+
[](const auto &) -> Fortran::lower::SymbolBox { return {}; });
1013+
10031014
return {};
10041015
}
10051016
if (Fortran::lower::SymbolBox v = symMap->lookupSymbol(sym))
@@ -1018,7 +1029,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
10181029
/// Find the symbol in one level up of symbol map such as for host-association
10191030
/// in OpenMP code or return null.
10201031
Fortran::lower::SymbolBox
1021-
lookupOneLevelUpSymbol(const Fortran::semantics::Symbol &sym) {
1032+
lookupOneLevelUpSymbol(const Fortran::semantics::Symbol &sym) override {
10221033
if (Fortran::lower::SymbolBox v = localSymbols.lookupOneLevelUpSymbol(sym))
10231034
return v;
10241035
return {};

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

+94-13
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,10 @@ void DataSharingProcessor::cloneSymbol(const Fortran::semantics::Symbol *sym) {
6666
}
6767

6868
void DataSharingProcessor::copyFirstPrivateSymbol(
69-
const Fortran::semantics::Symbol *sym) {
69+
const Fortran::semantics::Symbol *sym,
70+
mlir::OpBuilder::InsertPoint *copyAssignIP) {
7071
if (sym->test(Fortran::semantics::Symbol::Flag::OmpFirstPrivate))
71-
converter.copyHostAssociateVar(*sym);
72+
converter.copyHostAssociateVar(*sym, copyAssignIP);
7273
}
7374

7475
void DataSharingProcessor::copyLastPrivateSymbol(
@@ -307,14 +308,10 @@ void DataSharingProcessor::privatize() {
307308
for (const Fortran::semantics::Symbol *sym : privatizedSymbols) {
308309
if (const auto *commonDet =
309310
sym->detailsIf<Fortran::semantics::CommonBlockDetails>()) {
310-
for (const auto &mem : commonDet->objects()) {
311-
cloneSymbol(&*mem);
312-
copyFirstPrivateSymbol(&*mem);
313-
}
314-
} else {
315-
cloneSymbol(sym);
316-
copyFirstPrivateSymbol(sym);
317-
}
311+
for (const auto &mem : commonDet->objects())
312+
doPrivatize(&*mem);
313+
} else
314+
doPrivatize(sym);
318315
}
319316
}
320317

@@ -338,11 +335,95 @@ void DataSharingProcessor::defaultPrivatize() {
338335
!sym->GetUltimate().has<Fortran::semantics::NamelistDetails>() &&
339336
!symbolsInNestedRegions.contains(sym) &&
340337
!symbolsInParentRegions.contains(sym) &&
341-
!privatizedSymbols.contains(sym)) {
338+
!privatizedSymbols.contains(sym))
339+
doPrivatize(sym);
340+
}
341+
}
342+
343+
void DataSharingProcessor::doPrivatize(const Fortran::semantics::Symbol *sym) {
344+
if (!useDelayedPrivatization) {
345+
cloneSymbol(sym);
346+
copyFirstPrivateSymbol(sym);
347+
return;
348+
}
349+
350+
Fortran::lower::SymbolBox hsb = converter.lookupOneLevelUpSymbol(*sym);
351+
assert(hsb && "Host symbol box not found");
352+
353+
mlir::Type symType = hsb.getAddr().getType();
354+
mlir::Location symLoc = hsb.getAddr().getLoc();
355+
std::string privatizerName = sym->name().ToString() + ".privatizer";
356+
bool isFirstPrivate =
357+
sym->test(Fortran::semantics::Symbol::Flag::OmpFirstPrivate);
358+
359+
mlir::omp::PrivateClauseOp privatizerOp = [&]() {
360+
auto moduleOp = firOpBuilder.getModule();
361+
auto uniquePrivatizerName = fir::getTypeAsString(
362+
symType, converter.getKindMap(),
363+
converter.mangleName(*sym) +
364+
(isFirstPrivate ? "_firstprivate" : "_private"));
365+
366+
if (auto existingPrivatizer =
367+
moduleOp.lookupSymbol<mlir::omp::PrivateClauseOp>(
368+
uniquePrivatizerName))
369+
return existingPrivatizer;
370+
371+
auto ip = firOpBuilder.saveInsertionPoint();
372+
firOpBuilder.setInsertionPoint(&moduleOp.getBodyRegion().front(),
373+
moduleOp.getBodyRegion().front().begin());
374+
auto result = firOpBuilder.create<mlir::omp::PrivateClauseOp>(
375+
symLoc, uniquePrivatizerName, symType,
376+
isFirstPrivate ? mlir::omp::DataSharingClauseType::FirstPrivate
377+
: mlir::omp::DataSharingClauseType::Private);
378+
379+
symTable->pushScope();
380+
381+
// Populate the `alloc` region.
382+
{
383+
mlir::Region &allocRegion = result.getAllocRegion();
384+
mlir::Block *allocEntryBlock = firOpBuilder.createBlock(
385+
&allocRegion, /*insertPt=*/{}, symType, symLoc);
386+
387+
firOpBuilder.setInsertionPointToEnd(allocEntryBlock);
388+
symTable->addSymbol(*sym, allocRegion.getArgument(0));
389+
symTable->pushScope();
342390
cloneSymbol(sym);
343-
copyFirstPrivateSymbol(sym);
391+
firOpBuilder.create<mlir::omp::YieldOp>(
392+
hsb.getAddr().getLoc(),
393+
symTable->shallowLookupSymbol(*sym).getAddr());
394+
symTable->popScope();
344395
}
345-
}
396+
397+
// Populate the `copy` region if this is a `firstprivate`.
398+
if (isFirstPrivate) {
399+
mlir::Region &copyRegion = result.getCopyRegion();
400+
// First block argument corresponding to the original/host value while
401+
// second block argument corresponding to the privatized value.
402+
mlir::Block *copyEntryBlock = firOpBuilder.createBlock(
403+
&copyRegion, /*insertPt=*/{}, {symType, symType}, {symLoc, symLoc});
404+
firOpBuilder.setInsertionPointToEnd(copyEntryBlock);
405+
symTable->addSymbol(*sym, copyRegion.getArgument(0),
406+
/*force=*/true);
407+
symTable->pushScope();
408+
symTable->addSymbol(*sym, copyRegion.getArgument(1));
409+
auto ip = firOpBuilder.saveInsertionPoint();
410+
copyFirstPrivateSymbol(sym, &ip);
411+
412+
firOpBuilder.create<mlir::omp::YieldOp>(
413+
hsb.getAddr().getLoc(),
414+
symTable->shallowLookupSymbol(*sym).getAddr());
415+
symTable->popScope();
416+
}
417+
418+
symTable->popScope();
419+
firOpBuilder.restoreInsertionPoint(ip);
420+
return result;
421+
}();
422+
423+
delayedPrivatizationInfo.privatizers.push_back(
424+
mlir::SymbolRefAttr::get(privatizerOp));
425+
delayedPrivatizationInfo.originalAddresses.push_back(hsb.getAddr());
426+
delayedPrivatizationInfo.symbols.push_back(sym);
346427
}
347428

348429
} // namespace omp

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

+35-3
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,24 @@ namespace lower {
2323
namespace omp {
2424

2525
class DataSharingProcessor {
26+
public:
27+
/// Collects all the information needed for delayed privatization. This can be
28+
/// used by ops with data-sharing clauses to properly generate their regions
29+
/// (e.g. add region arguments) and map the original SSA values to their
30+
/// corresponding OMP region operands.
31+
struct DelayedPrivatizationInfo {
32+
// The list of symbols referring to delayed privatizer ops (i.e.
33+
// `omp.private` ops).
34+
llvm::SmallVector<mlir::SymbolRefAttr> privatizers;
35+
// SSA values that correspond to "original" values being privatized.
36+
// "Original" here means the SSA value outside the OpenMP region from which
37+
// a clone is created inside the region.
38+
llvm::SmallVector<mlir::Value> originalAddresses;
39+
// Fortran symbols corresponding to the above SSA values.
40+
llvm::SmallVector<const Fortran::semantics::Symbol *> symbols;
41+
};
42+
43+
private:
2644
bool hasLastPrivateOp;
2745
mlir::OpBuilder::InsertPoint lastPrivIP;
2846
mlir::OpBuilder::InsertPoint insPt;
@@ -36,6 +54,9 @@ class DataSharingProcessor {
3654
fir::FirOpBuilder &firOpBuilder;
3755
const Fortran::parser::OmpClauseList &opClauseList;
3856
Fortran::lower::pft::Evaluation &eval;
57+
bool useDelayedPrivatization;
58+
Fortran::lower::SymMap *symTable;
59+
DelayedPrivatizationInfo delayedPrivatizationInfo;
3960

4061
bool needBarrier();
4162
void collectSymbols(Fortran::semantics::Symbol::Flag flag);
@@ -47,21 +68,28 @@ class DataSharingProcessor {
4768
void collectDefaultSymbols();
4869
void privatize();
4970
void defaultPrivatize();
71+
void doPrivatize(const Fortran::semantics::Symbol *sym);
5072
void copyLastPrivatize(mlir::Operation *op);
5173
void insertLastPrivateCompare(mlir::Operation *op);
5274
void cloneSymbol(const Fortran::semantics::Symbol *sym);
53-
void copyFirstPrivateSymbol(const Fortran::semantics::Symbol *sym);
75+
void
76+
copyFirstPrivateSymbol(const Fortran::semantics::Symbol *sym,
77+
mlir::OpBuilder::InsertPoint *copyAssignIP = nullptr);
5478
void copyLastPrivateSymbol(const Fortran::semantics::Symbol *sym,
5579
mlir::OpBuilder::InsertPoint *lastPrivIP);
5680
void insertDeallocs();
5781

5882
public:
5983
DataSharingProcessor(Fortran::lower::AbstractConverter &converter,
6084
const Fortran::parser::OmpClauseList &opClauseList,
61-
Fortran::lower::pft::Evaluation &eval)
85+
Fortran::lower::pft::Evaluation &eval,
86+
bool useDelayedPrivatization = false,
87+
Fortran::lower::SymMap *symTable = nullptr)
6288
: hasLastPrivateOp(false), converter(converter),
6389
firOpBuilder(converter.getFirOpBuilder()), opClauseList(opClauseList),
64-
eval(eval) {}
90+
eval(eval), useDelayedPrivatization(useDelayedPrivatization),
91+
symTable(symTable) {}
92+
6593
// Privatisation is split into two steps.
6694
// Step1 performs cloning of all privatisation clauses and copying for
6795
// firstprivates. Step1 is performed at the place where process/processStep1
@@ -80,6 +108,10 @@ class DataSharingProcessor {
80108
assert(!loopIV && "Loop iteration variable already set");
81109
loopIV = iv;
82110
}
111+
112+
const DelayedPrivatizationInfo &getDelayedPrivatizationInfo() const {
113+
return delayedPrivatizationInfo;
114+
}
83115
};
84116

85117
} // namespace omp

0 commit comments

Comments
 (0)