[MLIR][OpenMP] Lower only target related pragmas for GPU MLIR

[DominikAdamski]

Flang generates two MLIR files for target-related pragmas: the host MLIR and the GPU MLIR. If a given source function contains both host and target-related pragmas, we must ensure that we do not lower the host-related pragmas for the GPU MLIR.

[skatrak]

Thank you Dominik for working on this. It does look close to what I was thinking could be the solution. I added a few comments if we go ahead with this approach. However, your check for this type of pattern made me think:

omp.parallel {
  omp.target { }
}

This was something I hadn't thought about until now. With this approach, we can't support this kind of pattern because we either check it and error out (what you've implemented) or we ignore the nested omp.target during device compilation because we don't lower the contents of omp.parallel as a result of it being a host operation (if we remove that check).

It seems like what we really want to do for the device is to lower the contents of any omp.target, regardless of whether it appears inside some other host OpenMP constructs, in addition to the applicable declare target functions. Maybe the proper solution would be to remove the check for a nested omp.target and, rather than just returning from convertOperation if the helper function returns false, to instead trigger codegen for all directly or indirectly nested omp.target ops.

I guess this would be problematic, because there may be some values defined inside the host ops that we wouldn't lower and that might be mapped, possibly breaking stuff during lowering. It's a tricky one, maybe early outlining really is the best long term option and this is the best we can do in the meantime.

[skatrak]

Nit: Saying an operation is "not allowed" to be lowered I think implies that it should be an error, which is not the intent. Maybe we can rename this to something more along the lines of "operation should be skipped" or "operation not applicable" or something like that. It doesn't have to be the name I suggest, but something along those lines.

Suggested change

	static bool isOpAllowedToBeLowered(Operation *opInst,
	static bool shouldSkipOpLowering(Operation *opInst,

If you agree to make this name change, note that the logic would be reversed.

[skatrak]

I guess this should be checked before the previous check that looks for a parent omp.target. Otherwise, this wouldn't trigger the error:

omp.target {
  omp.teams {
    omp.target { }
  }
}

That kind of thing might eventually appear for reverse-offload regions, though we're a long way off from supporting that yet.

[bhandarkar-pranav]

Question for @skatrak - What is the meaning of the a omp.target nested inside another omp.target? Is it a NOP (I am guessing not since you refer to "reverse offloading")

[skatrak]

Right, what happens when you enter an OpenMP target region is that execution moves to a target device. If a target device encounters a target region annotated with the ancestor device modifier, then this region must run in the host and the device would wait for it to finish before resuming. This is called reverse offloading, which is not implemented either in flang or clang as far as I know. Sections 13.2 and 13.8 of the 5.2 spec explain this if you'd like a bit more info about it.

Since it's not implemented yet, we can't be 100% sure of how we're going to represent this in MLIR, but I think it's likely that we would just follow the way it's done in the source and get something similar to this:

// runs in the host...
omp.target {
  // runs in the device...
  omp.target device(ancestor = 1) {
    // runs in the host...
  }
  // continues running in the device...
}
// continues running in the host...

This is where I think nested omp.target ops could come from in the future. Not sure if it's possible to create more target regions within a reverse-offload region, but if that's the case I guess it would be possible to have a multi-level nest of target regions at some point, interleaving host and device code.

[bhandarkar-pranav]

Thanks, I hope applications that do this have kernels that are big enough to make it worthwhile to do all the shuttling between host and device (or the cost of this shuttling is kept really low through zero-copy buffers and such things but still there can be such a thing as too much shuttling I think).

[skatrak]

Maybe the default if it's not inside of a function should be to lower it rather than ignoring it.

Suggested change

	return false;
	return true;

[skatrak]

Suggested change

	ompBuilder->Config.isGPU()) {
	ompBuilder->Config.isTargetDevice()) {

[skatrak]

I think this wouldn't be necessary, since the original code doesn't check this.

[skatrak]

Suggested change

	// Skip lowering of an OpenMP operation if it's context is not appropriate
	// Skip lowering of an OpenMP operation if its context is not appropriate.