[Mlir-commits] [mlir] [OpenMP][MLIR] Support LLVM translation for `distribute` with delayed privatization (PR #131564)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Mar 16 22:43:13 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm
Author: Kareem Ergawy (ergawy)
<details>
<summary>Changes</summary>
Adds support for tranlating delayed privatization (`private` and `firstprivate`) for `omp.distribute` ops.
---
Full diff: https://github.com/llvm/llvm-project/pull/131564.diff
3 Files Affected:
- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+32-1)
- (added) mlir/test/Target/LLVMIR/openmp-distribute-private.mlir (+94)
- (modified) mlir/test/Target/LLVMIR/openmp-todo.mlir (-15)
``````````diff
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 0010ca6630050..9198ee997be07 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -250,7 +250,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
checkAllocate(op, result);
checkDistSchedule(op, result);
checkOrder(op, result);
- checkPrivate(op, result);
})
.Case([&](omp::OrderedRegionOp op) { checkParLevelSimd(op, result); })
.Case([&](omp::SectionsOp op) {
@@ -4189,6 +4188,38 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
// DistributeOp has only one region associated with it.
builder.restoreIP(codeGenIP);
+ // TODO This is a recurring pattern in almost all ops that need
+ // privatization. Try to abstract it in a shared util/interface.
+ MutableArrayRef<BlockArgument> privateBlockArgs =
+ cast<omp::BlockArgOpenMPOpInterface>(*distributeOp)
+ .getPrivateBlockArgs();
+ SmallVector<mlir::Value> mlirPrivateVars;
+ SmallVector<llvm::Value *> llvmPrivateVars;
+ SmallVector<omp::PrivateClauseOp> privateDecls;
+ mlirPrivateVars.reserve(privateBlockArgs.size());
+ llvmPrivateVars.reserve(privateBlockArgs.size());
+ collectPrivatizationDecls(distributeOp, privateDecls);
+
+ for (mlir::Value privateVar : distributeOp.getPrivateVars())
+ mlirPrivateVars.push_back(privateVar);
+
+ llvm::Expected<llvm::BasicBlock *> afterAllocas = allocatePrivateVars(
+ builder, moduleTranslation, privateBlockArgs, privateDecls,
+ mlirPrivateVars, llvmPrivateVars, allocaIP);
+ if (handleError(afterAllocas, opInst).failed())
+ return llvm::make_error<PreviouslyReportedError>();
+
+ if (handleError(initPrivateVars(builder, moduleTranslation,
+ privateBlockArgs, privateDecls,
+ mlirPrivateVars, llvmPrivateVars),
+ opInst)
+ .failed())
+ return llvm::make_error<PreviouslyReportedError>();
+
+ if (failed(copyFirstPrivateVars(builder, moduleTranslation, mlirPrivateVars,
+ llvmPrivateVars, privateDecls)))
+ return llvm::make_error<PreviouslyReportedError>();
+
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
llvm::Expected<llvm::BasicBlock *> regionBlock =
diff --git a/mlir/test/Target/LLVMIR/openmp-distribute-private.mlir b/mlir/test/Target/LLVMIR/openmp-distribute-private.mlir
new file mode 100644
index 0000000000000..de4accc38a419
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-distribute-private.mlir
@@ -0,0 +1,94 @@
+// Test code-gen for `omp.distribute` ops with delayed privatizers (i.e. using
+// `omp.private` ops).
+
+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
+
+omp.private {type = private} @_QFEi_private_i32 : i32
+omp.private {type = private} @_QFEpriv_val_dist_private_f32 : f32
+
+llvm.func @_QQmain() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x f32 {bindc_name = "priv_val_dist"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %4 = llvm.mlir.constant(3.140000e+00 : f32) : f32
+ %5 = llvm.mlir.constant(1000 : i32) : i32
+ %6 = llvm.mlir.constant(1 : i32) : i32
+
+ omp.teams {
+ omp.distribute private(@_QFEpriv_val_dist_private_f32 %1 -> %arg0, @_QFEi_private_i32 %3 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2) : i32 = (%6) to (%5) inclusive step (%6) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ llvm.store %4, %arg0 : f32, !llvm.ptr
+ omp.yield
+ }
+ }
+ omp.terminator
+ }
+
+ llvm.return
+}
+
+// CHECK-LABEL: define void @_QQmain() {
+// CHECK: call void {{.*}} @__kmpc_fork_teams(ptr @{{.*}}, i32 0, ptr @[[TEAMS_FUNC:.*]])
+// CHECK-NEXT: br label %teams.exit
+// CHECK: }
+
+// CHECK: define internal void @[[TEAMS_FUNC]]({{.*}}) {
+// CHECK: call void @[[DIST_FUNC:.*]]()
+// CHECK-NEXT: br label %distribute.exit
+// CHECK: }
+
+// CHECK: define internal void @[[DIST_FUNC]]() {
+// CHECK: %[[PRIV_VAR_ALLOC:.*]] = alloca float, align 4
+// CHECK: %[[IV_ALLOC:.*]] = alloca i32, align 4
+
+// CHECK: omp.loop_nest.region:
+// CHECK-NEXT: store i32 %{{.*}}, ptr %[[IV_ALLOC]], align 4
+// CHECK-NEXT: store float 0x40091EB860000000, ptr %[[PRIV_VAR_ALLOC]], align 4
+// CHECK: }
+
+// -----
+
+omp.private {type = firstprivate} @_QFEpriv_val_dist_firstprivate_f32 : f32 copy {
+^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ %0 = llvm.load %arg0 : !llvm.ptr -> f32
+ llvm.store %0, %arg1 : f32, !llvm.ptr
+ omp.yield(%arg1 : !llvm.ptr)
+}
+
+llvm.func @_QQmain() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x f32 {bindc_name = "priv_val_dist"} : (i64) -> !llvm.ptr
+ %4 = llvm.mlir.constant(3.140000e+00 : f32) : f32
+ %6 = llvm.mlir.constant(1 : i32) : i32
+ omp.distribute private(@_QFEpriv_val_dist_firstprivate_f32 %1 -> %arg0 : !llvm.ptr) {
+ omp.loop_nest (%arg2) : i32 = (%6) to (%6) inclusive step (%6) {
+ llvm.store %4, %arg0 : f32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK-LABEL: define void @_QQmain() {
+// CHECK: %[[SHARED_VAR_ALLOC:.*]] = alloca float, i64 1, align 4
+// CHECK: %[[SHARED_VAR_PTR:.*]] = getelementptr { ptr }, ptr %[[DIST_PARAM:.*]], i32 0, i32 0
+// CHECK: store ptr %[[SHARED_VAR_ALLOC]], ptr %[[SHARED_VAR_PTR]], align 8
+// CHECK: call void @[[DIST_FUNC:.*]](ptr %[[DIST_PARAM]])
+// CHECK-NEXT: br label %distribute.exit
+// CHECK: }
+
+// CHECK: define internal void @[[DIST_FUNC]](ptr %[[DIST_ARG:.*]]) {
+// CHECK: %[[SHARED_VAR_GEP:.*]] = getelementptr { ptr }, ptr %[[DIST_ARG]], i32 0, i32 0
+// CHECK: %[[SHARED_VAR_PTR2:.*]] = load ptr, ptr %[[SHARED_VAR_GEP]], align 8
+// CHECK: %[[PRIV_VAR_ALLOC:.*]] = alloca float, align 4
+
+// CHECK: omp.private.copy:
+// CHECK-NEXT: %[[SHARED_VAR_VAL:.*]] = load float, ptr %[[SHARED_VAR_PTR2]], align 4
+// CHECK-NEXT: store float %[[SHARED_VAR_VAL]], ptr %[[PRIV_VAR_ALLOC]], align 4
+
+// CHECK: omp.loop_nest.region:
+// CHECK-NEXT: store float 0x40091EB860000000, ptr %[[PRIV_VAR_ALLOC]], align 4
+// CHECK: }
+
+
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 7a4d304f2d2d5..af31f8bab73ac 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -105,21 +105,6 @@ llvm.func @distribute_order(%lb : i32, %ub : i32, %step : i32) {
// -----
-omp.private {type = private} @x.privatizer : !llvm.ptr
-
-llvm.func @distribute_private(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
- // expected-error at below {{not yet implemented: Unhandled clause privatization in omp.distribute operation}}
- // expected-error at below {{LLVM Translation failed for operation: omp.distribute}}
- omp.distribute private(@x.privatizer %x -> %arg0 : !llvm.ptr) {
- omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
- omp.yield
- }
- }
- llvm.return
-}
-
-// -----
-
llvm.func @ordered_region_par_level_simd() {
// expected-error at below {{not yet implemented: Unhandled clause parallelization-level in omp.ordered.region operation}}
// expected-error at below {{LLVM Translation failed for operation: omp.ordered.region}}
``````````
</details>
https://github.com/llvm/llvm-project/pull/131564
More information about the Mlir-commits
mailing list