Skip to content

Commit 2525026

Browse files
ergawyergawy
committed
[flang][OpenMP] Process omp.atomic.update while translating scopes 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 | =============================================================================
1 parent 9b1f905 commit 2525026

File tree

2 files changed

+60
-1
lines changed

2 files changed

+60
-1
lines changed

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

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5323,6 +5323,43 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
53235323
if (auto blockArgsIface =
53245324
dyn_cast<omp::BlockArgOpenMPOpInterface>(oper))
53255325
forwardArgs(moduleTranslation, blockArgsIface);
5326+
else {
5327+
// For non-BlockArgOpenMPOpInterface ops, we also want to map
5328+
// their entry block arguments to their operands, if any. For the
5329+
// current support in the OpenMP dialect, the table below lists
5330+
// all ops that have arguments (SSA operands and/or attributes)
5331+
// and not target-related. Of all these ops, we need to only
5332+
// process `omp.atomic.update` since it is the only op that has
5333+
// SSA operands & an attached region. Therefore, the region's
5334+
// entry block arguments must be mapped to the op's operands in
5335+
// case they are referenced inside the region.
5336+
//
5337+
// clang-format off
5338+
// =============================================================================
5339+
// | op | operands? | region(s)? | parent is func? | processed? |
5340+
// =============================================================================
5341+
// | atomic.read | yes | no | yes | no |
5342+
// | atomic.write | yes | no | yes | no |
5343+
// | atomic.update | yes | yes | yes | yes |
5344+
// | critical | no | no | yes | no |
5345+
// | declare_mapper | no | yes | no | no |
5346+
// | declare_reduction | no | yes | no | no |
5347+
// | flush | yes | no | yes | no |
5348+
// | private | no | yes | yes | no |
5349+
// | threadprivate | yes | no | yes | no |
5350+
// | yield | yes | no | yes | no |
5351+
// =============================================================================
5352+
// clang-format on
5353+
if (isa<mlir::omp::AtomicUpdateOp>(oper))
5354+
for (auto [operand, arg] :
5355+
llvm::zip_equal(oper->getOperands(),
5356+
oper->getRegion(0).getArguments())) {
5357+
moduleTranslation.mapValue(
5358+
arg, builder.CreateLoad(
5359+
moduleTranslation.convertType(arg.getType()),
5360+
moduleTranslation.lookupValue(operand)));
5361+
}
5362+
}
53265363

53275364
if (auto loopNest = dyn_cast<omp::LoopNestOp>(oper)) {
53285365
assert(builder.GetInsertBlock() &&
@@ -5340,9 +5377,10 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
53405377
// translation of the OpenMP construct being converted (e.g. no
53415378
// OpenMP runtime calls will be generated). We just need this to
53425379
// prepare the kernel invocation args.
5380+
SmallVector<llvm::PHINode *> phis;
53435381
auto result = convertOmpOpRegions(
53445382
region, oper->getName().getStringRef().str() + ".fake.region",
5345-
builder, moduleTranslation);
5383+
builder, moduleTranslation, &phis);
53465384
if (failed(handleError(result, *oper)))
53475385
return WalkResult::interrupt();
53485386

‎mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,20 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true,
9797
llvm.return
9898
}
9999

100+
llvm.func @test_target_and_atomic_update(%x: !llvm.ptr, %expr : i32) {
101+
omp.target {
102+
omp.terminator
103+
}
104+
105+
omp.atomic.update %x : !llvm.ptr {
106+
^bb0(%xval: i32):
107+
%newval = llvm.add %xval, %expr : i32
108+
omp.yield(%newval : i32)
109+
}
110+
111+
llvm.return
112+
}
113+
100114
// CHECK-LABEL: define void @test_nested_target_in_parallel_with_private({{.*}}) {
101115
// CHECK: br label %omp.parallel.fake.region
102116
// CHECK: omp.parallel.fake.region:
@@ -132,4 +146,11 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true,
132146
// CHECK: call void @__kmpc_target_deinit()
133147
// CHECK: ret void
134148
// CHECK: }
149+
150+
// CHECK-LABEL: define {{.*}} amdgpu_kernel void @__omp_offloading_{{.*}}_test_target_and_atomic_update_{{.*}} {
151+
// CHECK: call i32 @__kmpc_target_init
152+
// CHECK: user_code.entry:
153+
// CHECK: call void @__kmpc_target_deinit()
154+
// CHECK: ret void
155+
// CHECK: }
135156
}

0 commit comments

Comments
 (0)