[flang-commits] [flang] [Flang][OpenMP] Add Lowering support for Collapse with Taskloop (PR #166791)

[Jack Styles via flang-commits]

https://github.com/Stylie777 updated https://github.com/llvm/llvm-project/pull/166791

>From caabd5e42c73914201a654b5276df76465241ffe Mon Sep 17 00:00:00 2001
From: Jack Styles <jack.styles at arm.com>
Date: Thu, 6 Nov 2025 15:48:45 +0000
Subject: [PATCH 1/2] [Flang][OpenMP] Add Lowering support for Collapse with
 Taskloop

Supprot for lowering collapse already exists within
`genLoopNestClauses`, which is called when lowering taskloop. However,
the TODO message still included the Collapse clause, so it was not
activated. By removing this, it enables lowering of the Collapse
clause in taskloop.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp             |  2 +-
 flang/test/Lower/OpenMP/taskloop-collapse.f90 | 34 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 flang/test/Lower/OpenMP/taskloop-collapse.f90

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index ad456d89bc432..e271a84bc5b4e 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -1773,7 +1773,7 @@ static void genTaskloopClauses(lower::AbstractConverter &converter,
   cp.processGrainsize(stmtCtx, clauseOps);
   cp.processNumTasks(stmtCtx, clauseOps);
 
-  cp.processTODO<clause::Allocate, clause::Collapse, clause::Default,
+  cp.processTODO<clause::Allocate, clause::Default,
                  clause::Final, clause::If, clause::InReduction,
                  clause::Lastprivate, clause::Mergeable, clause::Nogroup,
                  clause::Priority, clause::Reduction, clause::Shared,
diff --git a/flang/test/Lower/OpenMP/taskloop-collapse.f90 b/flang/test/Lower/OpenMP/taskloop-collapse.f90
new file mode 100644
index 0000000000000..48243640d07b9
--- /dev/null
+++ b/flang/test/Lower/OpenMP/taskloop-collapse.f90
@@ -0,0 +1,34 @@
+! Test the collapse clause when being used with the taskloop construct
+! RUN: bbc -emit-hlfir -fopenmp -fopenmp-version=45 %s -o - 2>&1 | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=45 %s -o - 2>&1 | FileCheck %s
+
+! CHECK-LABEL: omp.private
+! CHECK-SAME: {type = private} @[[J_PRIVATE:.*]] : i32
+! CHECK-LABEL: omp.private
+! CHECK-SAME: {type = private} @[[I_PRIVATE:.*]] : i32
+! CHECK-LABEL: omp.private
+! CHECK-SAME: {type = firstprivate} @[[SUM_FIRSTPRIVATE:.*]] : i32 copy
+
+! CHECK-LABEL: func.func @_QPtest()
+! CHECK: %[[ALLOCA_I:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFtestEi"}
+! CHECK: %[[DECLARE_I:.*]]:2 = hlfir.declare %1 {uniq_name = "_QFtestEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[ALLOCA_J:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFtestEj"}
+! CHECK: %[[DECLARE_J:.*]]:2 = hlfir.declare %3 {uniq_name = "_QFtestEj"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[ALLOCA_SUM:.*]] = fir.alloca i32 {bindc_name = "sum", uniq_name = "_QFtestEsum"}
+! CHECK: %[[DECLARE_SUM:.*]]:2 = hlfir.declare %5 {uniq_name = "_QFtestEsum"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+
+subroutine test()
+    integer :: i, j, sum
+
+    !$omp taskloop collapse(2)
+    ! CHECK-LABEL: omp.taskloop
+    ! CHECK-SAME: private(@_QFtestEsum_firstprivate_i32 %[[DECLARE_SUM]]#0 -> %arg0, @_QFtestEi_private_i32 %[[DECLARE_I]]#0 -> %arg1, @_QFtestEj_private_i32 %[[DECLARE_J]]#0 -> %arg2 : !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>)
+    ! CHECK-LABEL: omp.loop_nest
+    ! CHECK-SAME: (%arg3, %arg4) : i32 = (%c1_i32, %c1_i32_1) to (%c10_i32, %c5_i32) inclusive step (%c1_i32_0, %c1_i32_2) collapse(2)
+    do i = 1, 10
+        do j = 1, 5
+            sum = sum + i + j
+        end do
+    end do
+    !$omp end taskloop
+end subroutine

>From 543fad0232752003e4a481afd7bb613db068eb9e Mon Sep 17 00:00:00 2001
From: Jack Styles <jack.styles at arm.com>
Date: Thu, 6 Nov 2025 16:21:13 +0000
Subject: [PATCH 2/2] formatting

---
 flang/lib/Lower/OpenMP/OpenMP.cpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index e271a84bc5b4e..374fc62e7ff73 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -1773,11 +1773,11 @@ static void genTaskloopClauses(lower::AbstractConverter &converter,
   cp.processGrainsize(stmtCtx, clauseOps);
   cp.processNumTasks(stmtCtx, clauseOps);
 
-  cp.processTODO<clause::Allocate, clause::Default,
-                 clause::Final, clause::If, clause::InReduction,
-                 clause::Lastprivate, clause::Mergeable, clause::Nogroup,
-                 clause::Priority, clause::Reduction, clause::Shared,
-                 clause::Untied>(loc, llvm::omp::Directive::OMPD_taskloop);
+  cp.processTODO<clause::Allocate, clause::Default, clause::Final, clause::If,
+                 clause::InReduction, clause::Lastprivate, clause::Mergeable,
+                 clause::Nogroup, clause::Priority, clause::Reduction,
+                 clause::Shared, clause::Untied>(
+      loc, llvm::omp::Directive::OMPD_taskloop);
 }
 
 static void genTaskwaitClauses(lower::AbstractConverter &converter,