[MLIR][OpenMP] Should omp.target/target_data implement RegionBranchOpInterface?

MLIR
,
1

I am currently implementing a forward dense dataflow analysis over omp.target and omp.target_data, and I ran into a modeling issue.

It seems common for OpenMP lowering to produce nested MLIR structures such as:

omp.target_data map_entries(...) {
  omp.target map_entries(...) {
    ...
    omp.terminator
  }
  ...
  omp.terminator
}

For standalone mapping operations such as omp.target_enter_data, it is straightforward to model the transfer function using mlir::dataflow::DenseForwardDataFlowAnalysis::visitOperation.

However, for omp.target_data and omp.target, I would like to model the dataflow across the region boundaries. Conceptually, the transfer I want is:

before parent op -> region entry:
  apply map-entry semantics

region exit -> after parent op:
  apply map-exit semantics

This seems close to what DenseForwardDataFlowAnalysis::visitRegionBranchControlFlowTransfer is designed for. However, that mechanism requires the operation to implement RegionBranchOpInterface, and omp.target and omp.target_data currently do not implement that interface.

Would it make sense for omp.target and omp.target_data to implement RegionBranchOpInterface? Or is there another recommended way to model this kind of structured region entry/exit effect in a dense dataflow analysis over the OpenMP dialect?

2

Thank you for the proposal @MuyaoXiao, and excuse the delay taking a look at this. I’m no expert on the MLIR dataflow analysis framework, so do let me know if I’m misunderstanding or missing anything. I’ve been investigating this for a bit to hopefully help with some ideas.

From what I can tell, it seems like you’re right about the need for implementing the RegionBranchOpInterface if we want to enable dataflow analyses across some operation’s region boundaries without resetting to the pessimistic defaults.

There are some characteristics of omp.target that probably make it not very straightforward to handle, though. The issues there are: i) target regions run asynchronously if they have the nowait attribute, in which case an analysis that assumes that “exiting the omp.target region means immediately triggering all the map-exit semantics” would be wrong; and ii) their region executes in a separate device context independent of any analysis data that might have been gathered for the host.

What I’m thinking is that, in addition to implementing a simple RegionBranchOpInterface for omp.target representing its region as if executing inline with only host_eval values passed along as entry successor region operands (the rest of entry block arguments represent target region private copies there, not outside aliases), it probably would make sense to introduce the MemoryEffectOpInterface and the RecursiveMemoryEffects trait for it to model the memory transfers too. The former would be the one responsible for modeling map- and firstprivate-related reads and writes, and possibly forcing a pessimistic read+write effect on every mapped value to prevent invalid order-based optimizations in the nowait case, and the latter would make sure that e.g. map-less target regions wouldn’t get DCE’d when they had side-effecting device code in them. The SideEffects::EffectInstance stage could probably be set to tag region entry and exit memory effects separately, if your analysis needs to handle them separately.

I think that this should let a forward dense dataflow analysis go though omp.target while keeping track of host memory reads/writes (and also keep the device-side analysis of each target region isolated), and be pessimistically handled whenever nowait is present. Your analysis could query these data movements by looking at the memory effects of the relevant operation.

omp.target_data is simpler in this case, as it only contains host code and it doesn’t take a nowait clause or host_eval operands. The RegionBranchOpInterface and MemoryEffectOpInterface would end up being simplified versions of those for omp.target. I don’t know if perhaps it would make sense in your case for ops such as omp.target_enter_data to also properly state their memory effects.

I think we’re currently quite conservative in the omp dialect in general with regards to memory effects (relatively few operations actually define any), and this seems like a situation where that is shielding us from potential useful optimizations / analyses. Other dialect ops will still trip these up and be excessively conservative, but I guess we have to start somewhere with properly defining memory effects. However, we need to be really careful while doing it, because the current sub-optimal situation is preferable to a slightly wrong modeling of memory effects causing illegal transformations.

Do let me know, though, if you had a simpler approach in mind or if perhaps you think going into the full memory effect modeling path for this isn’t a good idea.

3

Thanks for the explanation, I think now I understand the distinction more clearly.
My original motivation was limited to the requirements of my own analyses, without sufficiently considering how generic analyses and transformations would interpret the interface.
In my analysis, I was particularly interested in map_entries. When entering an omp.target region, I created a a separate lattice entry for each mapped region argument. The entry state is initialized by explicitly interpreting the corresponding omp.map_info. Operations inside the target region modify the new lattice entry rather than the host-side entry.
When exiting the target region, the analysis again inspects the corresponding omp.map_info to decide whether and how the resulting state should be propagated back to the host-side lattice entry. It depends on the map type and capture type, as well as the reference count range recorded in my lattice. My analysis currently tracks reference counts only for temporary variables declared within a subroutine; reference counts for other variables are treated as unknown and are therefore handled conservatively.

My remaining concern is that actual transfer may also depend on the exiting device mapping and its reference count, which can be accumulated and reduced across omp.target_eneter_data, omp.target_exit_data, and the entry and exit edges of omp.target and omp.target_data. I do not think the behavior can be represented precisely by MemoryEffectOpInterface alone.

My current understanding is that the three mechanisms have complementary roles:

  • RegionBranchOpInterface provides the structure control flow and host_eval forwarding. It also makes it possible for OpenMP specific data flow analysis to implement the more detailed transfer semantics for mapped structures.
  • MemroyEffectOpInterface could provide conservative may-read and may-write effects for generic analyses.
  • RecursiveMemoryEffects should accompany MemoryEffectOpInterface so that effects from operations inside target region are not hidden.

Would it therefore be sound to implement RegionBranchOpInterface independently, with only host_eval represented as forwarded successor operands?
Or do you think MemoryEffectOpInterface and RecursiveMemoryEffetcs could be introduced in the same time?

Please feel free to correct me if any part of my understanding is inaccurate.

4

Thank you again, I don’t think I have any big concerns about your proposal in this RFC.

I think it should still be safe for these ops to remain conservative in terms of memory effects, and keep the current optimization-preventing behavior, while properly modeling how control flows through their regions. So I wouldn’t be against adding that interface on its own. In my opinion, you should feel free to make a PR for this.

I think I understand now what you’re trying. If the purpose of your analysis is to calculate the potential runtime reference counts for mapped variables to figure out when memory movements might trigger, then it does seem like something different to what we could model with memory effects. We can think about memory effects modeling independently later on.

I guess this is the main thing I’d be careful with. If I understand it correctly, this causes your analysis to assume that all updates to memory performed by the target region are completed before the next operation. However, this would break down for a nowait target region that e.g. takes a long time to run or if it depends on another task to complete before it gets launched. In that case, only reaching a barrier (explicit or implicit) would guarantee that these memory updates are effective.

5

Thanks, and sorry for a delayed reply.

I agree nowait is the main complication.

If I understand it correctly, this causes your analysis to assume that all updates to memory performed by the target region are completed before the next operation.

Exactly, and therefore in my analysis, I will need to handle nowait conservatively.

I have created a draft PR that only adds the region control-flow interfaces and forwards host_eval values for omp.target, without modeling memory effects or map-entry/exit semantics: