From 096f79523816bf0e593e50e9976d2faed6b699d5 Mon Sep 17 00:00:00 2001
From: Leandro Lupori <leandro.lupori@linaro.org>
Date: Tue, 4 Jun 2024 17:52:25 +0200
Subject: [PATCH] [flang][OpenMP] Fix privatization when critical is present

When a critical construct is present inside another construct where
privatizations may occur, such as a parallel construct, some
privatizations are skipped if the corresponding symbols are defined
inside the critical section only (see the example below).

This happens because, while critical constructs have a "body", they
don't have a separate scope (which makes sense, since no
privatizations can occur in them). Because of this, in semantics
phase, it's not possible to insert a new host association symbol,
but instead the symbol from the enclosing context is used directly.

This makes symbol collection in DataSharingProcessor consider the
new symbol to be defined by the critical construct, instead of by
the enclosing one, which causes the privatization to be skipped.

Example:
```
!$omp parallel default(firstprivate)
  !$omp critical
     i = 200
  !$omp end critical
!$omp end parallel
```

This patch fixes this by identifying constructs where
privatizations may not happen and skipping them during the
collection of nested symbols. Currently, this seems to happen only
with critical constructs, but others can be easily added to the
skip list, if needed.

Fixes https://github.com/llvm/llvm-project/issues/75767
---
 .../lib/Lower/OpenMP/DataSharingProcessor.cpp | 11 ++++++++-
 flang/test/Lower/OpenMP/critical.f90          | 24 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
index 557a9685024c5..ba8cd44ab9150 100644
--- a/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/DataSharingProcessor.cpp
@@ -262,6 +262,15 @@ void DataSharingProcessor::insertLastPrivateCompare(mlir::Operation *op) {
   }
 }
 
+static bool mayHavePrivatizations(const lower::pft::Evaluation &eval) {
+  if (const parser::OpenMPConstruct *construct =
+          eval.getIf<parser::OpenMPConstruct>()) {
+    if (std::holds_alternative<parser::OpenMPCriticalConstruct>(construct->u))
+      return false;
+  }
+  return true;
+}
+
 static const parser::CharBlock *
 getSource(const semantics::SemanticsContext &semaCtx,
           const lower::pft::Evaluation &eval) {
@@ -305,7 +314,7 @@ void DataSharingProcessor::collectSymbolsInNestedRegions(
     llvm::SetVector<const semantics::Symbol *> &symbolsInNestedRegions) {
   for (lower::pft::Evaluation &nestedEval : eval.getNestedEvaluations()) {
     if (nestedEval.hasNestedEvaluations()) {
-      if (nestedEval.isConstruct())
+      if (nestedEval.isConstruct() || !mayHavePrivatizations(nestedEval))
         // Recursively look for OpenMP constructs within `nestedEval`'s region
         collectSymbolsInNestedRegions(nestedEval, flag, symbolsInNestedRegions);
       else {
diff --git a/flang/test/Lower/OpenMP/critical.f90 b/flang/test/Lower/OpenMP/critical.f90
index d62c58b3081ad..53c7866c46658 100644
--- a/flang/test/Lower/OpenMP/critical.f90
+++ b/flang/test/Lower/OpenMP/critical.f90
@@ -51,3 +51,27 @@ subroutine predetermined_privatization()
   end do
   !$omp end parallel do
 end
+
+! https://github.com/llvm/llvm-project/issues/75767
+!CHECK-LABEL: func @_QPnested_privatization(
+subroutine nested_privatization()
+  integer :: i
+
+  !CHECK: %[[I:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFnested_privatizationEi"}
+  !CHECK: %[[I_DECL:.*]]:2 = hlfir.declare %[[I]] {uniq_name = "_QFnested_privatizationEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+  !CHECK: omp.parallel {
+  !CHECK:   %[[PRIV_I:.*]] = fir.alloca i32 {bindc_name = "i", pinned, uniq_name = "_QFnested_privatizationEi"}
+  !CHECK:   %[[PRIV_I_DECL:.*]]:2 = hlfir.declare %[[PRIV_I]] {uniq_name = "_QFnested_privatizationEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+  !CHECK:   %[[TEMP:.*]] = fir.load %[[I_DECL]]#0 : !fir.ref<i32>
+  !CHECK:   hlfir.assign %[[TEMP]] to %[[PRIV_I_DECL]]#0 temporary_lhs : i32, !fir.ref<i32>
+  !$omp parallel default(firstprivate)
+    !CHECK: omp.critical {
+    !$omp critical
+      !CHECK: %[[C200:.*]] = arith.constant 200 : i32
+      !CHECK: hlfir.assign %[[C200]] to %[[PRIV_I_DECL]]#0 : i32, !fir.ref<i32>
+      i = 200
+    !CHECK: }
+    !$omp end critical
+  !CHECK: }
+  !$omp end parallel
+end subroutine