[Mlir-commits] [mlir] [flang][OpenMP] Process `omp.atomic.update` while translating scopes for target device (PR #132165)
Kareem Ergawy
llvmlistbot at llvm.org
Thu Mar 20 06:51:26 PDT 2025
https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/132165
>From 252502666e8b94c25d75d64fd3743129fd6c9246 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 20 Mar 2025 03:40:32 -0500
Subject: [PATCH 1/2] [flang][OpenMP] Process `omp.atomic.update` while
translating scopes for target device
Fixes a bug introduced by https://github.com/llvm/llvm-project/pull/130078.
For non-BlockArgOpenMPOpInterface ops, we also want to map their entry block arguments to their operands, if any. For the current support in the OpenMP dialect, the table below lists all ops that have arguments (SSA operands and/or attributes) and not target-related. Of all these ops, we need to only process `omp.atomic.update` since it is the only op that has SSA operands & an attached region. Therefore, the region's entry block arguments must be mapped to the op's operands in case they are referenced inside the region.
=============================================================================
| op | operands? | region(s)? | parent is func? | processed? |
=============================================================================
| atomic.read | yes | no | yes | no |
| atomic.write | yes | no | yes | no |
| atomic.update | yes | yes | yes | yes |
| critical | no | no | yes | no |
| declare_mapper | no | yes | no | no |
| declare_reduction | no | yes | no | no |
| flush | yes | no | yes | no |
| private | no | yes | yes | no |
| threadprivate | yes | no | yes | no |
| yield | yes | no | yes | no |
=============================================================================
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 40 ++++++++++++++++++-
.../openmp-target-nesting-in-host-ops.mlir | 21 ++++++++++
2 files changed, 60 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 722107c7ec6d7..cfc4e98ad3341 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5323,6 +5323,43 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
if (auto blockArgsIface =
dyn_cast<omp::BlockArgOpenMPOpInterface>(oper))
forwardArgs(moduleTranslation, blockArgsIface);
+ else {
+ // For non-BlockArgOpenMPOpInterface ops, we also want to map
+ // their entry block arguments to their operands, if any. For the
+ // current support in the OpenMP dialect, the table below lists
+ // all ops that have arguments (SSA operands and/or attributes)
+ // and not target-related. Of all these ops, we need to only
+ // process `omp.atomic.update` since it is the only op that has
+ // SSA operands & an attached region. Therefore, the region's
+ // entry block arguments must be mapped to the op's operands in
+ // case they are referenced inside the region.
+ //
+ // clang-format off
+ // =============================================================================
+ // | op | operands? | region(s)? | parent is func? | processed? |
+ // =============================================================================
+ // | atomic.read | yes | no | yes | no |
+ // | atomic.write | yes | no | yes | no |
+ // | atomic.update | yes | yes | yes | yes |
+ // | critical | no | no | yes | no |
+ // | declare_mapper | no | yes | no | no |
+ // | declare_reduction | no | yes | no | no |
+ // | flush | yes | no | yes | no |
+ // | private | no | yes | yes | no |
+ // | threadprivate | yes | no | yes | no |
+ // | yield | yes | no | yes | no |
+ // =============================================================================
+ // clang-format on
+ if (isa<mlir::omp::AtomicUpdateOp>(oper))
+ for (auto [operand, arg] :
+ llvm::zip_equal(oper->getOperands(),
+ oper->getRegion(0).getArguments())) {
+ moduleTranslation.mapValue(
+ arg, builder.CreateLoad(
+ moduleTranslation.convertType(arg.getType()),
+ moduleTranslation.lookupValue(operand)));
+ }
+ }
if (auto loopNest = dyn_cast<omp::LoopNestOp>(oper)) {
assert(builder.GetInsertBlock() &&
@@ -5340,9 +5377,10 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
// translation of the OpenMP construct being converted (e.g. no
// OpenMP runtime calls will be generated). We just need this to
// prepare the kernel invocation args.
+ SmallVector<llvm::PHINode *> phis;
auto result = convertOmpOpRegions(
region, oper->getName().getStringRef().str() + ".fake.region",
- builder, moduleTranslation);
+ builder, moduleTranslation, &phis);
if (failed(handleError(result, *oper)))
return WalkResult::interrupt();
diff --git a/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir b/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir
index 71a4c29eaf0aa..c618b68d52aaf 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir
@@ -97,6 +97,20 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true,
llvm.return
}
+ llvm.func @test_target_and_atomic_update(%x: !llvm.ptr, %expr : i32) {
+ omp.target {
+ omp.terminator
+ }
+
+ omp.atomic.update %x : !llvm.ptr {
+ ^bb0(%xval: i32):
+ %newval = llvm.add %xval, %expr : i32
+ omp.yield(%newval : i32)
+ }
+
+ llvm.return
+ }
+
// CHECK-LABEL: define void @test_nested_target_in_parallel_with_private({{.*}}) {
// CHECK: br label %omp.parallel.fake.region
// CHECK: omp.parallel.fake.region:
@@ -132,4 +146,11 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true,
// CHECK: call void @__kmpc_target_deinit()
// CHECK: ret void
// CHECK: }
+
+// CHECK-LABEL: define {{.*}} amdgpu_kernel void @__omp_offloading_{{.*}}_test_target_and_atomic_update_{{.*}} {
+// CHECK: call i32 @__kmpc_target_init
+// CHECK: user_code.entry:
+// CHECK: call void @__kmpc_target_deinit()
+// CHECK: ret void
+// CHECK: }
}
>From 533a643b245f465299c56d4817d8a88ec0d08d5c Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 20 Mar 2025 08:51:01 -0500
Subject: [PATCH 2/2] summarize comment
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 29 ++-----------------
1 file changed, 3 insertions(+), 26 deletions(-)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index cfc4e98ad3341..831d260bbc4bb 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5324,32 +5324,9 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
dyn_cast<omp::BlockArgOpenMPOpInterface>(oper))
forwardArgs(moduleTranslation, blockArgsIface);
else {
- // For non-BlockArgOpenMPOpInterface ops, we also want to map
- // their entry block arguments to their operands, if any. For the
- // current support in the OpenMP dialect, the table below lists
- // all ops that have arguments (SSA operands and/or attributes)
- // and not target-related. Of all these ops, we need to only
- // process `omp.atomic.update` since it is the only op that has
- // SSA operands & an attached region. Therefore, the region's
- // entry block arguments must be mapped to the op's operands in
- // case they are referenced inside the region.
- //
- // clang-format off
- // =============================================================================
- // | op | operands? | region(s)? | parent is func? | processed? |
- // =============================================================================
- // | atomic.read | yes | no | yes | no |
- // | atomic.write | yes | no | yes | no |
- // | atomic.update | yes | yes | yes | yes |
- // | critical | no | no | yes | no |
- // | declare_mapper | no | yes | no | no |
- // | declare_reduction | no | yes | no | no |
- // | flush | yes | no | yes | no |
- // | private | no | yes | yes | no |
- // | threadprivate | yes | no | yes | no |
- // | yield | yes | no | yes | no |
- // =============================================================================
- // clang-format on
+ // Here we map entry block arguments of
+ // non-BlockArgOpenMPOpInterface ops if they can be encountered
+ // inside of a function and they define any of these arguments.
if (isa<mlir::omp::AtomicUpdateOp>(oper))
for (auto [operand, arg] :
llvm::zip_equal(oper->getOperands(),
More information about the Mlir-commits
mailing list