[Flang][OpenMP] Fix lowering for threadprivate with HLFIR

HLFIR inserts hlfir.declare for threadprivate variables. Seek over
the hlfir.declare to find the threadprivate global symbol.

Reviewed By: tblah

Differential Revision: https://reviews.llvm.org/D157730

flang/lib/Lower/OpenMP.cpp

@@ -20,6 +20,7 @@
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Semantics/tools.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
@@ -1845,13 +1846,15 @@ static void threadPrivatizeVars(Fortran::lower::AbstractConverter &converter,
   firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
 
   // Get the original ThreadprivateOp corresponding to the symbol and use the
-  // symbol value from that opeartion to create one ThreadprivateOp copy
+  // symbol value from that operation to create one ThreadprivateOp copy
   // operation inside the parallel region.
   auto genThreadprivateOp = [&](Fortran::lower::SymbolRef sym) -> mlir::Value {
     mlir::Value symOriThreadprivateValue = converter.getSymbolAddress(sym);
     mlir::Operation *op = symOriThreadprivateValue.getDefiningOp();
+    if (auto declOp = mlir::dyn_cast<hlfir::DeclareOp>(op))
+      op = declOp.getMemref().getDefiningOp();
     assert(mlir::isa<mlir::omp::ThreadprivateOp>(op) &&
-           "The threadprivate operation not created");
+           "Threadprivate operation not created");
     mlir::Value symValue =
         mlir::dyn_cast<mlir::omp::ThreadprivateOp>(op).getSymAddr();
     return firOpBuilder.create<mlir::omp::ThreadprivateOp>(

flang/test/Lower/OpenMP/threadprivate-hlfir.f90

@@ -0,0 +1,26 @@
+! Simple test for lowering of OpenMP Threadprivate Directive with HLFIR.
+
+!RUN: %flang_fc1 -flang-experimental-hlfir -emit-hlfir -fopenmp %s -o - | FileCheck %s
+!RUN: bbc -hlfir -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+!CHECK-LABEL: @_QPsub
+!CHECK:    %[[ADDR:.*]] = fir.address_of(@_QFsubEa) : !fir.ref<i32>
+!CHECK:    %[[DECL:.*]]:2 = hlfir.declare %[[ADDR]] {uniq_name = "_QFsubEa"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:    %[[TP:.*]] = omp.threadprivate %[[DECL]]#1 : !fir.ref<i32> -> !fir.ref<i32>
+!CHECK:    %[[TP_DECL:.*]]:2 = hlfir.declare %[[TP:.*]] {uniq_name = "_QFsubEa"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:    omp.parallel   {
+!CHECK:      %[[TP_PARALLEL:.*]] = omp.threadprivate %[[DECL]]#1 : !fir.ref<i32> -> !fir.ref<i32>
+!CHECK:      %[[TP_PARALLEL_DECL:.*]]:2 = hlfir.declare %[[TP_PARALLEL]] {uniq_name = "_QFsubEa"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:      %[[TP_VAL:.*]] = fir.load %[[TP_PARALLEL_DECL]]#0 : !fir.ref<i32>
+!CHECK:      %{{.*}} = fir.call @_FortranAioOutputInteger32(%{{.*}}, %[[TP_VAL]]) fastmath<contract> : (!fir.ref<i8>, i32) -> i1
+!CHECK:      omp.terminator
+
+!CHECK:  fir.global internal @_QFsubEa : i32
+
+subroutine sub()
+  integer, save:: a
+  !$omp threadprivate(a)
+  !$omp parallel
+    print *, a
+  !$omp end parallel
+end subroutine