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

[kiranchandramohan]

Summary

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

[tblah]

LGTM, thanks for this!

flang/lib/Lower/OpenMP.cpp

 #include "flang/Lower/Bridge.h"
 #include "flang/Lower/ConvertExpr.h"
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/StatementContext.h"
 #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"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
 using DeclareTargetCapturePair =
     std::pair<mlir::omp::DeclareTargetCaptureClause,
               Fortran::semantics::Symbol>;

 static void threadPrivatizeVars(Fortran::lower::AbstractConverter &converter,
                                 Fortran::lower::pft::Evaluation &eval) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   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 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>(
         currentLocation, symValue.getType(), symValue);
   };
 
   llvm::SetVector<const Fortran::semantics::Symbol *> threadprivateSyms;
   converter.collectSymbolSet(

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

+! 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