Skip to content

Commit 134e3f6

Browse files
ergawyronlieb
authored andcommitted
[flang][OpenMP] Process omp.atomic.update while translating scopes 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 |
1 parent 3b052d2 commit 134e3f6

File tree

2 files changed

+37
-1
lines changed

2 files changed

+37
-1
lines changed

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5822,6 +5822,20 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
58225822
if (auto blockArgsIface =
58235823
dyn_cast<omp::BlockArgOpenMPOpInterface>(oper))
58245824
forwardArgs(moduleTranslation, blockArgsIface);
5825+
else {
5826+
// Here we map entry block arguments of
5827+
// non-BlockArgOpenMPOpInterface ops if they can be encountered
5828+
// inside of a function and they define any of these arguments.
5829+
if (isa<mlir::omp::AtomicUpdateOp>(oper))
5830+
for (auto [operand, arg] :
5831+
llvm::zip_equal(oper->getOperands(),
5832+
oper->getRegion(0).getArguments())) {
5833+
moduleTranslation.mapValue(
5834+
arg, builder.CreateLoad(
5835+
moduleTranslation.convertType(arg.getType()),
5836+
moduleTranslation.lookupValue(operand)));
5837+
}
5838+
}
58255839

58265840
if (auto loopNest = dyn_cast<omp::LoopNestOp>(oper)) {
58275841
assert(builder.GetInsertBlock() &&
@@ -5839,9 +5853,10 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
58395853
// translation of the OpenMP construct being converted (e.g. no
58405854
// OpenMP runtime calls will be generated). We just need this to
58415855
// prepare the kernel invocation args.
5856+
SmallVector<llvm::PHINode *> phis;
58425857
auto result = convertOmpOpRegions(
58435858
region, oper->getName().getStringRef().str() + ".fake.region",
5844-
builder, moduleTranslation);
5859+
builder, moduleTranslation, &phis);
58455860
if (failed(handleError(result, *oper)))
58465861
return WalkResult::interrupt();
58475862

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)