[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 01:44:27 PDT 2025
https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/132165
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 |
>From db9c57c743e3bc3685cd17df0e5f6fc7e39c0e8a 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] [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: }
}
More information about the Mlir-commits
mailing list