[flang][OpenMP] Process omp.atomic.update while translating scopes for target device
#132165
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-openmp Author: Kareem Ergawy (ergawy) ChangesFixes a bug introduced by #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
Full diff: https://github.com/llvm/llvm-project/pull/132165.diff 2 Files Affected:
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: }
}
|
ronlieb
approved these changes
Mar 20, 2025
skatrak
approved these changes
Mar 20, 2025
Comment on lines
5327
to
5352
| // 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 |
There was a problem hiding this comment.
I think having this detailed description of how things are at this point in time is very good information for the commit message. However, I don't think we want this as a comment in the code because it can very easily become outdated. The comment could be summarized to something along the lines of "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.".
| auto result = convertOmpOpRegions( | ||
| region, oper->getName().getStringRef().str() + ".fake.region", | ||
| builder, moduleTranslation); | ||
| builder, moduleTranslation, &phis); |
There was a problem hiding this comment.
I see phis isn't used. Does passing the vector on its own cause a behavior change?
There was a problem hiding this comment.
Yes, since omp.atomic.update yields values (which are not used in this context), convertOmpOpRegions asserts unless we pass a vector to collect the phis.
There was a problem hiding this comment.
Ah I see, and I guess not using these PHI nodes later doesn't cause issues. Thanks for the explanation!
…for target device Fixes a bug introduced by #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 | =============================================================================
ergawy
force-pushed
the
users/ergawy/translate_atomic_update
branch
from
db9c57c to
533a643
Compare
dpalermo
approved these changes
Mar 20, 2025
searlmc1
pushed a commit
to ROCm/llvm-project
that referenced
this pull request
Mar 21, 2025
…for target device (llvm#132165) Fixes a bug introduced by llvm#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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fixes a bug introduced by #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.updatesince 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.