[flang-commits] [flang] [llvm] [mlir] [flang][mlir][OpenMP] Allow integer host_eval uses in effect-free ops (PR #223475)

[Caroline Newcombe via flang-commits]

cenewcombe wrote:

Here’s the MLIR produced from the first subroutine in collapse-target-intervening.f90 in a host-only compilation (annotated by AI). The six `host_eval` operands are the collapsed loop's `(lb1, lb2, ub1, ub2, step1, step2)`, so `%arg4`, `%arg6` and `%arg8` are the inner lower bound, upper bound and step. `%8` and `%9` are where the host evaluates `n` and `m`.

All six `host_eval` block arguments (`%arg3`–`%arg8`) are used inside the region by `omp.loop_nest` as its bounds — the use `TargetOp::verifyRegions()` requires on host compilations. The guard at `%24`–`%28` additionally uses the inner three (`%arg4, %arg6, %arg8`) to compute `lb + ((ub-lb)/step)*step` and compare it against the inner IV, detecting the last inner iteration. The `fir.if` body then recomputes that value plus one step to get the terminal value of j, and runs the intervening `x = x + j`.

<details><summary>Details</summary>
```
  func.func @_QPrepro(%arg0: !fir.ref<i32> {fir.bindc_name = "n", fir.read_only}, %arg1: !fir.ref<i32> {fir.bindc_name = "m", fir.read_only}, %arg2: !fir.ref<i32> {fir.bindc_name = "x"}) {
    %0 = fir.dummy_scope : !fir.dscope
    %1 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFreproEi"}
    %2:2 = hlfir.declare %1 {uniq_name = "_QFreproEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
    %3 = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFreproEj"}
    %4:2 = hlfir.declare %3 {uniq_name = "_QFreproEj"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
    %5:2 = hlfir.declare %arg1 dummy_scope %0 arg 2 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QFreproEm"} : (!fir.ref<i32>, !fir.dscope) -> (!fir.ref<i32>, !fir.ref<i32>)
    %6:2 = hlfir.declare %arg0 dummy_scope %0 arg 1 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QFreproEn"} : (!fir.ref<i32>, !fir.dscope) -> (!fir.ref<i32>, !fir.ref<i32>)
    %7:2 = hlfir.declare %arg2 dummy_scope %0 arg 3 {fortran_attrs = #fir.var_attrs<intent_inout>, uniq_name = "_QFreproEx"} : (!fir.ref<i32>, !fir.dscope) -> (!fir.ref<i32>, !fir.ref<i32>)


    // The six loop bounds, evaluated here on the host.
    %c1_i32 = arith.constant 1 : i32     // outer lb
    %8 = fir.load %6#0 : !fir.ref<i32>   // outer ub = n
    %c1_i32_0 = arith.constant 1 : i32   // outer step
    %c1_i32_1 = arith.constant 1 : i32   // inner lb
    %9 = fir.load %5#0 : !fir.ref<i32>   // inner ub = m
    %c1_i32_2 = arith.constant 1 : i32   // inner step

    %10 = omp.map.info var_ptr(%7#1 : !fir.ref<i32>, i32) map_clauses(tofrom) capture(ByRef) name("x") -> !fir.ref<i32>
    %11 = omp.map.info var_ptr(%2#1 : !fir.ref<i32>, i32) map_clauses(implicit) capture(ByCopy) name("i") -> !fir.ref<i32>
    %12 = omp.map.info var_ptr(%6#1 : !fir.ref<i32>, i32) map_clauses(implicit) capture(ByCopy) name("n") -> !fir.ref<i32>
    %13 = omp.map.info var_ptr(%4#1 : !fir.ref<i32>, i32) map_clauses(implicit) capture(ByCopy) name("j") -> !fir.ref<i32>
    %14 = omp.map.info var_ptr(%5#1 : !fir.ref<i32>, i32) map_clauses(implicit) capture(ByCopy) name("m") -> !fir.ref<i32>

    // host_eval order is (lb1, lb2, ub1, ub2, step1, step2), so the inner
    // bounds arrive as %arg4 (lb), %arg6 (ub = m) and %arg8 (step).
    omp.target kernel_type(spmd) host_eval(%c1_i32 -> %arg3, %c1_i32_1 -> %arg4, %8 -> %arg5, %9 -> %arg6, %c1_i32_0 -> %arg7, %c1_i32_2 -> %arg8 : i32, i32, i32, i32, i32, i32) map_entries(%10 -> %arg9, %11 -> %arg10, %12 -> %arg11, %13 -> %arg12, %14 -> %arg13 : !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>) {
      %15:2 = hlfir.declare %arg9 {fortran_attrs = #fir.var_attrs<intent_inout>, uniq_name = "_QFreproEx"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
      %16:2 = hlfir.declare %arg10 {uniq_name = "_QFreproEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
      %17:2 = hlfir.declare %arg11 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QFreproEn"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
      %18:2 = hlfir.declare %arg12 {uniq_name = "_QFreproEj"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
      // Mapped copy of m, available in-region: this is what reading the bound
      // from the map/firstprivate binding would load from.
      %19:2 = hlfir.declare %arg13 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QFreproEm"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
      omp.teams {
        omp.parallel private(@_QFreproEi_private_i32 %16#0 -> %arg14, @_QFreproEj_private_i32 %18#0 -> %arg15 : !fir.ref<i32>, !fir.ref<i32>) {
          %20:2 = hlfir.declare %arg14 {uniq_name = "_QFreproEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
          %21:2 = hlfir.declare %arg15 {uniq_name = "_QFreproEj"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
          omp.distribute {
            omp.wsloop {
              // First in-region use of the host_eval values: all six are the
              // loop_nest bounds, which TargetOp::verifyRegions() requires.
              omp.loop_nest (%arg16, %arg17) : i32 = (%arg3, %arg4) to (%arg5, %arg6) inclusive step (%arg7, %arg8) collapse(2) {
                hlfir.assign %arg16 to %20#0 : i32, !fir.ref<i32>
                hlfir.assign %arg17 to %21#0 : i32, !fir.ref<i32>

                // x = x + 1
                %22 = fir.load %15#0 : !fir.ref<i32>
                %c1_i32_3 = arith.constant 1 : i32
                %23 = arith.addi %22, %c1_i32_3 : i32
                hlfir.assign %23 to %15#0 : i32, !fir.ref<i32>

                // Second use: guard the intervening code on the last inner
                // iteration, lb + ((ub-lb)/step)*step. Needs all three inner
                // bounds, and this is the arithmetic the verifier rejects.
                %24 = arith.subi %arg6, %arg4 : i32
                %25 = arith.divsi %24, %arg8 : i32
                %26 = arith.muli %25, %arg8 : i32
                %27 = arith.addi %arg4, %26 : i32
                %28 = arith.cmpi eq, %arg17, %27 : i32
                fir.if %28 {
                  // Fortran terminal value of j: the above plus one step.
                  %29 = arith.subi %arg6, %arg4 : i32
                  %30 = arith.divsi %29, %arg8 : i32
                  %31 = arith.muli %30, %arg8 : i32
                  %32 = arith.addi %arg4, %31 : i32
                  %33 = arith.addi %32, %arg8 : i32
                  hlfir.assign %33 to %21#0 : i32, !fir.ref<i32>

                  // Intervening statement: x = x + j
                  %34 = fir.load %15#0 : !fir.ref<i32>
                  %35 = fir.load %21#0 : !fir.ref<i32>
                  %36 = arith.addi %34, %35 : i32
                  hlfir.assign %36 to %15#0 : i32, !fir.ref<i32>
                }
                omp.yield
              }
            } {omp.composite}
          } {omp.composite}
          omp.terminator
        } {omp.composite}
        omp.terminator
      } {omp.combined}
      omp.terminator
    } {omp.combined}
    return
  }
​​```
</details> 

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