Skip to content

Commit 32e696f

Browse files
skatrakskatrak
committed
[MLIR][OpenMP] Support target SPMD
This patch implements MLIR to LLVM IR translation of host-evaluated loop bounds, completing initial support for `target teams distribute parallel do [simd]` and `target teams distribute [simd]`.
1 parent 033091e commit 32e696f

File tree

3 files changed

+159
-44
lines changed

3 files changed

+159
-44
lines changed

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

Lines changed: 63 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -173,15 +173,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
173173
if (op.getHint())
174174
op.emitWarning("hint clause discarded");
175175
};
176-
auto checkHostEval = [](auto op, LogicalResult &result) {
177-
// Host evaluated clauses are supported, except for loop bounds.
178-
for (BlockArgument arg :
179-
cast<omp::BlockArgOpenMPOpInterface>(*op).getHostEvalBlockArgs())
180-
for (Operation *user : arg.getUsers())
181-
if (isa<omp::LoopNestOp>(user))
182-
result = op.emitError("not yet implemented: host evaluation of loop "
183-
"bounds in omp.target operation");
184-
};
185176
auto checkInReduction = [&todo](auto op, LogicalResult &result) {
186177
if (!op.getInReductionVars().empty() || op.getInReductionByref() ||
187178
op.getInReductionSyms())
@@ -318,7 +309,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
318309
checkBare(op, result);
319310
checkDevice(op, result);
320311
checkHasDeviceAddr(op, result);
321-
checkHostEval(op, result);
322312
checkInReduction(op, result);
323313
checkIsDevicePtr(op, result);
324314
checkPrivate(op, result);
@@ -4058,9 +4048,13 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
40584048
///
40594049
/// Loop bounds and steps are only optionally populated, if output vectors are
40604050
/// provided.
4061-
static void extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
4062-
Value &numTeamsLower, Value &numTeamsUpper,
4063-
Value &threadLimit) {
4051+
static void
4052+
extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
4053+
Value &numTeamsLower, Value &numTeamsUpper,
4054+
Value &threadLimit,
4055+
llvm::SmallVectorImpl<Value> *lowerBounds = nullptr,
4056+
llvm::SmallVectorImpl<Value> *upperBounds = nullptr,
4057+
llvm::SmallVectorImpl<Value> *steps = nullptr) {
40644058
auto blockArgIface = llvm::cast<omp::BlockArgOpenMPOpInterface>(*targetOp);
40654059
for (auto item : llvm::zip_equal(targetOp.getHostEvalVars(),
40664060
blockArgIface.getHostEvalBlockArgs())) {
@@ -4085,11 +4079,26 @@ static void extractHostEvalClauses(omp::TargetOp targetOp, Value &numThreads,
40854079
llvm_unreachable("unsupported host_eval use");
40864080
})
40874081
.Case([&](omp::LoopNestOp loopOp) {
4088-
// TODO: Extract bounds and step values. Currently, this cannot be
4089-
// reached because translation would have been stopped earlier as a
4090-
// result of `checkImplementationStatus` detecting and reporting
4091-
// this situation.
4092-
llvm_unreachable("unsupported host_eval use");
4082+
auto processBounds =
4083+
[&](OperandRange opBounds,
4084+
llvm::SmallVectorImpl<Value> *outBounds) -> bool {
4085+
bool found = false;
4086+
for (auto [i, lb] : llvm::enumerate(opBounds)) {
4087+
if (lb == blockArg) {
4088+
found = true;
4089+
if (outBounds)
4090+
(*outBounds)[i] = hostEvalVar;
4091+
}
4092+
}
4093+
return found;
4094+
};
4095+
bool found =
4096+
processBounds(loopOp.getLoopLowerBounds(), lowerBounds);
4097+
found = processBounds(loopOp.getLoopUpperBounds(), upperBounds) ||
4098+
found;
4099+
found = processBounds(loopOp.getLoopSteps(), steps) || found;
4100+
if (!found)
4101+
llvm_unreachable("unsupported host_eval use");
40934102
})
40944103
.Default([](Operation *) {
40954104
llvm_unreachable("unsupported host_eval use");
@@ -4226,6 +4235,7 @@ initTargetDefaultAttrs(omp::TargetOp targetOp,
42264235
combinedMaxThreadsVal = maxThreadsVal;
42274236

42284237
// Update kernel bounds structure for the `OpenMPIRBuilder` to use.
4238+
attrs.ExecFlags = targetOp.getKernelExecFlags();
42294239
attrs.MinTeams = minTeamsVal;
42304240
attrs.MaxTeams.front() = maxTeamsVal;
42314241
attrs.MinThreads = 1;
@@ -4243,9 +4253,15 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
42434253
LLVM::ModuleTranslation &moduleTranslation,
42444254
omp::TargetOp targetOp,
42454255
llvm::OpenMPIRBuilder::TargetKernelRuntimeAttrs &attrs) {
4256+
omp::LoopNestOp loopOp = castOrGetParentOfType<omp::LoopNestOp>(
4257+
targetOp.getInnermostCapturedOmpOp());
4258+
unsigned numLoops = loopOp ? loopOp.getNumLoops() : 0;
4259+
42464260
Value numThreads, numTeamsLower, numTeamsUpper, teamsThreadLimit;
4261+
llvm::SmallVector<Value> lowerBounds(numLoops), upperBounds(numLoops),
4262+
steps(numLoops);
42474263
extractHostEvalClauses(targetOp, numThreads, numTeamsLower, numTeamsUpper,
4248-
teamsThreadLimit);
4264+
teamsThreadLimit, &lowerBounds, &upperBounds, &steps);
42494265

42504266
// TODO: Handle constant 'if' clauses.
42514267
if (Value targetThreadLimit = targetOp.getThreadLimit())
@@ -4265,7 +4281,34 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
42654281
if (numThreads)
42664282
attrs.MaxThreads = moduleTranslation.lookupValue(numThreads);
42674283

4268-
// TODO: Populate attrs.LoopTripCount if it is target SPMD.
4284+
if (targetOp.getKernelExecFlags() != llvm::omp::OMP_TGT_EXEC_MODE_GENERIC) {
4285+
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
4286+
attrs.LoopTripCount = nullptr;
4287+
4288+
// To calculate the trip count, we multiply together the trip counts of
4289+
// every collapsed canonical loop. We don't need to create the loop nests
4290+
// here, since we're only interested in the trip count.
4291+
for (auto [loopLower, loopUpper, loopStep] :
4292+
llvm::zip_equal(lowerBounds, upperBounds, steps)) {
4293+
llvm::Value *lowerBound = moduleTranslation.lookupValue(loopLower);
4294+
llvm::Value *upperBound = moduleTranslation.lookupValue(loopUpper);
4295+
llvm::Value *step = moduleTranslation.lookupValue(loopStep);
4296+
4297+
llvm::OpenMPIRBuilder::LocationDescription loc(builder);
4298+
llvm::Value *tripCount = ompBuilder->calculateCanonicalLoopTripCount(
4299+
loc, lowerBound, upperBound, step, /*IsSigned=*/true,
4300+
loopOp.getLoopInclusive());
4301+
4302+
if (!attrs.LoopTripCount) {
4303+
attrs.LoopTripCount = tripCount;
4304+
continue;
4305+
}
4306+
4307+
// TODO: Enable UndefinedSanitizer to diagnose an overflow here.
4308+
attrs.LoopTripCount = builder.CreateMul(attrs.LoopTripCount, tripCount,
4309+
{}, /*HasNUW=*/true);
4310+
}
4311+
}
42694312
}
42704313

42714314
static LogicalResult

‎mlir/test/Target/LLVMIR/openmp-target-spmd.mlir

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
// RUN: split-file %s %t
2+
// RUN: mlir-translate -mlir-to-llvmir %t/host.mlir | FileCheck %s --check-prefix=HOST
3+
// RUN: mlir-translate -mlir-to-llvmir %t/device.mlir | FileCheck %s --check-prefix=DEVICE
4+
5+
//--- host.mlir
6+
7+
module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-amd-amdhsa"]} {
8+
llvm.func @main(%x : i32) {
9+
omp.target host_eval(%x -> %lb, %x -> %ub, %x -> %step : i32, i32, i32) {
10+
omp.teams {
11+
omp.parallel {
12+
omp.distribute {
13+
omp.wsloop {
14+
omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
15+
omp.yield
16+
}
17+
} {omp.composite}
18+
} {omp.composite}
19+
omp.terminator
20+
} {omp.composite}
21+
omp.terminator
22+
}
23+
omp.terminator
24+
}
25+
llvm.return
26+
}
27+
}
28+
29+
// HOST-LABEL: define void @main
30+
// HOST: %omp_loop.tripcount = {{.*}}
31+
// HOST-NEXT: br label %[[ENTRY:.*]]
32+
// HOST: [[ENTRY]]:
33+
// HOST-NEXT: %[[TRIPCOUNT:.*]] = zext i32 %omp_loop.tripcount to i64
34+
// HOST: %[[TRIPCOUNT_KARG:.*]] = getelementptr inbounds nuw %struct.__tgt_kernel_arguments, ptr %[[KARGS:.*]], i32 0, i32 8
35+
// HOST-NEXT: store i64 %[[TRIPCOUNT]], ptr %[[TRIPCOUNT_KARG]]
36+
// HOST: %[[RESULT:.*]] = call i32 @__tgt_target_kernel({{.*}}, ptr %[[KARGS]])
37+
// HOST-NEXT: %[[CMP:.*]] = icmp ne i32 %[[RESULT]], 0
38+
// HOST-NEXT: br i1 %[[CMP]], label %[[OFFLOAD_FAILED:.*]], label %{{.*}}
39+
// HOST: [[OFFLOAD_FAILED]]:
40+
// HOST: call void @[[TARGET_OUTLINE:.*]]({{.*}})
41+
42+
// HOST: define internal void @[[TARGET_OUTLINE]]
43+
// HOST: call void{{.*}}@__kmpc_fork_teams({{.*}}, ptr @[[TEAMS_OUTLINE:.*]], {{.*}})
44+
45+
// HOST: define internal void @[[TEAMS_OUTLINE]]
46+
// HOST: call void{{.*}}@__kmpc_fork_call({{.*}}, ptr @[[PARALLEL_OUTLINE:.*]], {{.*}})
47+
48+
// HOST: define internal void @[[PARALLEL_OUTLINE]]
49+
// HOST: call void @[[DISTRIBUTE_OUTLINE:.*]]({{.*}})
50+
51+
// HOST: define internal void @[[DISTRIBUTE_OUTLINE]]
52+
// HOST: call void @__kmpc_dist_for_static_init{{.*}}(ptr {{.*}}, i32 {{.*}}, i32 34, ptr {{.*}}, ptr {{.*}}, ptr {{.*}}, ptr {{.*}}, ptr {{.*}}, i32 {{.*}}, i32 {{.*}})
53+
54+
//--- device.mlir
55+
56+
module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_target_device = true, omp.is_gpu = true} {
57+
llvm.func @main(%x : i32) {
58+
omp.target host_eval(%x -> %lb, %x -> %ub, %x -> %step : i32, i32, i32) {
59+
omp.teams {
60+
omp.parallel {
61+
omp.distribute {
62+
omp.wsloop {
63+
omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
64+
omp.yield
65+
}
66+
} {omp.composite}
67+
} {omp.composite}
68+
omp.terminator
69+
} {omp.composite}
70+
omp.terminator
71+
}
72+
omp.terminator
73+
}
74+
llvm.return
75+
}
76+
}
77+
78+
// DEVICE: @[[KERNEL_NAME:.*]]_exec_mode = weak protected constant i8 2
79+
// DEVICE: @llvm.compiler.used = appending global [1 x ptr] [ptr @[[KERNEL_NAME]]_exec_mode], section "llvm.metadata"
80+
// DEVICE: @[[KERNEL_NAME]]_kernel_environment = weak_odr protected constant %struct.KernelEnvironmentTy {
81+
// DEVICE-SAME: %struct.ConfigurationEnvironmentTy { i8 0, i8 1, i8 [[EXEC_MODE:2]], {{.*}}},
82+
// DEVICE-SAME: ptr @{{.*}}, ptr @{{.*}} }
83+
84+
// DEVICE: define weak_odr protected amdgpu_kernel void @[[KERNEL_NAME]]({{.*}})
85+
// DEVICE: %{{.*}} = call i32 @__kmpc_target_init(ptr @[[KERNEL_NAME]]_kernel_environment, {{.*}})
86+
// DEVICE: call void @[[TARGET_OUTLINE:.*]]({{.*}})
87+
// DEVICE: call void @__kmpc_target_deinit()
88+
89+
// DEVICE: define internal void @[[TARGET_OUTLINE]]({{.*}})
90+
// DEVICE: call void @__kmpc_parallel_51(ptr {{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}}, ptr @[[PARALLEL_OUTLINE:.*]], ptr {{.*}}, ptr {{.*}}, i64 {{.*}})
91+
92+
// DEVICE: define internal void @[[PARALLEL_OUTLINE]]({{.*}})
93+
// DEVICE: call void @[[DISTRIBUTE_OUTLINE:.*]]({{.*}})
94+
95+
// DEVICE: define internal void @[[DISTRIBUTE_OUTLINE]]({{.*}})
96+
// DEVICE: call void @__kmpc_distribute_for_static_loop{{.*}}({{.*}})

‎mlir/test/Target/LLVMIR/openmp-todo.mlir

Lines changed: 0 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -319,30 +319,6 @@ llvm.func @target_has_device_addr(%x : !llvm.ptr) {
319319

320320
// -----
321321

322-
llvm.func @target_host_eval(%x : i32) {
323-
// expected-error@below {{not yet implemented: host evaluation of loop bounds in omp.target operation}}
324-
// expected-error@below {{LLVM Translation failed for operation: omp.target}}
325-
omp.target host_eval(%x -> %lb, %x -> %ub, %x -> %step : i32, i32, i32) {
326-
omp.teams {
327-
omp.parallel {
328-
omp.distribute {
329-
omp.wsloop {
330-
omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
331-
omp.yield
332-
}
333-
} {omp.composite}
334-
} {omp.composite}
335-
omp.terminator
336-
} {omp.composite}
337-
omp.terminator
338-
}
339-
omp.terminator
340-
}
341-
llvm.return
342-
}
343-
344-
// -----
345-
346322
omp.declare_reduction @add_f32 : f32
347323
init {
348324
^bb0(%arg: f32):

0 commit comments

Comments
 (0)