[Mlir-commits] [mlir] [OpenMP][MLIR] Add omp.distribute op to the OMP dialect (PR #67720)

[Jan Leyonberg]

jsjodin wrote:

> > > Since the canonical loop proposal might take some time before it's finished, in my opinion a good approach could be to mimic `omp.wsloop` and produce new loop index variables as entry block arguments to this op's region. I guess in that case we'd have to think about what the `omp.distribute` loop bounds and steps would be for a given nested Do loop, and how they might interoperate with the corresponding `omp.wsloop` operation.
> > 
> > 
> > Are you suggesting a different approach here? To have `omp.distribute` the same as `omp.wsloop`, i.e a loop-like operation? Currently `omp.wsloop` subsumes the nested `Do loop`.
> 
> I am proposing a different approach. The way I see it it makes sense that the omp.distribute is a wrapper op because it does not affect the code of the loop only the "schedule" of the loop, either run all iterations for each team, or a portion of the iterations per team.
> 
> > In the lowering in #67798, an `omp.wsloop` is created for `!$omp distribute`. Is this always correct to do as per the standard? Is it based on existing Clang lowering?
> 
> I think a omp.wsloop is created from parallel do, Afaik 'distribute' is associated with teams, so that each team takes a chunk of the iterations instead of all teams taking all (duplicating) iterations.
> 

I am backtracking a bit on this, The reason for this is that unless there is some value that is passed from the wrapping op to the operations inside the inner there is nothing that would prevent MLIR optimizations from hoisting code outside the region, unless we add some side effect to the op. I think the better approach is to have a new op that looks like omp.wsloop, or to add an attribute to the existing omp.wsloop to indicate that it is a distribute workshare loop. I don't know if there is a preference for either?

Example user code:
````
!$omp target teams distribute parallel do
do i=1,10
  do j=1, 10
  end do
end do
!$omp end target teams distribute parallel do
````
 Will be translated to the MLIR code (new op):
````
omp.target {
  omp.teams {
    omp.parallel {
      omp.dwsloop (...) {}
    }
  }
}
````
or (wsloop with flag):
````
omp.target {
  omp.teams {
    omp.parallel {
      omp.wsloop (type=DistributeFor) {}
    }
  }
}
````
 

https://github.com/llvm/llvm-project/pull/67720