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

[Kiran Chandramohan]

kiranchandramohan wrote:

> > I was originally thinking about other regular operations, however considering that there might be transformation ops there could be other ops than just a omp.canonical_loop inside a omp.distribute region.
> > ```
> > omp.distribute {
> >    %cli = omp.cli
> >    canonical_loop ... %cli
> >    omp.tile(%cli, ... )
> > }
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > Also with the original proposal for omp.canonical_loop this would be possible.
> 
> The `omp.new_cli` operation is strictly not necessary. They could just be block arguments. We could go ahead with the nesting approach (with block arguments) and @jsjodin 's approach of providing CLI as operands of canonical loops and loop transformation ops. An example is given below
> 
> ```
> omp.distribute loops(%tloop)
> bb0 (%tloop : !omp.cli):   
>   omp.tile loops(%outer, %inner), construct(%tloop:!omp.cli) {
>   bb0 (%outer, %inner : !omp.cli, !omp.cli):
>     omp.canonical_loop %iv1 : i32 in [0, %tripcount), construct(%outer : !omp.cli){
>       omp.canonical_loop %iv2 : i32 in [0, %tc), construct(%inner : !omp.cli) {
>         %a = load %arrA[%iv1, %iv2] : memref<?x?xf32>
>         store %a, %arrB[%iv1, %iv2] : memref<?x?xf32>
>       }
>     }
>   }
> ```

I think tile was the long construct, I meant a collapse here.

```
omp.distribute loops(%tloop)
bb0 (%tloop : !omp.cli):   
  omp.collapse loops(%outer, %inner), construct(%tloop:!omp.cli) {
  bb0 (%outer, %inner : !omp.cli, !omp.cli):
    omp.canonical_loop %iv1 : i32 in [0, %tripcount), construct(%outer : !omp.cli){
      omp.canonical_loop %iv2 : i32 in [0, %tc), construct(%inner : !omp.cli) {
        %a = load %arrA[%iv1, %iv2] : memref<?x?xf32>
        store %a, %arrB[%iv1, %iv2] : memref<?x?xf32>
      }
    }
  }
  ```
`loops` refer to the loops on which the construct will operate, `construct` points back to the construct that will modify this loop or loop transformation construct.

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