Problem statement
According to the OpenMP specification, the num_teams and thread_limit clauses on teams constructs combined or nested inside of a target construct are evaluated on the host device on entry to the target construct. Additionally, the loop trip count for target SPMD kernels (target teams distribute parallel {do,for} or semantically equivalent nesting of constructs) is also expected to be evaluated in the host so that the OpenMP runtime kernel launch can be set up properly.
Even though I have not been able to find explicit mentions for this in the specification, clang also evaluates the num_threads clause of the parallel construct on the host prior to launching target SPMD kernels as well, similarly to what is done for thread_limit.
The issue is that the omp.target operation is IsolatedFromAbove, so the following representation would not be legal due to live-ins present in the target region:
// Initialize %0, %1, %2, %3
// ...
omp.target thread_limit(%0 : i32) {
omp.teams num_teams(%1 : i32) thread_limit(%2 : i32) {
omp.parallel num_threads(%3 : i32) {
omp.distribute {
omp.wsloop {
omp.loop_nest ... {
...
omp.yield
}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
}
omp.terminator
}
Furthermore, the loop trip count is implicit in the current representation of omp.loop_nest, so for it to be evaluated on the host, bounds and step variables would also have to be accessible in the host in the same way.
Potential solutions
This RFC is to decide on a representation at the dialect level to address these situations. I have thought of a few options, but I’d be open to other alternatives as well.
Attaching information to omp.target
This alternative consists in adding arguments for each of these pieces of data to the omp.target operation. Doing this would make the MLIR to LLVM IR translation trivial, because it ensures these values are lowered outside of the target region, in the host, and it is very easy to identify the value that represents each clause and the trip count. However, it also results in this operation holding information that applies to another directive, breaking one of the main design rules of the dialect.
Since the thread_limit clause can be attached to both omp.teams and omp.target, for this approach we would have to either introduce a teams_thread_limit clause to avoid collisions when it is specified for both constructs or force users creating the MLIR representation to address any conflicting values before setting a single value in the omp.target operation.
// Initialize %0, %1, %2, %3
// ...
omp.target num_teams(%0 : i32) num_threads(%1 : i32) thread_limit(%2 : i32)
teams_thread_limit(%3 : i32) {
omp.teams {
omp.parallel {
omp.distribute {
omp.wsloop {
omp.loop_nest ... {
...
omp.yield
}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
}
omp.terminator
}
Map values
One way of ensuring certain values present in the host are available inside of the target region is through the use of the map clause. New mapped variables could be added to the omp.target operation to be able to refer to them for these host-evaluated clauses. From the MLIR creation perspective, this is not much different from the previous alternative, although it would require the addition of omp.map.info operations, as well as storing / loading from pointer-like types. Once delayed privatization support for omp.target is implemented, it might be possible to simplify this by using firstprivate clauses instead.
One of the main disadvantages of this approach is that, unless some cleanup pass is introduced with it, it will result in variables being mapped that are not needed inside of the target region. It is also more difficult when translating the omp.target operation to LLVM IR for the host to track which values represent each of these clauses.
// Initialize %0, %1, %2, %3
// ...
omp.target map_entries(%0 -> %arg0, %1 -> %arg1, %2 -> %arg2 : !llvm.ptr, !llvm.ptr, !llvm.ptr)
thread_limit(%3 : i32) {
%num_teams = llvm.load %arg0 : !llvm.ptr -> i32
%thread_limit = llvm.load %arg1 : !llvm.ptr -> i32
%num_threads = llvm.load %arg2 : !llvm.ptr -> i32
omp.teams num_teams(%num_teams : i32) thread_limit(%thread_limit : i32) {
omp.parallel num_threads(%num_threads : i32) {
omp.distribute {
omp.wsloop {
omp.loop_nest ... {
...
omp.yield
}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
}
omp.terminator
}
Add passthrough map-style argument to omp.target
This approach is similar to the previous one, but instead of introducing mappings for variables that might not actually be used inside of the target region, it would consist in creating a new map-like passthrough argument. This would only be used for cases where we need to be able to match a host value to a clause for an operation inside of the target region to be evaluated in the host. It would be illegal to use these values for any other purpose, since they would not exist inside of the target region.
Even though this would also require some work at the MLIR to LLVM IR level to match each clause in a nested operation to be evaluated in the host with its host value, it would be potentially simpler to do because there would be no local allocations or load operations, etc. to jump through.
// Initialize %0, %1, %2, %3
// ...
omp.target passthrough(%0 -> %arg0, %1 -> %arg1, %2 -> %arg2 : i32, i32, i32)
thread_limit(%3 : i32) {
omp.teams num_teams(%arg0 : i32) thread_limit(%arg1 : i32) {
omp.parallel num_threads(%arg2 : i32) {
omp.distribute {
omp.wsloop {
omp.loop_nest ... {
...
omp.yield
}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
}
omp.terminator
}
Keep MLIR representation unchanged
Another possibility is to make transformations during the MLIR to LLVM IR translation stage by cloning or hoisting the initialization of these clauses prior to the kernel launch call, while avoiding the introduction of any changes to the MLIR representation. The main advantage of this is that the MLIR representation remains the same, making it easier for users of the dialect. However, the initialization of these clauses would exist inside of the target region (and other nested operations), which would not match where that is actually done.
The fact these initializations would be based on mapped or private values of the target region, possibly shared with some other operations, would make the MLIR to LLVM IR translation harder to implement and maintain, and easier to break.
// Initialize %0
// ...
omp.target thread_limit(%0 : i32) {
// Initialize %1, %2
// ...
omp.teams num_teams(%1 : i32) thread_limit(%2 : i32) {
// Initialize %3
// ...
omp.parallel num_threads(%3 : i32) {
omp.distribute {
omp.wsloop {
omp.loop_nest ... {
...
omp.yield
}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
} {omp.composite}
omp.terminator
}
omp.terminator
}
Handling trip count
The trip count is a property of a collapsed loop nest that must be able to be evaluated in advance to executing the loop. This restriction, in the case of a target SPMD loop, extends to being evaluated outside of the target region, on the host. This enables the trip count to be passed to the runtime when launching a kernel.
For this to be achieved, a solution based on what is decided with regards to the num_teams and thread_limit clauses should probably be followed as well, unless having different approaches makes more sense. The peculiarity of this case is that the trip count is currently not explicitly represented like these clauses are, but instead it can be calculated from the set of bounds and steps of the omp.loop_nest operation.
In that case, the decision to make would be whether to create an argument to hold the calculated value of the trip count, placed in the omp.target or omp.loop_nest depending on the approach taken, and add MLIR code before the omp.target operation to calculate it from the bounds and steps. Alternatively, the whole list of bounds and steps could be added as arguments or mapped/forwarded through the omp.target operation. The first option has the disadvantage of adding redundancy, so it creates the potential of the trip count argument not matching the number of iterations represented by the omp.loop_nest. The second option would introduce the need for more values to be mapped/forwarded or arguments being added to omp.target, depending on the approach taken.
If the decision is to keep the MLIR representation unchanged, there would be no need for a decision on this point.
Initial thoughts
At this time, I think the passthrough approach and forwarding bounds and step in place of the trip count seems to be the best tradeoff between MLIR representation clarity, MLIR creation complexity and MLIR to LLVM IR lowering complexity, but I would like to hear your thoughts.