[Flang][OpenMP] Add Lowering support for Collapse with Taskloop

[Stylie777]

Support 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.

[llvmbot]

@llvm/pr-subscribers-flang-semantics
@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-flang-fir-hlfir

Author: Jack Styles (Stylie777)

Changes

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.

---

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

2 Files Affected:

• (modified) flang/lib/Lower/OpenMP/OpenMP.cpp (+1-1)
• (added) flang/test/Lower/OpenMP/taskloop-collapse.f90 (+34)

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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[kaviya2510]

Thankyou for the patch. LGTM!

[tblah]

How does this interact with #166903? If that patch were merged would collapse automatically work end-to-end or would this require an extra TODO message?

[Stylie777]

How does this interact with #166903? If that patch were merged would collapse automatically work end-to-end or would this require an extra TODO message?
@tblah The way that Collapse support is working here is that the Collapse clause gets attached to the omp.loop_nest rather than the omp.taskloop in FIR.

Take the following example

subroutine f00
  integer :: i, j, sum

  sum = 0
  !$omp taskloop collapse(2)
  do i = 1,10
    do j = 1,5
    sum = sum + i + j
    end do
  end do
  !$omp end taskloop

end subroutine

This creates the taskloop and loopnest as

omp.taskloop private(@_QFf00Esum_firstprivate_i32 %6#0 -> %arg0, @_QFf00Ei_private_i32 %2#0 -> %arg1, @_QFf00Ej_private_i32 %4#0 -> %arg2 : !fir.ref<i32>, !fir.ref<i32>, !fir.ref<i32>) {
      omp.loop_nest (%arg3, %arg4) : i32 = (%c1_i32, %c1_i32_1) to (%c10_i32, %c5_i32) inclusive step (%c1_i32_0, %c1_i32_2) collapse(2) {

It looks like we already have support for Collapse when lowering omp.loop_nest from FIR to LLVM IR, so it interfaces well with the patch from @kaviya2510 to add taskloop support to MLIR.

[tblah]

Thanks

[kaviya2510]

How does this interact with #166903? If that patch were merged would collapse automatically work end-to-end or would this require an extra TODO message?

Yes, it is a valid scenario. I tested the testcase shared by @Stylie777 with taskloop translation patch and it is not generating the expected result. I thought that collapse clause requires some adjustments specific to taskloop, so can we mark it as a TODO in translation patch using getCollapseNumLoops() and address it later?

[tblah]

Can we mark it as a TODO in translation patch using getCollapseNumLoops() and address it later?

+1

[Stylie777]

If that is the case, I will go ahead and fix the merge conflict and merge this commit so you can properly test the TODO message for Collapse @kaviya2510.
Thanks for the information on this, its good to know!