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

[llvmlistbot at llvm.org]

Author: Jason Van Beusekom
Date: 2026-01-21T11:56:36Z
New Revision: 0bdbf01e4e5dbb867bbca0747b327b86f8bb4f43

URL: https://github.com/llvm/llvm-project/commit/0bdbf01e4e5dbb867bbca0747b327b86f8bb4f43
DIFF: https://github.com/llvm/llvm-project/commit/0bdbf01e4e5dbb867bbca0747b327b86f8bb4f43.diff

LOG: [OpenMP][Flang][MLIR] Skip trip count calculation when bounds are null (#176469)

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

Added: 
    

Modified: 
    flang/test/Lower/OpenMP/host-eval.f90
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Removed: 
    


################################################################################
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 fa907233612e2..781928e79db35 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6319,6 +6319,11 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
       llvm::Value *upperBound = moduleTranslation.lookupValue(loopUpper);
       llvm::Value *step = moduleTranslation.lookupValue(loopStep);
 
+      if (!lowerBound || !upperBound || !step) {
+        attrs.LoopTripCount = nullptr;
+        break;
+      }
+
       llvm::OpenMPIRBuilder::LocationDescription loc(builder);
       llvm::Value *tripCount = ompBuilder->calculateCanonicalLoopTripCount(
           loc, lowerBound, upperBound, step, /*IsSigned=*/true,