Skip to content

Commit 2b22a17

Browse files
ergawyergawy
committed
[OpenMP][MLIR] Support LLVM translation for distribute with delayed privatization
Adds support for tranlating delayed privatization (`private` and `firstprivate`) for `omp.distribute` ops.
1 parent 87916f8 commit 2b22a17

File tree

3 files changed

+126
-16
lines changed

3 files changed

+126
-16
lines changed

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
250250
checkAllocate(op, result);
251251
checkDistSchedule(op, result);
252252
checkOrder(op, result);
253-
checkPrivate(op, result);
254253
})
255254
.Case([&](omp::OrderedRegionOp op) { checkParLevelSimd(op, result); })
256255
.Case([&](omp::SectionsOp op) {
@@ -4189,6 +4188,38 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
41894188
// DistributeOp has only one region associated with it.
41904189
builder.restoreIP(codeGenIP);
41914190

4191+
// TODO This is a recurring pattern in almost all ops that need
4192+
// privatization. Try to abstract it in a shared util/interface.
4193+
MutableArrayRef<BlockArgument> privateBlockArgs =
4194+
cast<omp::BlockArgOpenMPOpInterface>(*distributeOp)
4195+
.getPrivateBlockArgs();
4196+
SmallVector<mlir::Value> mlirPrivateVars;
4197+
SmallVector<llvm::Value *> llvmPrivateVars;
4198+
SmallVector<omp::PrivateClauseOp> privateDecls;
4199+
mlirPrivateVars.reserve(privateBlockArgs.size());
4200+
llvmPrivateVars.reserve(privateBlockArgs.size());
4201+
collectPrivatizationDecls(distributeOp, privateDecls);
4202+
4203+
for (mlir::Value privateVar : distributeOp.getPrivateVars())
4204+
mlirPrivateVars.push_back(privateVar);
4205+
4206+
llvm::Expected<llvm::BasicBlock *> afterAllocas = allocatePrivateVars(
4207+
builder, moduleTranslation, privateBlockArgs, privateDecls,
4208+
mlirPrivateVars, llvmPrivateVars, allocaIP);
4209+
if (handleError(afterAllocas, opInst).failed())
4210+
return llvm::make_error<PreviouslyReportedError>();
4211+
4212+
if (handleError(initPrivateVars(builder, moduleTranslation,
4213+
privateBlockArgs, privateDecls,
4214+
mlirPrivateVars, llvmPrivateVars),
4215+
opInst)
4216+
.failed())
4217+
return llvm::make_error<PreviouslyReportedError>();
4218+
4219+
if (failed(copyFirstPrivateVars(builder, moduleTranslation, mlirPrivateVars,
4220+
llvmPrivateVars, privateDecls)))
4221+
return llvm::make_error<PreviouslyReportedError>();
4222+
41924223
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
41934224
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
41944225
llvm::Expected<llvm::BasicBlock *> regionBlock =

‎mlir/test/Target/LLVMIR/openmp-distribute-private.mlir

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
// Test code-gen for `omp.distribute` ops with delayed privatizers (i.e. using
2+
// `omp.private` ops).
3+
4+
// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
5+
6+
omp.private {type = private} @_QFEi_private_i32 : i32
7+
omp.private {type = private} @_QFEpriv_val_dist_private_f32 : f32
8+
9+
llvm.func @_QQmain() {
10+
%0 = llvm.mlir.constant(1 : i64) : i64
11+
%1 = llvm.alloca %0 x f32 {bindc_name = "priv_val_dist"} : (i64) -> !llvm.ptr
12+
%3 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
13+
%4 = llvm.mlir.constant(3.140000e+00 : f32) : f32
14+
%5 = llvm.mlir.constant(1000 : i32) : i32
15+
%6 = llvm.mlir.constant(1 : i32) : i32
16+
17+
omp.teams {
18+
omp.distribute private(@_QFEpriv_val_dist_private_f32 %1 -> %arg0, @_QFEi_private_i32 %3 -> %arg1 : !llvm.ptr, !llvm.ptr) {
19+
omp.loop_nest (%arg2) : i32 = (%6) to (%5) inclusive step (%6) {
20+
llvm.store %arg2, %arg1 : i32, !llvm.ptr
21+
llvm.store %4, %arg0 : f32, !llvm.ptr
22+
omp.yield
23+
}
24+
}
25+
omp.terminator
26+
}
27+
28+
llvm.return
29+
}
30+
31+
// CHECK-LABEL: define void @_QQmain() {
32+
// CHECK: call void {{.*}} @__kmpc_fork_teams(ptr @{{.*}}, i32 0, ptr @[[TEAMS_FUNC:.*]])
33+
// CHECK-NEXT: br label %teams.exit
34+
// CHECK: }
35+
36+
// CHECK: define internal void @[[TEAMS_FUNC]]({{.*}}) {
37+
// CHECK: call void @[[DIST_FUNC:.*]]()
38+
// CHECK-NEXT: br label %distribute.exit
39+
// CHECK: }
40+
41+
// CHECK: define internal void @[[DIST_FUNC]]() {
42+
// CHECK: %[[PRIV_VAR_ALLOC:.*]] = alloca float, align 4
43+
// CHECK: %[[IV_ALLOC:.*]] = alloca i32, align 4
44+
45+
// CHECK: omp.loop_nest.region:
46+
// CHECK-NEXT: store i32 %{{.*}}, ptr %[[IV_ALLOC]], align 4
47+
// CHECK-NEXT: store float 0x40091EB860000000, ptr %[[PRIV_VAR_ALLOC]], align 4
48+
// CHECK: }
49+
50+
// -----
51+
52+
omp.private {type = firstprivate} @_QFEpriv_val_dist_firstprivate_f32 : f32 copy {
53+
^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
54+
%0 = llvm.load %arg0 : !llvm.ptr -> f32
55+
llvm.store %0, %arg1 : f32, !llvm.ptr
56+
omp.yield(%arg1 : !llvm.ptr)
57+
}
58+
59+
llvm.func @_QQmain() {
60+
%0 = llvm.mlir.constant(1 : i64) : i64
61+
%1 = llvm.alloca %0 x f32 {bindc_name = "priv_val_dist"} : (i64) -> !llvm.ptr
62+
%4 = llvm.mlir.constant(3.140000e+00 : f32) : f32
63+
%6 = llvm.mlir.constant(1 : i32) : i32
64+
omp.distribute private(@_QFEpriv_val_dist_firstprivate_f32 %1 -> %arg0 : !llvm.ptr) {
65+
omp.loop_nest (%arg2) : i32 = (%6) to (%6) inclusive step (%6) {
66+
llvm.store %4, %arg0 : f32, !llvm.ptr
67+
omp.yield
68+
}
69+
}
70+
llvm.return
71+
}
72+
73+
// CHECK-LABEL: define void @_QQmain() {
74+
// CHECK: %[[SHARED_VAR_ALLOC:.*]] = alloca float, i64 1, align 4
75+
// CHECK: %[[SHARED_VAR_PTR:.*]] = getelementptr { ptr }, ptr %[[DIST_PARAM:.*]], i32 0, i32 0
76+
// CHECK: store ptr %[[SHARED_VAR_ALLOC]], ptr %[[SHARED_VAR_PTR]], align 8
77+
// CHECK: call void @[[DIST_FUNC:.*]](ptr %[[DIST_PARAM]])
78+
// CHECK-NEXT: br label %distribute.exit
79+
// CHECK: }
80+
81+
// CHECK: define internal void @[[DIST_FUNC]](ptr %[[DIST_ARG:.*]]) {
82+
// CHECK: %[[SHARED_VAR_GEP:.*]] = getelementptr { ptr }, ptr %[[DIST_ARG]], i32 0, i32 0
83+
// CHECK: %[[SHARED_VAR_PTR2:.*]] = load ptr, ptr %[[SHARED_VAR_GEP]], align 8
84+
// CHECK: %[[PRIV_VAR_ALLOC:.*]] = alloca float, align 4
85+
86+
// CHECK: omp.private.copy:
87+
// CHECK-NEXT: %[[SHARED_VAR_VAL:.*]] = load float, ptr %[[SHARED_VAR_PTR2]], align 4
88+
// CHECK-NEXT: store float %[[SHARED_VAR_VAL]], ptr %[[PRIV_VAR_ALLOC]], align 4
89+
90+
// CHECK: omp.loop_nest.region:
91+
// CHECK-NEXT: store float 0x40091EB860000000, ptr %[[PRIV_VAR_ALLOC]], align 4
92+
// CHECK: }
93+
94+

‎mlir/test/Target/LLVMIR/openmp-todo.mlir

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -105,21 +105,6 @@ llvm.func @distribute_order(%lb : i32, %ub : i32, %step : i32) {
105105

106106
// -----
107107

108-
omp.private {type = private} @x.privatizer : !llvm.ptr
109-
110-
llvm.func @distribute_private(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
111-
// expected-error@below {{not yet implemented: Unhandled clause privatization in omp.distribute operation}}
112-
// expected-error@below {{LLVM Translation failed for operation: omp.distribute}}
113-
omp.distribute private(@x.privatizer %x -> %arg0 : !llvm.ptr) {
114-
omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
115-
omp.yield
116-
}
117-
}
118-
llvm.return
119-
}
120-
121-
// -----
122-
123108
llvm.func @ordered_region_par_level_simd() {
124109
// expected-error@below {{not yet implemented: Unhandled clause parallelization-level in omp.ordered.region operation}}
125110
// expected-error@below {{LLVM Translation failed for operation: omp.ordered.region}}

0 commit comments

Comments
 (0)