[flang][OpenMP] Enable lastprivate with collapse #93778

NimishMishra · 2024-05-30T07:45:34Z

This PR enables the lowering of lastprivate in presence of collapse clause.

llvmbot · 2024-05-30T07:46:06Z

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

Author: None (NimishMishra)

Changes

This PR enables the lowering of lastprivate in presence of collapse clause.

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

2 Files Affected:

(modified) flang/lib/Lower/OpenMP/DataSharingProcessor.cpp (+1-6)
(modified) flang/test/Lower/OpenMP/wsloop-collapse.f90 (+33)

diff --git a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
index 557a9685024c5..8aa37035ddee5 100644
--- a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
@@ -143,7 +143,7 @@ void DataSharingProcessor::collectOmpObjectListSymbol(
 }
 
 void DataSharingProcessor::collectSymbolsForPrivatization() {
-  bool hasCollapse = false;
+
   for (const omp::Clause &clause : clauses) {
     if (const auto &privateClause =
             std::get_if<omp::clause::Private>(&clause.u)) {
@@ -157,16 +157,11 @@ void DataSharingProcessor::collectSymbolsForPrivatization() {
       const ObjectList &objects = std::get<ObjectList>(lastPrivateClause->t);
       collectOmpObjectListSymbol(objects, explicitlyPrivatizedSymbols);
       hasLastPrivateOp = true;
-    } else if (std::get_if<omp::clause::Collapse>(&clause.u)) {
-      hasCollapse = true;
     }
   }
 
   for (auto *sym : explicitlyPrivatizedSymbols)
     allPrivatizedSymbols.insert(sym);
-
-  if (hasCollapse && hasLastPrivateOp)
-    TODO(converter.getCurrentLocation(), "Collapse clause with lastprivate");
 }
 
 bool DataSharingProcessor::needBarrier() {
diff --git a/flang/test/Lower/OpenMP/wsloop-collapse.f90 b/flang/test/Lower/OpenMP/wsloop-collapse.f90
index 67351ca275efb..3e862c25a8089 100644
--- a/flang/test/Lower/OpenMP/wsloop-collapse.f90
+++ b/flang/test/Lower/OpenMP/wsloop-collapse.f90
@@ -94,3 +94,36 @@ program wsloop_collapse
 !CHECK:         return
 !CHECK:       }
 end program wsloop_collapse
+
+subroutine collapse_with_lastprivate
+!CHECK: %[[VAL_I:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFcollapse_with_lastprivateEi"}
+!CHECK: %[[VAL_I_DECLARE:.*]]:2 = hlfir.declare %[[VAL_I]] {{.*}}
+!CHECK: %[[VAL_J:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFcollapse_with_lastprivateEj"}
+!CHECK: %[[VAL_J_DECLARE:.*]]:2 = hlfir.declare %[[VAL_J]] {{.*}}
+!CHECK: %[[VAL_K:.*]] = fir.alloca i32 {bindc_name = "k", uniq_name = "_QFcollapse_with_lastprivateEk"}
+!CHECK: %[[VAL_K_DECLARE:.*]]:2 = hlfir.declare %[[VAL_K]] {{.}}
+!CHECK: %[[VAL_X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFcollapse_with_lastprivateEx"}
+!CHECK: %[[VAL_X_DECLARE:.*]]:2 = hlfir.declare %[[VAL_X]] {{.*}}
+
+        integer :: x, i, j, k
+
+!CHECK: omp.parallel {
+!CHECK: %[[INNER_X:.*]] = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFcollapse_with_lastprivateEx"}
+!CHECK: %[[INNER_X_DECLARE:.*]]:2 = hlfir.declare %[[INNER_X]] {{.*}}
+        !$omp parallel do lastprivate(x) collapse(2)
+!CHECK: omp.wsloop {
+!CHECK: omp.loop_nest (%[[ARG0:.*]], %[[ARG1:.*]]) : i32 = (%[[LB1:.*]], %[[LB2:.*]]) to (%[[UB1:.*]], %[[UB2:.*]]) inclusive step (%[[STEP1:.*]], %[[STEP2:.*]]) {
+           do i = 1, 10
+             do j = 1, 20
+!CHECK: %[[LOOP_CONTROL:.*]]:2 = fir.do_loop %[[ARG2:.*]] = {{.*}}
+               do k = 1, 30
+                 x = x + 1
+               end do
+             end do
+           end do
+!CHECK: fir.if {{.*}} {
+!CHECK: %[[LOADED_X:.*]] = fir.load %[[INNER_X_DECLARE]]#0 : !fir.ref<i32>
+!CHECK: hlfir.assign %[[LOADED_X]] to %[[VAL_X_DECLARE]]#0 temporary_lhs : i32, !fir.ref<i32>
+!CHECK: }
+        !$omp end parallel do
+end subroutine

llvmbot · 2024-05-30T07:46:06Z

@llvm/pr-subscribers-flang-openmp

Author: None (NimishMishra)

Changes

This PR enables the lowering of lastprivate in presence of collapse clause.

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

2 Files Affected:

(modified) flang/lib/Lower/OpenMP/DataSharingProcessor.cpp (+1-6)
(modified) flang/test/Lower/OpenMP/wsloop-collapse.f90 (+33)

diff --git a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
index 557a9685024c5..8aa37035ddee5 100644
--- a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
@@ -143,7 +143,7 @@ void DataSharingProcessor::collectOmpObjectListSymbol(
 }
 
 void DataSharingProcessor::collectSymbolsForPrivatization() {
-  bool hasCollapse = false;
+
   for (const omp::Clause &clause : clauses) {
     if (const auto &privateClause =
             std::get_if<omp::clause::Private>(&clause.u)) {
@@ -157,16 +157,11 @@ void DataSharingProcessor::collectSymbolsForPrivatization() {
       const ObjectList &objects = std::get<ObjectList>(lastPrivateClause->t);
       collectOmpObjectListSymbol(objects, explicitlyPrivatizedSymbols);
       hasLastPrivateOp = true;
-    } else if (std::get_if<omp::clause::Collapse>(&clause.u)) {
-      hasCollapse = true;
     }
   }
 
   for (auto *sym : explicitlyPrivatizedSymbols)
     allPrivatizedSymbols.insert(sym);
-
-  if (hasCollapse && hasLastPrivateOp)
-    TODO(converter.getCurrentLocation(), "Collapse clause with lastprivate");
 }
 
 bool DataSharingProcessor::needBarrier() {
diff --git a/flang/test/Lower/OpenMP/wsloop-collapse.f90 b/flang/test/Lower/OpenMP/wsloop-collapse.f90
index 67351ca275efb..3e862c25a8089 100644
--- a/flang/test/Lower/OpenMP/wsloop-collapse.f90
+++ b/flang/test/Lower/OpenMP/wsloop-collapse.f90
@@ -94,3 +94,36 @@ program wsloop_collapse
 !CHECK:         return
 !CHECK:       }
 end program wsloop_collapse
+
+subroutine collapse_with_lastprivate
+!CHECK: %[[VAL_I:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFcollapse_with_lastprivateEi"}
+!CHECK: %[[VAL_I_DECLARE:.*]]:2 = hlfir.declare %[[VAL_I]] {{.*}}
+!CHECK: %[[VAL_J:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFcollapse_with_lastprivateEj"}
+!CHECK: %[[VAL_J_DECLARE:.*]]:2 = hlfir.declare %[[VAL_J]] {{.*}}
+!CHECK: %[[VAL_K:.*]] = fir.alloca i32 {bindc_name = "k", uniq_name = "_QFcollapse_with_lastprivateEk"}
+!CHECK: %[[VAL_K_DECLARE:.*]]:2 = hlfir.declare %[[VAL_K]] {{.}}
+!CHECK: %[[VAL_X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFcollapse_with_lastprivateEx"}
+!CHECK: %[[VAL_X_DECLARE:.*]]:2 = hlfir.declare %[[VAL_X]] {{.*}}
+
+        integer :: x, i, j, k
+
+!CHECK: omp.parallel {
+!CHECK: %[[INNER_X:.*]] = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFcollapse_with_lastprivateEx"}
+!CHECK: %[[INNER_X_DECLARE:.*]]:2 = hlfir.declare %[[INNER_X]] {{.*}}
+        !$omp parallel do lastprivate(x) collapse(2)
+!CHECK: omp.wsloop {
+!CHECK: omp.loop_nest (%[[ARG0:.*]], %[[ARG1:.*]]) : i32 = (%[[LB1:.*]], %[[LB2:.*]]) to (%[[UB1:.*]], %[[UB2:.*]]) inclusive step (%[[STEP1:.*]], %[[STEP2:.*]]) {
+           do i = 1, 10
+             do j = 1, 20
+!CHECK: %[[LOOP_CONTROL:.*]]:2 = fir.do_loop %[[ARG2:.*]] = {{.*}}
+               do k = 1, 30
+                 x = x + 1
+               end do
+             end do
+           end do
+!CHECK: fir.if {{.*}} {
+!CHECK: %[[LOADED_X:.*]] = fir.load %[[INNER_X_DECLARE]]#0 : !fir.ref<i32>
+!CHECK: hlfir.assign %[[LOADED_X]] to %[[VAL_X_DECLARE]]#0 temporary_lhs : i32, !fir.ref<i32>
+!CHECK: }
+        !$omp end parallel do
+end subroutine

tblah

LGTM, thanks!

tblah · 2024-05-31T10:17:42Z

flang/lib/Lower/OpenMP/DataSharingProcessor.cpp

@@ -143,7 +143,7 @@ void DataSharingProcessor::collectOmpObjectListSymbol(
 }

 void DataSharingProcessor::collectSymbolsForPrivatization() {
-  bool hasCollapse = false;
+


super nit: an empty line at the start of a function looks wrong to me (feel free to ignore)

kiranchandramohan · 2024-05-31T16:17:20Z

Have you checked whether the condition generation for the last iteration is correct when collapse is specified?

luporl

As @kiranchandramohan mentioned, DataSharingProcessor::insertLastPrivateCompare is currently assuming collapse == 1.

While this patch may produce correct results, it would result in the lastprivate variable being assigned to on every iteration of the inner loops, when on the last iteration of the outer loop, which doesn't seem right.

However, checking all related iteration variables during many iterations would not be very efficient. IIRC, some time ago it was suggested that lastprivate handling should be moved to omp dialect, to make it possible to update the lastprivate variable in the end of the loop and avoid iteration variable comparisons.

kiranchandramohan · 2024-06-03T09:26:59Z

As @kiranchandramohan mentioned, DataSharingProcessor::insertLastPrivateCompare is currently assuming collapse == 1.

While this patch may produce correct results, it would result in the lastprivate variable being assigned to on every iteration of the inner loops, when on the last iteration of the outer loop, which doesn't seem right.

However, checking all related iteration variables during many iterations would not be very efficient. IIRC, some time ago it was suggested that lastprivate handling should be moved to omp dialect, to make it possible to update the lastprivate variable in the end of the loop and avoid iteration variable comparisons.

The previous suggestion was to create an omp.lastprivate operation that takes the place of the fir.if that we insert for the lastprivate comparison. Now that @ergawy has made progress with the delayed privatization, you might be able to use the infrastructure he created to do this. The idea is to use the %.omp.is_last variable (which identifies whether this is the last iteration or not) that is given as input to the __kmpc_for_static_init_4 openmp runtime call. You can use this in the exit block of the loop and do a comparison and branch to either do the lastprivate update (if it is the last iteration) or exit. All this to be done during translation or in the OpenMP IRBuilder.

call void @__kmpc_for_static_init_4(ptr @1, i32 %2, i32 34, ptr %.omp.is_last, ptr %.omp.lb, ptr %.omp.ub, ptr %.omp.stride, i32 1, i32 1)
; ... intervening code
omp.loop.exit:                                    ; preds = %omp.inner.for.end
  call void @__kmpc_for_static_fini(ptr @1, i32 %2)
  %14 = load i32, ptr %.omp.is_last, align 4
  %15 = icmp ne i32 %14, 0
  br i1 %15, label %.omp.lastprivate.then, label %.omp.lastprivate.done

.omp.lastprivate.then:                            ; preds = %omp.loop.exit
  %16 = load i32, ptr %x2, align 4
  store i32 %16, ptr %0, align 4
  br label %.omp.lastprivate.done

.omp.lastprivate.done:                            ; preds = %.omp.lastprivate.then, %omp.loop.exit
  ret void

ergawy

LG, just one nit.

ergawy · 2024-06-07T09:36:51Z

flang/test/Lower/OpenMP/wsloop-collapse.f90

+!CHECK: %[[LOADED_X:.*]] = fir.load %[[INNER_X_DECLARE]]#0 : !fir.ref<i32>
+!CHECK: hlfir.assign %[[LOADED_X]] to %[[VAL_X_DECLARE]]#0 temporary_lhs : i32, !fir.ref<i32>
+!CHECK: }
+        !$omp end parallel do


nit: similar to #93786 (comment), can you add the terminators to make the region structure more obvious?

DavidTruby · 2024-07-18T14:33:45Z

Hi, just removing the current check doesn't quite do the right thing, as the value of the lastprivate variable will be a random one of the last values of the inner loops, rather than the definite final one. To implement this we need to check that we have reached the final value of all the loop indices, not just the outer one.

I've submitted a patch with this approach here: #99500

NimishMishra · 2024-07-18T14:47:35Z

Hi, just removing the current check doesn't quite do the right thing, as the value of the lastprivate variable will be a random one of the last values of the inner loops, rather than the definite final one. To implement this we need to check that we have reached the final value of all the loop indices, not just the outer one.

I've submitted a patch with this approach here: #99500

Hi @DavidTruby

Thanks for the PR. I see what you mean. I'll close this then.

I've an internal reprioritization for a few days; I do plan to return to working on this and implementing the delayed privatization approach, as comments here (and your PR) mention. But in the meantime, we can go ahead with your PR.

NimishMishra · 2024-07-18T14:50:04Z

Closing as the issue is covered by #99500

[flang][OpenMP] Enable lastprivate with collapse

77c0e18

llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir flang:openmp labels May 30, 2024

NimishMishra requested review from kiranchandramohan, ergawy, tblah, harishch4 and kiranktp May 30, 2024 07:56

tblah approved these changes May 31, 2024

View reviewed changes

kiranchandramohan requested a review from luporl May 31, 2024 16:16

luporl reviewed Jun 3, 2024

View reviewed changes

ergawy approved these changes Jun 7, 2024

View reviewed changes

NimishMishra closed this Jul 18, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[flang][OpenMP] Enable lastprivate with collapse #93778

[flang][OpenMP] Enable lastprivate with collapse #93778

Uh oh!

NimishMishra commented May 30, 2024

Uh oh!

llvmbot commented May 30, 2024

Uh oh!

llvmbot commented May 30, 2024

Uh oh!

tblah left a comment

Uh oh!

tblah May 31, 2024

Uh oh!

kiranchandramohan commented May 31, 2024

Uh oh!

luporl left a comment

Uh oh!

kiranchandramohan commented Jun 3, 2024

Uh oh!

ergawy left a comment

Uh oh!

ergawy Jun 7, 2024

Uh oh!

DavidTruby commented Jul 18, 2024

Uh oh!

NimishMishra commented Jul 18, 2024

Uh oh!

NimishMishra commented Jul 18, 2024

Uh oh!

Uh oh!

[flang][OpenMP] Enable lastprivate with collapse #93778

[flang][OpenMP] Enable lastprivate with collapse #93778

Uh oh!

Conversation

NimishMishra commented May 30, 2024

Uh oh!

llvmbot commented May 30, 2024

Uh oh!

llvmbot commented May 30, 2024

Uh oh!

tblah left a comment

Choose a reason for hiding this comment

Uh oh!

tblah May 31, 2024

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan commented May 31, 2024

Uh oh!

luporl left a comment

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan commented Jun 3, 2024

Uh oh!

ergawy left a comment

Choose a reason for hiding this comment

Uh oh!

ergawy Jun 7, 2024

Choose a reason for hiding this comment

Uh oh!

DavidTruby commented Jul 18, 2024

Uh oh!

NimishMishra commented Jul 18, 2024

Uh oh!

NimishMishra commented Jul 18, 2024

Uh oh!

Uh oh!