Delayed privatization for `omp.wsloop` (or the need for pure `alloca` regions in `omp.private`)

Subprojects Flang
,
1

I am trying to enable delayed privatization by default for omp.wsloop ops and I am facing a problem. Consider the following subroutine:

subroutine test2 (xyz)
  integer :: i
  integer :: xyz(:)

  !$omp target map(from:xyz)
    !$omp do private(xyz)
      do i = 1, 10
        ...
      end do
  !$omp end target
end subroutine

The omp.private alloc region looks like this:

  omp.private {type = private} @_QFtest2Exyz_private_box_Uxi32 : !fir.box<!fir.array<?xi32>> alloc {
  ^bb0(%arg0: !fir.box<!fir.array<?xi32>>):
    ...
    %1:3 = fir.box_dims %arg0, %c0_0 : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
    %2 = fir.alloca !fir.array<?xi32>, %1#1 {bindc_name = "xyz", pinned, uniq_name = "_QFtest2Exyz"}
    %3 = fir.shape %1#1 : (index) -> !fir.shape<1>
    ...
  }

Note that %arg0 is used to retrieve the shape needed for the allocation.

If we look at the corresponding omp.target op in which the privatizer is used:

    omp.target map_entries(%59 -> %arg1, %63 -> %arg2, %64 -> %arg3, %58 -> %arg4, %62 -> %arg5 : !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr) {
      ...
      %66 = llvm.alloca %65 x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
      ...
      llvm.store %104, %66 : !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)>, !llvm.ptr
      omp.wsloop private(@_QFtest2Exyz_private_box_Uxi32 %66 -> %arg6, @_QFtest2Ei_private_ref_i32 %arg3 -> %arg7 : !llvm.ptr, !llvm.ptr) {
        ...
      }
      omp.terminator
    }

The problem occurs when we try to inline the alloc region of the omp.private op in the outlined target region. If we choose an alloc block like we do now, we end up with invalid IR since there is an operation in the alloc block (that op being the one corresponding to fir.box dims in the alloc region of the privatizer) that references an SSA value defined in a later block (the value being %66 in this example and the block is the first block in the target region).

I think this boils down to (possibly) a wrong design decision in the omp.private op. I think we should have introduced a pure alloc block that does not take any arguments and moved the rest of the logic to a separate block (I think similar to reductions’ alloc and init blocks). Not sure if this possible in case of omp.private though since they accept lesss restricted range of values compared to reductions (I think).

I have a solution that is less than ideal which is implemented by this abandoned PR: [flang][OpenMP] Allow saving first block of an OMP region for allocas by ergawy · Pull Request #121886 · llvm/llvm-project · GitHub.

Any input on how you prefer to proceed with this issue?

I think we have 2 options:

  1. Accept the above PR as a final solution to the issue.
  2. Temporarily apply the above PR and add a TODO to extend omp.private with a pure alloc region + a new init region similar to reductions.

People might have other better informed options, so suggestions are welcome.

WDYT?

cc @kiranchandramohan

2

I am already working on a pure alloc region separate from initialization as part of [RFC][OpenMP] supporting delayed task execution with firstprivate variables - #7 by tblah

I think [flang][OpenMP] Fix reduction init region block management by ergawy · Pull Request #122079 · llvm/llvm-project · GitHub is the right solution. Maybe target should be adapted to treat the current alloc region as an initialization region for the time being.

2 Likes
3

Awesome, I did not notice your work on this (I should check discourse more often). Thanks!