Skip to content

Commit 7dd8122

Browse files
[Flang][MLIR][OpenMP] - Add support for firstprivate when translating omp.target ops from MLIR to LLVMIR (llvm#131213)
This patch adds support to translate `firstprivate` clauses on `omp.target` ops when translating from MLIR to LLVMIR. Presently, this PR is restricted to supporting only included tasks, i.e `#omp target nowait firstprivate(some_variable)` will likely not work correctly even if it produces object code.
1 parent c88b537 commit 7dd8122

File tree

10 files changed

+207
-69
lines changed

10 files changed

+207
-69
lines changed

flang/lib/Lower/OpenMP/OpenMP.cpp

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1719,13 +1719,14 @@ static void genTargetClauses(
17191719
cp.processNowait(clauseOps);
17201720
cp.processThreadLimit(stmtCtx, clauseOps);
17211721

1722-
cp.processTODO<clause::Allocate, clause::Defaultmap, clause::Firstprivate,
1723-
clause::InReduction, clause::UsesAllocators>(
1724-
loc, llvm::omp::Directive::OMPD_target);
1722+
cp.processTODO<clause::Allocate, clause::Defaultmap, clause::InReduction,
1723+
clause::UsesAllocators>(loc,
1724+
llvm::omp::Directive::OMPD_target);
17251725

17261726
// `target private(..)` is only supported in delayed privatization mode.
17271727
if (!enableDelayedPrivatizationStaging)
1728-
cp.processTODO<clause::Private>(loc, llvm::omp::Directive::OMPD_target);
1728+
cp.processTODO<clause::Firstprivate, clause::Private>(
1729+
loc, llvm::omp::Directive::OMPD_target);
17291730
}
17301731

17311732
static void genTargetDataClauses(

flang/lib/Optimizer/OpenMP/MapsForPrivatizedSymbols.cpp

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,14 @@
3838
#include <type_traits>
3939

4040
#define DEBUG_TYPE "omp-maps-for-privatized-symbols"
41-
41+
#define PDBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "]: ")
4242
namespace flangomp {
4343
#define GEN_PASS_DEF_MAPSFORPRIVATIZEDSYMBOLSPASS
4444
#include "flang/Optimizer/OpenMP/Passes.h.inc"
4545
} // namespace flangomp
46+
4647
using namespace mlir;
48+
4749
namespace {
4850
class MapsForPrivatizedSymbolsPass
4951
: public flangomp::impl::MapsForPrivatizedSymbolsPassBase<
@@ -60,14 +62,14 @@ class MapsForPrivatizedSymbolsPass
6062
// We want the first result of the hlfir.declare op because our goal
6163
// is to map the descriptor (fir.box or fir.boxchar) and the first
6264
// result for hlfir.declare is the descriptor if a the symbol being
63-
// decalred needs a descriptor.
65+
// declared needs a descriptor.
6466
// Some types are boxed immediately before privatization. These have other
6567
// operations in between the privatization and the declaration. It is safe
6668
// to use var directly here because they will be boxed anyway.
6769
if (auto declOp = llvm::dyn_cast_if_present<hlfir::DeclareOp>(definingOp))
6870
varPtr = declOp.getBase();
6971

70-
// If we do not have a reference to descritor, but the descriptor itself
72+
// If we do not have a reference to a descriptor but the descriptor itself,
7173
// then we need to store that on the stack so that we can map the
7274
// address of the descriptor.
7375
if (mlir::isa<fir::BaseBoxType>(varPtr.getType()) ||
@@ -81,6 +83,15 @@ class MapsForPrivatizedSymbolsPass
8183
builder.create<fir::StoreOp>(loc, varPtr, alloca);
8284
varPtr = alloca;
8385
}
86+
assert(mlir::isa<omp::PointerLikeType>(varPtr.getType()) &&
87+
"Dealing with a varPtr that is not a PointerLikeType");
88+
89+
// Figure out the bounds because knowing the bounds will help the subsequent
90+
// MapInfoFinalizationPass map the underlying data of the descriptor.
91+
llvm::SmallVector<mlir::Value> boundsOps;
92+
if (needsBoundsOps(varPtr))
93+
genBoundsOps(builder, varPtr, boundsOps);
94+
8495
return builder.create<omp::MapInfoOp>(
8596
loc, varPtr.getType(), varPtr,
8697
TypeAttr::get(llvm::cast<omp::PointerLikeType>(varPtr.getType())
@@ -92,7 +103,7 @@ class MapsForPrivatizedSymbolsPass
92103
/*varPtrPtr=*/Value{},
93104
/*members=*/SmallVector<Value>{},
94105
/*member_index=*/mlir::ArrayAttr{},
95-
/*bounds=*/ValueRange{},
106+
/*bounds=*/boundsOps,
96107
/*mapperId=*/mlir::FlatSymbolRefAttr(), /*name=*/StringAttr(),
97108
builder.getBoolAttr(false));
98109
}
@@ -143,8 +154,8 @@ class MapsForPrivatizedSymbolsPass
143154
omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder);
144155
mapInfoOps.push_back(mapInfoOp);
145156

146-
LLVM_DEBUG(llvm::dbgs() << "MapsForPrivatizedSymbolsPass created ->\n");
147-
LLVM_DEBUG(mapInfoOp.dump());
157+
LLVM_DEBUG(PDBGS() << "MapsForPrivatizedSymbolsPass created ->\n"
158+
<< mapInfoOp << "\n");
148159
}
149160
if (!mapInfoOps.empty()) {
150161
mapInfoOpsForTarget.insert({targetOp.getOperation(), mapInfoOps});
@@ -158,5 +169,52 @@ class MapsForPrivatizedSymbolsPass
158169
}
159170
}
160171
}
172+
// As the name suggests, this function examines var to determine if
173+
// it has dynamic size. If true, this pass'll have to extract these
174+
// bounds from descriptor of var and add the bounds to the resultant
175+
// MapInfoOp.
176+
bool needsBoundsOps(mlir::Value var) {
177+
assert(mlir::isa<omp::PointerLikeType>(var.getType()) &&
178+
"needsBoundsOps can deal only with pointer types");
179+
mlir::Type t = fir::unwrapRefType(var.getType());
180+
// t could be a box, so look inside the box
181+
auto innerType = fir::dyn_cast_ptrOrBoxEleTy(t);
182+
if (innerType)
183+
return fir::hasDynamicSize(innerType);
184+
return fir::hasDynamicSize(t);
185+
}
186+
187+
// TODO: Remove this in favor of fir::factory::genImplicitBoundsOps
188+
// in a subsequent PR.
189+
void genBoundsOps(fir::FirOpBuilder &builder, mlir::Value var,
190+
llvm::SmallVector<mlir::Value> &boundsOps) {
191+
if (!fir::isBoxAddress(var.getType()))
192+
return;
193+
194+
unsigned int rank = 0;
195+
rank = fir::getBoxRank(fir::unwrapRefType(var.getType()));
196+
mlir::Location loc = var.getLoc();
197+
mlir::Type idxTy = builder.getIndexType();
198+
mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
199+
mlir::Value zero = builder.createIntegerConstant(loc, idxTy, 0);
200+
mlir::Type boundTy = builder.getType<omp::MapBoundsType>();
201+
mlir::Value box = builder.create<fir::LoadOp>(loc, var);
202+
for (unsigned int i = 0; i < rank; ++i) {
203+
mlir::Value dimNo = builder.createIntegerConstant(loc, idxTy, i);
204+
auto dimInfo =
205+
builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy, box, dimNo);
206+
mlir::Value lb = dimInfo.getLowerBound();
207+
mlir::Value extent = dimInfo.getExtent();
208+
mlir::Value byteStride = dimInfo.getByteStride();
209+
mlir::Value ub = builder.create<mlir::arith::SubIOp>(loc, extent, one);
210+
211+
mlir::Value boundsOp = builder.create<omp::MapBoundsOp>(
212+
loc, boundTy, /*lower_bound=*/zero,
213+
/*upper_bound=*/ub, /*extent=*/extent, /*stride=*/byteStride,
214+
/*stride_in_bytes = */ true, /*start_idx=*/lb);
215+
LLVM_DEBUG(PDBGS() << "Created BoundsOp " << boundsOp << "\n");
216+
boundsOps.push_back(boundsOp);
217+
}
218+
}
161219
};
162220
} // namespace

