[Mlir-commits] [mlir] [mlir][OpenMP] Don't use label prefixes on linear variable rewrite (PR #200900)

[Leandro Lupori]

luporl wrote:

> ```
>   llvm.func @wsloop_linear_post_use(%lb : i32, %ub : i32, %step : i32,
>                                     %x : !llvm.ptr, %out : !llvm.ptr) {
>     omp.wsloop linear(%x : !llvm.ptr = %step : i32) {
>       omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
>         %cur = llvm.load %x : !llvm.ptr -> i32
>         llvm.store %cur, %out : i32, !llvm.ptr
>         omp.yield
>       }
>     } {linear_var_types = [i32]}
>     %after = llvm.load %x : !llvm.ptr -> i32
>     llvm.store %after, %out : i32, !llvm.ptr
>     llvm.return
>   }
> ```
> 
> The new traversal collects every successor reachable from startBB until it pops endBB, so it also collects blocks outside the actual loop-body region. (I think) This can be seen for translation of this reproducer.
> 

The use after the loop (`%after = llvm.load %x : !llvm.ptr -> i32`) was actually being translated correctly, as `%x` was not being replaced by the privatized linear variable. However, uses outside of the loop body, such as on `omp_loop.preheader` and `omp_loop.linear_lastiter_exit`, were incorrectly being replaced, which wasn't occurring before this PR.

> I think one solution would be to use the `llvm::CanonicalLoopInfo` for the start and end blocks: as I understand it, `loopInfo->getBody()` gives the single entry block for the loop body, and `loopInfo->getLatch()` is the block reached after the body. All arbitrary control flow inside of the loop body should always eventually branch to the latch so those should work as safer start and end blocks for the walk.

I have used this approach, but I had to move the rewrite to before `applyWorkshareLoop`/`applySimd`, as they invalidate `loopInfo`. This should be safe, as apparently they don't manipulate the iteration variable specified by the user, but only the one from `omp.loop_nest`.

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