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

[Jason Van Beusekom]

https://github.com/Jason-Van-Beusekom updated https://github.com/llvm/llvm-project/pull/176469

>From 1755897b08437b71ca3c70529a019382bb26b734 Mon Sep 17 00:00:00 2001
From: Jason Van Beusekom <jason.van-beusekom at hpe.com>
Date: Fri, 16 Jan 2026 12:45:14 -0600
Subject: [PATCH] [OpenMP][Flang][MLIR] Skip trip count calculation when bounds
 are null 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.

---
 flang/test/Lower/OpenMP/host-eval.f90         | 47 +++++++++++++++++++
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      |  5 ++
 2 files changed, 52 insertions(+)

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..9e5896d3a01c4 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6314,6 +6314,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,