flang/lib/Optimizer/Passes/Pipelines.cpp

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -298,8 +298,14 @@ void createOpenMPFIRPassPipeline(mlir::PassManager &pm,
298298
pm.addPass(flangomp::createDoConcurrentConversionPass(
299299
opts.doConcurrentMappingKind == DoConcurrentMappingKind::DCMK_Device));
300300

301-
pm.addPass(flangomp::createMapInfoFinalizationPass());
301+
// The MapsForPrivatizedSymbols pass needs to run before
302+
// MapInfoFinalizationPass because the former creates new
303+
// MapInfoOp instances, typically for descriptors.
304+
// MapInfoFinalizationPass adds MapInfoOp instances for the descriptors
305+
// underlying data which is necessary to access the data on the offload
306+
// target device.
302307
pm.addPass(flangomp::createMapsForPrivatizedSymbolsPass());
308+
pm.addPass(flangomp::createMapInfoFinalizationPass());
303309
pm.addPass(flangomp::createMarkDeclareTargetPass());
304310
pm.addPass(flangomp::createGenericLoopConversionPass());
305311
if (opts.isTargetDevice)

flang/test/Lower/OpenMP/DelayedPrivatization/target-private-allocatable.f90

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,9 @@ end subroutine target_allocatable
5858
! CHECK: %[[VAR_ALLOC:.*]] = fir.alloca [[DESC_TYPE]]
5959
! CHECK-SAME: {bindc_name = "alloc_var", {{.*}}}
6060
! CHECK: %[[VAR_DECL:.*]]:2 = hlfir.declare %[[VAR_ALLOC]]
61+
! CHECK: %[[BASE_ADDR:.*]] = fir.box_offset %[[VAR_DECL]]#0 base_addr : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> [[MEMBER_TYPE:.*]]
62+
! CHECK: %[[MEMBER:.*]] = omp.map.info var_ptr(%[[VAR_DECL]]#0 : [[TYPE]], i32) map_clauses(to) capture(ByRef) var_ptr_ptr(%[[BASE_ADDR]] : [[MEMBER_TYPE:.*]]) -> {{.*}}
63+
! CHECK: %[[MAP_VAR:.*]] = omp.map.info var_ptr(%[[VAR_DECL]]#0 : [[TYPE]], [[DESC_TYPE]]) map_clauses(to) capture(ByRef) members(%[[MEMBER]] : [0] : !fir.llvm_ptr<!fir.ref<i32>>) -> !fir.ref<!fir.box<!fir.heap<i32>>>
6164

62-
! CHECK: %[[MAP_VAR:.*]] = omp.map.info var_ptr(%[[VAR_DECL]]#0 : [[TYPE]], [[DESC_TYPE]])
63-
! CHECK-SAME: map_clauses(to) capture(ByRef) -> [[TYPE]]
64-
! CHECK: omp.target map_entries(%[[MAP_VAR]] -> %arg0 : [[TYPE]]) private(
65-
! CHECK-SAME: @[[VAR_PRIVATIZER_SYM]] %[[VAR_DECL]]#0 -> %{{.*}} : [[TYPE]]) {
65+
! CHECK: omp.target map_entries(%[[MAP_VAR]] -> %arg0, %[[MEMBER]] -> %arg1 : [[TYPE]], [[MEMBER_TYPE]]) private(
66+
! CHECK-SAME: @[[VAR_PRIVATIZER_SYM]] %[[VAR_DECL]]#0 -> %{{.*}} [map_idx=0] : [[TYPE]]) {

flang/test/Lower/OpenMP/DelayedPrivatization/target-private-multiple-variables.f90

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -140,17 +140,19 @@ end subroutine target_allocatable
140140
! CHECK: %[[REAL_ARR_DECL:.*]]:2 = hlfir.declare %[[REAL_ARR_ALLOC]]({{.*}})
141141
! CHECK: fir.store %[[REAL_ARR_DECL]]#0 to %[[REAL_ARR_DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xf32>>>
142142
! CHECK: %[[MAPPED_MI0:.*]] = omp.map.info var_ptr(%[[MAPPED_DECL]]#1 : !fir.ref<i32>, i32) {{.*}}
143-
! CHECK: %[[ALLOC_VAR_MAP:.*]] = omp.map.info var_ptr(%[[ALLOC_VAR_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<i32>>>, !fir.box<!fir.heap<i32>>)
144-
! CHECK: %[[REAL_ARR_DESC_MAP:.*]] = omp.map.info var_ptr(%[[REAL_ARR_DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>)
143+
! CHECK: %[[ALLOC_VAR_MEMBER:.*]] = omp.map.info var_ptr(%[[ALLOC_VAR_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<i32>>>, i32)
144+
! CHECK: %[[ALLOC_VAR_MAP:.*]] = omp.map.info var_ptr(%[[ALLOC_VAR_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<i32>>>, !fir.box<!fir.heap<i32>>) {{.*}} members(%[[ALLOC_VAR_MEMBER]] :
145+
! CHECK: %[[REAL_ARR_MEMBER:.*]] = omp.map.info var_ptr(%[[REAL_ARR_DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, f32)
146+
! CHECK: %[[REAL_ARR_DESC_MAP:.*]] = omp.map.info var_ptr(%[[REAL_ARR_DESC_ALLOCA]] : !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.box<!fir.array<?xf32>>) {{.*}} members(%[[REAL_ARR_MEMBER]] :
145147
! CHECK: fir.store %[[CHAR_VAR_DECL]]#0 to %[[CHAR_VAR_DESC_ALLOCA]] : !fir.ref<!fir.boxchar<1>>
146148
! CHECK: %[[CHAR_VAR_DESC_MAP:.*]] = omp.map.info var_ptr(%[[CHAR_VAR_DESC_ALLOCA]] : !fir.ref<!fir.boxchar<1>>, !fir.boxchar<1>)
147149
! CHECK: omp.target
148150
! CHECK-SAME: map_entries(
149151
! CHECK-SAME: %[[MAPPED_MI0]] -> %[[MAPPED_ARG0:[^,]+]],
150152
! CHECK-SAME: %[[ALLOC_VAR_MAP]] -> %[[MAPPED_ARG1:[^,]+]]
151153
! CHECK-SAME: %[[REAL_ARR_DESC_MAP]] -> %[[MAPPED_ARG2:[^,]+]]
152-
! CHECK-SAME: %[[CHAR_VAR_DESC_MAP]] -> %[[MAPPED_ARG3:.[^,]+]] :
153-
! CHECK-SAME: !fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.ref<!fir.boxchar<1>>)
154+
! CHECK-SAME: %[[CHAR_VAR_DESC_MAP]] -> %[[MAPPED_ARG3:.[^,]+]]
155+
! CHECK-SAME: !fir.ref<i32>, !fir.ref<!fir.box<!fir.heap<i32>>>, !fir.ref<!fir.box<!fir.array<?xf32>>>, !fir.ref<!fir.boxchar<1>>, !fir.llvm_ptr<!fir.ref<i32>>, !fir.llvm_ptr<!fir.ref<!fir.array<?xf32>>>
154156
! CHECK-SAME: private(
155157
! CHECK-SAME: @[[ALLOC_PRIVATIZER_SYM]] %{{[^[:space:]]+}}#0 -> %[[ALLOC_ARG:[^,]+]] [map_idx=1],
156158
! CHECK-SAME: @[[REAL_PRIVATIZER_SYM]] %{{[^[:space:]]+}}#0 -> %[[REAL_ARG:[^,]+]],

flang/test/Lower/OpenMP/Todo/firstprivate-target.f90

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
integer :: i
55
! CHECK: not yet implemented: Unhandled clause FIRSTPRIVATE in TARGET construct
6-
!$omp target firstprivate(i)
6+
!$omp target firstprivate(i) nowait
77
!$omp end target
88

99
end program

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove, RecipeInterface]>
8181

8282
There are no restrictions on the body except for:
8383
- The `dealloc` regions has a single argument.
84-
- The `init & `copy` regions have 2 arguments.
84+
- The `init` & `copy` regions have 2 arguments.
8585
- All three regions are terminated by `omp.yield` ops.
8686
The above restrictions and other obvious restrictions (e.g. verifying the
8787
type of yielded values) are verified by the custom op verifier. The actual
@@ -90,15 +90,15 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove, RecipeInterface]>
9090
Instances of this op would then be used by ops that model directives that
9191
accept data-sharing attribute clauses.
9292

93-
The $sym_name attribute provides a symbol by which the privatizer op can be
93+
The `sym_name` attribute provides a symbol by which the privatizer op can be
9494
referenced by other dialect ops.
9595

96-
The $type attribute is the type of the value being privatized. This type
96+
The `type` attribute is the type of the value being privatized. This type
9797
will be implicitly allocated in MLIR->LLVMIR conversion and passed as the
9898
second argument to the init region. Therefore the type of arguments to
99-
the regions should be a type which represents a pointer to $type.
99+
the regions should be a type which represents a pointer to `type`.
100100

101-
The $data_sharing_type attribute specifies whether privatizer corresponds
101+
The `data_sharing_type` attribute specifies whether privatizer corresponds
102102
to a `private` or a `firstprivate` clause.
103103
}];
104104

@@ -161,9 +161,10 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove, RecipeInterface]>
161161
/// needsMap returns true if the value being privatized should additionally
162162
/// be mapped to the target region using a MapInfoOp. This is most common
163163
/// when an allocatable is privatized. In such cases, the descriptor is used
164-
/// in privatization and needs to be mapped on to the device.
164+
/// in privatization and needs to be mapped on to the device. The use of
165+
/// firstprivate also creates the need to map the host variable to the device.
165166
bool needsMap() {
166-
return initReadsFromMold();
167+
return readsFromMold();
167168
}
168169

169170
/// Get the type for arguments to nested regions. This should

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Lines changed: 18 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -216,19 +216,9 @@ static LogicalResult checkImplementationStatus(Operation &op) {
216216
};
217217
auto checkPrivate = [&todo](auto op, LogicalResult &result) {
218218
if constexpr (std::is_same_v<std::decay_t<decltype(op)>, omp::TargetOp>) {
219-
// Privatization clauses are supported, except on some situations, so we
220-
// need to check here whether any of these unsupported cases are being
221-
// translated.
222-
if (std::optional<ArrayAttr> privateSyms = op.getPrivateSyms()) {
223-
for (Attribute privatizerNameAttr : *privateSyms) {
224-
omp::PrivateClauseOp privatizer = findPrivatizer(
225-
op.getOperation(), cast<SymbolRefAttr>(privatizerNameAttr));
226-
227-
if (privatizer.getDataSharingType() ==
228-
omp::DataSharingClauseType::FirstPrivate)
229-
result = todo("firstprivate");
230-
}
231-
}
219+
// Privatization is supported only for included target tasks.
220+
if (!op.getPrivateVars().empty() && op.getNowait())
221+
result = todo("privatization for deferred target tasks");
232222
} else {
233223
if (!op.getPrivateVars().empty() || op.getPrivateSyms())
234224
result = todo("privatization");
@@ -1487,6 +1477,7 @@ allocatePrivateVars(llvm::IRBuilderBase &builder,
14871477
assert(allocaTerminator->getNumSuccessors() == 1 &&
14881478
"This is an unconditional branch created by splitBB");
14891479

1480+
llvm::DataLayout dataLayout = builder.GetInsertBlock()->getDataLayout();
14901481
llvm::BasicBlock *afterAllocas = allocaTerminator->getSuccessor(0);
14911482

14921483
unsigned int allocaAS =
@@ -1513,12 +1504,12 @@ allocatePrivateVars(llvm::IRBuilderBase &builder,
15131504
return afterAllocas;
15141505
}
15151506

1516-
static LogicalResult
1517-
copyFirstPrivateVars(llvm::IRBuilderBase &builder,
1518-
LLVM::ModuleTranslation &moduleTranslation,
1519-
SmallVectorImpl<mlir::Value> &mlirPrivateVars,
1520-
ArrayRef<llvm::Value *> llvmPrivateVars,
1521-
SmallVectorImpl<omp::PrivateClauseOp> &privateDecls) {
1507+
static LogicalResult copyFirstPrivateVars(
1508+
llvm::IRBuilderBase &builder, LLVM::ModuleTranslation &moduleTranslation,
1509+
SmallVectorImpl<mlir::Value> &mlirPrivateVars,
1510+
ArrayRef<llvm::Value *> llvmPrivateVars,
1511+
SmallVectorImpl<omp::PrivateClauseOp> &privateDecls,
1512+
llvm::DenseMap<Value, Value> *mappedPrivateVars = nullptr) {
15221513
// Apply copy region for firstprivate.
15231514
bool needsFirstprivate =
15241515
llvm::any_of(privateDecls, [](omp::PrivateClauseOp &privOp) {
@@ -1542,7 +1533,8 @@ copyFirstPrivateVars(llvm::IRBuilderBase &builder,
15421533
Region &copyRegion = decl.getCopyRegion();
15431534

15441535
// map copyRegion rhs arg
1545-
llvm::Value *nonPrivateVar = moduleTranslation.lookupValue(mlirVar);
1536+
llvm::Value *nonPrivateVar = findAssociatedValue(
1537+
mlirVar, builder, moduleTranslation, mappedPrivateVars);
15461538
assert(nonPrivateVar);
15471539
moduleTranslation.mapValue(decl.getCopyMoldArg(), nonPrivateVar);
15481540

@@ -5129,6 +5121,12 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
51295121
.failed())
51305122
return llvm::make_error<PreviouslyReportedError>();
51315123

5124+
if (failed(copyFirstPrivateVars(
5125+
builder, moduleTranslation, privateVarsInfo.mlirVars,
5126+
privateVarsInfo.llvmVars, privateVarsInfo.privatizers,
5127+
&mappedPrivateVars)))
5128+
return llvm::make_error<PreviouslyReportedError>();
5129+
51325130
SmallVector<Region *> privateCleanupRegions;
51335131
llvm::transform(privateVarsInfo.privatizers,
51345132
std::back_inserter(privateCleanupRegions),

0 commit comments

Comments
 (0)