[Mlir-commits] [mlir] [MLIR][OpenMP] Add omp.loop_nest operation (PR #87083)

[Sergio Afonso]

skatrak wrote:

> > > What is the mechanism we now have to avoid code appearing between loop_nest and the wrapper operation?
> > 
> > 
> > The `LoopWrapperInterface` (in the follow-up PR: #87232) defines an `isWrapper()` method that checks that the operation region only contains a single block with a single wrapper or `omp.loop_nest` and a terminator. Then, the wrapper op's verifier can be updated to make sure this is the case (e.g. [#87239](https://github.com/llvm/llvm-project/pull/87239/files#diff-a897370ad8f5ad37e8c1adb3c145c2304aaa38da3227bc1d02ac701ee8dc0754), [#87253](https://github.com/llvm/llvm-project/pull/87253/files#diff-a897370ad8f5ad37e8c1adb3c145c2304aaa38da3227bc1d02ac701ee8dc0754), [#87365](https://github.com/llvm/llvm-project/pull/87365/files#diff-a897370ad8f5ad37e8c1adb3c145c2304aaa38da3227bc1d02ac701ee8dc0754)).
> 
> Would the code motion passes still end up moving the operation and then subsequently the verifier errors out? Have you tried out a code sinking or hoisting pass to see what is the behaviour?

I'm currently working on making `omp.wsloop` a wrapper, so I was able to make a small test to check this. I've noticed that CSE works as expected and doesn't try to insert anything inside of the wrapper, and the inlining pass seems to skip trying to inline anything inside of the loop nest region. Maybe you can suggest better ways to test this, but this is what I've tried so far:

```mlir
// cat test.mlir
func.func private @bar(i32) -> ()
func.func @foo(%arg0 : i32) -> () {
  func.call @bar(%arg0) : (i32) -> ()
  return
}
func.func @test_cse(%x : i32) -> () {
  %c1_0 = arith.constant 1 : i32
  %x_plus_1_0 = arith.addi %x, %c1_0 : i32
  func.call @foo(%x_plus_1_0) : (i32) -> ()
  %c2_0 = arith.constant 2 : i32
  %x_plus_2_0 = arith.addi %x, %c2_0 : i32
  omp.wsloop {
    omp.loop_nest (%i) : i32 = (%x) to (%x_plus_2_0) step (%c1_0) {
      %c2_1 = arith.constant 2 : i32
      %x_plus_2_1 = arith.addi %x, %c2_1 : i32
      func.call @foo(%x_plus_2_1) : (i32) -> ()
      omp.yield
    }
    omp.terminator
  }
  return
}

// mlir-opt test.mlir -inline -cse -o -
func.func private @bar(i32)
func.func @foo(%arg0: i32) {
  call @bar(%arg0) : (i32) -> ()
  return
}
func.func @test_cse(%arg0: i32) {
  %c2_i32 = arith.constant 2 : i32
  %c1_i32 = arith.constant 1 : i32
  %0 = arith.addi %arg0, %c1_i32 : i32
  call @bar(%0) : (i32) -> ()
  %1 = arith.addi %arg0, %c2_i32 : i32
  omp.wsloop {
    omp.loop_nest (%arg1) : i32 = (%arg0) to (%1) step (%c1_i32) {
      func.call @foo(%1) : (i32) -> ()
      omp.yield
    }
    omp.terminator
  }
  return
}
```

> For omp.wsloop where will collapse be represented? In the loop_nest operation?

Yes, so it will be possible to represent `collapse` for any loop construct, composite or otherwise (support would also need MLIR to LLVM IR translation implemented for that case). The way it's represented is the same as it's currently done for `omp.wsloop`, by having multiple IVs and range/step variables.

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