[Mlir-commits] [flang] [mlir] [OpenMP][Flang][MLIR] Skip trip count calculation when bounds are null (PR #176469)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-openmp

Author: Jason Van Beusekom (Jason-Van-Beusekom)

<details>
<summary>Changes</summary>

Fixes a segfault when trip count values are null by skipping trip count calculation when we cannot determine if it is safe to hoist out the values.

Of note I originally tried to modify `extractOnlyOmpNestedDir` to return the first OpenMPConstruct directive, skipping over any earlier directives (ie stores), which did work for the below generic test case:

```fortran
program minimal_repro
  implicit none

  integer :: i, m
  integer :: res(10) = 0

!$omp target teams map(from:m,res) private(m)
  m = 5
!$omp distribute parallel do
  do i = 1, 10
    res(i) = 5 + i
  end do
!$omp end distribute parallel do
!$omp end target teams

end program minimal_repro
```

But that led to incorrect output in this test case as the trip count was hoisted out and calculated by m(1000000) instead of m(1)
```fortran
program minimal_repro
  implicit none

  integer :: i, x
  integer :: m(1) = 0
  integer :: res(10) = 0
  m(1) = 10
  x = 1000000
!$omp target teams map(res)
  x = 1
!$omp distribute parallel do
  do i = 1, m(x)
    res(i) = 5 + i
  end do
!$omp end distribute parallel do
!$omp end target teams

print *, "Test completed successfully m =", m, " res=", res

end program minimal_repro
```
Leading to a segfault, due to the loop bounds being calculated with m(1000000)
```mlir
    %c1000000_i32 = arith.constant 1000000 : i32
    hlfir.assign %c1000000_i32 to %10#<!-- -->0 : i32, !fir.ref<i32>
    %c1_i32 = arith.constant 1 : i32
    %12 = fir.load %10#<!-- -->0 : !fir.ref<i32>
    %13 = fir.convert %12 : (i32) -> i64
    %14 = hlfir.designate %5#<!-- -->0 (%13)  : (!fir.ref<!fir.array<1xi32>>, i64) -> !fir.ref<i32>
    %15 = fir.load %14 : !fir.ref<i32>
    ...
    omp.target host_eval(%c1_i32 -> %arg0, %15 -> %arg1, %c1_i32_1 -> %arg2 : i32, i32, i32) map_entries(%18 -> %arg3, %19 -> %arg4, %20 -> %arg5, %23 -> %arg6 : !fir.ref<!fir.array<10xi32>>, !fir.ref<i32>, !fir.ref<i32>, !fir.ref<!fir.array<1xi32>>) {
      ...
      omp.teams {
              ...
              omp.loop_nest (%arg8) : i32 = (%arg0) to (%arg1) inclusive step (%arg2) {
 
```

The wip commit for this change is here: https://github.com/Jason-Van-Beusekom/llvm-project/commit/beafeae39600bd575e52cb942996c3789c461e25

We would need to have some sort of intelligent hoisting for these cases, to allow hoisting, but for now I just created this PR to fix the bug. 

Fixes: #<!-- -->176030

---
Full diff: https://github.com/llvm/llvm-project/pull/176469.diff


2 Files Affected:

- (modified) flang/test/Lower/OpenMP/host-eval.f90 (+47) 
- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+3) 


``````````diff
diff --git a/flang/test/Lower/OpenMP/host-eval.f90 b/flang/test/Lower/OpenMP/host-eval.f90
index cd759a988e4f5..735e8b8c811ed 100644
--- a/flang/test/Lower/OpenMP/host-eval.f90
+++ b/flang/test/Lower/OpenMP/host-eval.f90
@@ -283,3 +283,50 @@ subroutine loop()
   end do
   !$omp end target teams
 end subroutine loop
+
+! BOTH-LABEL: func.func @_QPdistribute_parallel_do_with_modified_trip
+subroutine distribute_parallel_do_with_modified_trip()
+  integer :: i, x
+  integer :: m(1)
+  integer :: res(10)
+  m(1) = 10
+  x = 1000000
+
+  ! BOTH: omp.target
+  ! BOTH-NOT: host_eval({{.*}})
+  ! BOTH-SAME: {
+  ! BOTH: omp.teams
+  !$omp target teams map(res)
+  x = 1
+  ! BOTH: omp.parallel
+  ! BOTH: omp.distribute
+  ! BOTH-NEXT: omp.wsloop
+  !$omp distribute parallel do
+  do i = 1, m(x)
+    res(i) = 5 + i
+  end do
+  !$omp end distribute parallel do
+  !$omp end target teams
+end subroutine distribute_parallel_do_with_modified_trip
+
+! BOTH-LABEL: func.func @_QPdistribute_parallel_do_with_assignment
+subroutine distribute_parallel_do_with_assignment()
+  integer :: i, m
+  integer :: res(10)
+
+  ! BOTH: omp.target
+  ! BOTH-NOT: host_eval({{.*}})
+  ! BOTH-SAME: {
+  ! BOTH: omp.teams
+  !$omp target teams map(from:m,res) private(m)
+  m = 5
+  ! BOTH: omp.parallel
+  ! BOTH: omp.distribute
+  ! BOTH-NEXT: omp.wsloop
+  !$omp distribute parallel do
+  do i = 1, 10
+    res(i) = 5 + i
+  end do
+  !$omp end distribute parallel do
+  !$omp end target teams
+end subroutine distribute_parallel_do_with_assignment
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 4e7942e382c8b..6d6ee38c468c4 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6314,6 +6314,9 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
       llvm::Value *upperBound = moduleTranslation.lookupValue(loopUpper);
       llvm::Value *step = moduleTranslation.lookupValue(loopStep);
 
+      if (!lowerBound || !upperBound || !step)
+        continue;
+
       llvm::OpenMPIRBuilder::LocationDescription loc(builder);
       llvm::Value *tripCount = ompBuilder->calculateCanonicalLoopTripCount(
           loc, lowerBound, upperBound, step, /*IsSigned=*/true,

``````````

</details>


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