[flang][OpenMP] Fix privatization of threadprivate common block

[luporl]

In some cases, when privatizing a threadprivate common block, the
original symbol will correspond to the common block, instead of
its threadprivate version. This can happen, for instance, with a
common block, declared in a separate module, used by a parent
procedure and privatized in its child procedure. In this case,
symbol lookup won't find a symbol in the parent procedure, but
only in the module where the common block was defined.

Fixes #65028

[llvmbot]

@llvm/pr-subscribers-flang-openmp

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

Author: Leandro Lupori (luporl)

Changes

In some cases, when privatizing a threadprivate common block, the
original symbol will correspond to the common block, instead of
its threadprivate version. This can happen, for instance, with a
common block, declared in a separate module, used by a parent
procedure and privatized in its child procedure. In this case,
symbol lookup won't find a symbol in the parent procedure, but
only in the module where the common block was defined.

Fixes #65028

---

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

2 Files Affected:

• (modified) flang/lib/Lower/OpenMP.cpp (+11-9)
• (added) flang/test/Lower/OpenMP/threadprivate-commonblock-use.f90 (+29)

diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index c3a570bf15ea0d..ce2ee4befddcf1 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2029,18 +2029,20 @@ static void threadPrivatizeVars(Fortran::lower::AbstractConverter &converter,
   mlir::OpBuilder::InsertPoint insPt = firOpBuilder.saveInsertionPoint();
   firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
 
-  // Get the original ThreadprivateOp corresponding to the symbol and use the
-  // symbol value from that operation to create one ThreadprivateOp copy
-  // operation inside the parallel region.
+  // If the symbol corresponds to the original ThreadprivateOp, use the symbol
+  // value from that operation to create one ThreadprivateOp copy operation
+  // inside the parallel region.
+  // In some cases, however, the symbol will correspond to the original,
+  // non-threadprivate variable. This can happen, for instance, with a common
+  // block, declared in a separate module, used by a parent procedure and
+  // privatized in its child procedure.
   auto genThreadprivateOp = [&](Fortran::lower::SymbolRef sym) -> mlir::Value {
-    mlir::Value symOriThreadprivateValue = converter.getSymbolAddress(sym);
-    mlir::Operation *op = symOriThreadprivateValue.getDefiningOp();
+    mlir::Value symValue = converter.getSymbolAddress(sym);
+    mlir::Operation *op = symValue.getDefiningOp();
     if (auto declOp = mlir::dyn_cast<hlfir::DeclareOp>(op))
       op = declOp.getMemref().getDefiningOp();
-    assert(mlir::isa<mlir::omp::ThreadprivateOp>(op) &&
-           "Threadprivate operation not created");
-    mlir::Value symValue =
-        mlir::dyn_cast<mlir::omp::ThreadprivateOp>(op).getSymAddr();
+    if (mlir::isa<mlir::omp::ThreadprivateOp>(op))
+      symValue = mlir::dyn_cast<mlir::omp::ThreadprivateOp>(op).getSymAddr();
     return firOpBuilder.create<mlir::omp::ThreadprivateOp>(
         currentLocation, symValue.getType(), symValue);
   };
diff --git a/flang/test/Lower/OpenMP/threadprivate-commonblock-use.f90 b/flang/test/Lower/OpenMP/threadprivate-commonblock-use.f90
new file mode 100644
index 00000000000000..28616f7595a0d6
--- /dev/null
+++ b/flang/test/Lower/OpenMP/threadprivate-commonblock-use.f90
@@ -0,0 +1,29 @@
+! This test checks lowering of OpenMP Threadprivate Directive.
+! Test for common block, defined in one module, used in a subroutine of
+! another module and privatized in a nested subroutine.
+
+!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+!CHECK: fir.global common @cmn_(dense<0> : vector<4xi8>) : !fir.array<4xi8>
+module m0
+  common /cmn/ k1
+  !$omp threadprivate(/cmn/)
+end
+
+module  m1
+contains
+  subroutine ss1
+    use m0
+  contains
+!CHECK-LABEL: func @_QMm1Fss1Pss2
+!CHECK: %[[CMN:.*]] = fir.address_of(@cmn_) : !fir.ref<!fir.array<4xi8>>
+!CHECK: omp.parallel
+!CHECK: %{{.*}} = omp.threadprivate %[[CMN]] : !fir.ref<!fir.array<4xi8>> -> !fir.ref<!fir.array<4xi8>>
+    subroutine ss2
+      !$omp parallel copyin (k1)
+      !$omp end parallel
+    end subroutine ss2
+  end subroutine ss1
+end
+
+end

[kiranchandramohan]

There was a patch some time back that fixed a similar issue. That patch looked at it from the point of creating the missing symbols. https://reviews.llvm.org/D145684

[NimishMishra]

There was a patch some time back that fixed a similar issue. That patch looked at it from the point of creating the missing symbols. https://reviews.llvm.org/D145684

@kiranchandramohan Which of the two approaches seems better? I think this patch handles it better than explicit symbol creation in https://reviews.llvm.org/D145684. I generated the IR for some tests (including the one in https://reviews.llvm.org/D145684) with this PR. All tests dumped similar IR:

fir.global @<module_var>

func.func @<function_name> attributes ... {
    %0 = fir.address_of(module_var)
    %1 = fir.declare %0 ...
    omp.parallel {
          %2 = omp.threadprivate %1
          %3 = fir.declare %2
          fir.load %1
          fir.store %4 to %3
    }
}

Even host-association tests in the linked patch (below) works ok. For instance,

module tpmod
  real :: val = 2.0
end module

subroutine tptest2
  use tpmod
  call tps()
contains
  subroutine tps()
    !$omp parallel
    !$omp single
      val = 1.0
    !$omp end single
    !$omp end parallel
  end subroutine
end subroutine

have the following IR with this PR:

fir.global @_QMtpmodEval : f32 {
   %cst = arith.constant 2.000000e+00 : f32
   fir.has_value %cst : f32
}

func.func @_QFtotest2Ptps() attributes : {
   %cst = arith.constant 1.000000e+00 : f32
   %0 = fir.address_of(@_QMtpmodEval)
   %1 = fir.declare %0 
   omp.parallel {
      omp.single {
          fir.store %cst to %1
          omp.terminator
      }
     omp.terminator
   }
   return
}

Which seems like the intended behaviour: capturing the globally declared variable, and updating it within single as required.

So I wonder if we can do with the current approach, and not by creating extraneous symbols.

[luporl]

@kiranchandramohan I also have the same question as @NimishMishra. Do you think it would be better to create symbols in cases similar to #65028, instead of using this patch's approach? What would be the advantage of having these new symbols created?

[kiranchandramohan]

In general, having symbols created correctly helps the frontend perform semantic checks accurately. Originally, it was also supposed to help lowering, but at some point lowering started using nested symbol tables instead of relying on separate symbols per region. I cant think of anything that we might fail to catch.

The patch definitely looks simpler than https://reviews.llvm.org/D145684 and we can go ahead with it.

[kiranchandramohan]

Is it possible to have some other assertion to ensure that we are operating on a symbol which has a threadprivate directive?

[luporl]

I've added an assert to ensure the symbol has the threadprivate flag set in it.

[luporl]

Thanks for the review @kiranchandramohan!
If something comes up in the future that we are failing to catch we may revisit this part.