[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 d8a773c57c0975a940beb4b699bdfd36c6eeb11e 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 1/2] [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      |  3 ++
 2 files changed, 50 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..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,

>From d8c881bdec38ca9881a697331e0e8021aeac9bd3 Mon Sep 17 00:00:00 2001
From: Jason Van Beusekom <jason.van-beusekom at hpe.com>
Date: Tue, 20 Jan 2026 10:10:23 -0600
Subject: [PATCH 2/2] Update
 mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Co-authored-by: Sergio Afonso <safonsof at amd.com>
---
 .../LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp     | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 6d6ee38c468c4..9e5896d3a01c4 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -6314,8 +6314,10 @@ initTargetRuntimeAttrs(llvm::IRBuilderBase &builder,
       llvm::Value *upperBound = moduleTranslation.lookupValue(loopUpper);
       llvm::Value *step = moduleTranslation.lookupValue(loopStep);
 
-      if (!lowerBound || !upperBound || !step)
-        continue;
+      if (!lowerBound || !upperBound || !step) {
+        attrs.LoopTripCount = nullptr;
+        break;
+      }
 
       llvm::OpenMPIRBuilder::LocationDescription loc(builder);
       llvm::Value *tripCount = ompBuilder->calculateCanonicalLoopTripCount(