[Flang][OpenMP] Fix loop index privatisation in Parallel region in HLFIR

[kiranchandramohan]

Summary

HLFIR lowering always adds hlfir.declare when symbols are bound to their
address allocated on the stack. Ensure that the declare is placed along
with the alloca if it is hoisted. And always return the mlir value that
is bound to the symbol (i.e the alloca in FIR lowering and the declare
in HLFIR lowering).

Context: Loop index variables in OpenMP parallel regions should be
privatised to work correctly.

[tblah]

LGTM

flang/include/flang/Optimizer/Builder/FIRBuilder.h

mlir::Value allocateLocal(mlir::Location loc, mlir::Type ty,
                             llvm::ArrayRef<mlir::Value> lenParams,
                             bool asTarget = false);
   mlir::Value allocateLocal(mlir::Location loc, mlir::Type ty,
                             llvm::StringRef uniqName, llvm::StringRef name,
                             llvm::ArrayRef<mlir::Value> shape,
                             llvm::ArrayRef<mlir::Value> lenParams,
                             bool asTarget = false);
 
+  /// Create a temporary using `fir.alloca`. This function does not hoist.
+  /// It is the callers responsibility to set the insertion point if
+  /// hoisting is required.
+  mlir::Value
+  createTemporaryAlloc(mlir::Location loc, mlir::Type type,
+                       llvm::StringRef name, mlir::ValueRange lenParams = {},
+                       mlir::ValueRange shape = {},
+                       llvm::ArrayRef<mlir::NamedAttribute> attrs = {});
+
   /// Create a temporary. A temp is allocated using `fir.alloca` and can be read
   /// and written using `fir.load` and `fir.store`, resp.  The temporary can be
   /// given a name via a front-end `Symbol` or a `StringRef`.
   mlir::Value createTemporary(mlir::Location loc, mlir::Type type,
                               llvm::StringRef name = {},
                               mlir::ValueRange shape = {},
                               mlir::ValueRange lenParams = {},
                               llvm::ArrayRef<mlir::NamedAttribute> attrs = {});

flang/lib/Lower/Bridge.cpp

private:
   mlir::Value genLoopVariableAddress(mlir::Location loc,
                                      const Fortran::semantics::Symbol &sym,
                                      bool isUnordered) {
     if (isUnordered || sym.has<Fortran::semantics::HostAssocDetails>() ||
         sym.has<Fortran::semantics::UseDetails>()) {
       if (!shallowLookupSymbol(sym)) {
         // Do concurrent loop variables are not mapped yet since they are local
         // to the Do concurrent scope (same for OpenMP loops).
-        auto newVal = builder->createTemporary(loc, genType(sym),
-                                               toStringRef(sym.name()));
-        bindIfNewSymbol(sym, newVal);
-        return newVal;
+        mlir::OpBuilder::InsertPoint insPt = builder->saveInsertionPoint();
+        builder->setInsertionPointToStart(builder->getAllocaBlock());
+        mlir::Type tempTy = genType(sym);
+        mlir::Value temp =
+            builder->createTemporaryAlloc(loc, tempTy, toStringRef(sym.name()));
+        bindIfNewSymbol(sym, temp);
+        builder->restoreInsertionPoint(insPt);
       }
     }
     auto entry = lookupSymbol(sym);
     (void)entry;
     assert(entry && "loop control variable must already be in map");
     Fortran::lower::StatementContext stmtCtx;
     return fir::getBase(
         genExprAddr(Fortran::evaluate::AsGenericExpr(sym).value(), stmtCtx));

flang/lib/Optimizer/Builder/FIRBuilder.cpp

 
 /// Get the block for adding Allocas.
 mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
   auto iface =
       getRegion().getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>();
   return iface ? iface.getAllocaBlock() : getEntryBlock();
 }
 
+mlir::Value fir::FirOpBuilder::createTemporaryAlloc(
+    mlir::Location loc, mlir::Type type, llvm::StringRef name,
+    mlir::ValueRange lenParams, mlir::ValueRange shape,
+    llvm::ArrayRef<mlir::NamedAttribute> attrs) {
+  assert(!type.isa<fir::ReferenceType>() && "cannot be a reference");
+  // If the alloca is inside an OpenMP Op which will be outlined then pin
+  // the alloca here.
+  const bool pinned =
+      getRegion().getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>();
+  mlir::Value temp =
+      create<fir::AllocaOp>(loc, type, /*unique_name=*/llvm::StringRef{}, name,
+                            pinned, lenParams, shape, attrs);
+  return temp;
+}
+
 /// Create a temporary variable on the stack. Anonymous temporaries have no
 /// `name` value. Temporaries do not require a uniqued name.
 mlir::Value
 fir::FirOpBuilder::createTemporary(mlir::Location loc, mlir::Type type,
                                    llvm::StringRef name, mlir::ValueRange shape,
                                    mlir::ValueRange lenParams,
                                    llvm::ArrayRef<mlir::NamedAttribute> attrs) {
   llvm::SmallVector<mlir::Value> dynamicShape =
       elideExtentsAlreadyInType(type, shape);
   llvm::SmallVector<mlir::Value> dynamicLength =
       elideLengthsAlreadyInType(type, lenParams);
   InsertPoint insPt;
   const bool hoistAlloc = dynamicShape.empty() && dynamicLength.empty();
   if (hoistAlloc) {
     insPt = saveInsertionPoint();
     setInsertionPointToStart(getAllocaBlock());
   }
 
-  // If the alloca is inside an OpenMP Op which will be outlined then pin the
-  // alloca here.
-  const bool pinned =
-      getRegion().getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>();
-  assert(!type.isa<fir::ReferenceType>() && "cannot be a reference");
-  auto ae =
-      create<fir::AllocaOp>(loc, type, /*unique_name=*/llvm::StringRef{}, name,
-                            pinned, dynamicLength, dynamicShape, attrs);
+  mlir::Value ae =
+      createTemporaryAlloc(loc, type, name, dynamicLength, dynamicShape, attrs);
+
   if (hoistAlloc)
     restoreInsertionPoint(insPt);
   return ae;
 }
 
 mlir::Value fir::FirOpBuilder::createHeapTemporary(
     mlir::Location loc, mlir::Type type, llvm::StringRef name,
     mlir::ValueRange shape, mlir::ValueRange lenParams,

flang/test/Lower/OpenMP/hlfir-seqloop-parallel.f90

+! This test checks lowering of sequential loops in OpenMP parallel.
+! The loop indices of these loops should be privatised.
+
+! RUN: bbc -hlfir -fopenmp -emit-hlfir %s -o - | FileCheck %s
+! RUN: %flang_fc1 -emit-hlfir -flang-experimental-hlfir -fopenmp %s -o - | FileCheck %s
+
+
+subroutine sb1
+  integer i
+  !$omp parallel
+  do i=1,10
+  end do
+  !$omp end parallel
+
+end subroutine
+
+!CHECK-LABEL:  @_QPsb1
+!CHECK:    %[[I_ADDR:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsb1Ei"}
+!CHECK:    %[[I_DECL:.*]]:2 = hlfir.declare %[[I_ADDR]] {uniq_name = "_QFsb1Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:    omp.parallel   {
+!CHECK:      %[[I_PVT_ADDR:.*]] = fir.alloca i32 {bindc_name = "i", pinned}
+!CHECK:      %[[I_PVT_DECL:.*]]:2 = hlfir.declare %[[I_PVT_ADDR]] {uniq_name = "_QFsb1Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:      %[[I_FINAL_VAL:.*]]:2 = fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[I_VAL:.*]] = %{{.*}}) -> (index, i32) {
+!CHECK:        fir.store %[[I_VAL]] to %[[I_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:      }
+!CHECK:      fir.store %[[I_FINAL_VAL]]#1 to %[[I_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:      omp.terminator
+!CHECK:    }
+!CHECK:    return
+!CHECK:  }
+
+subroutine sb2
+  integer i, j, k
+  !$omp parallel
+  do j=1,10
+    if (k .eq. 1) then
+      do i=20, 30
+      end do
+    endif
+
+    do i=40,50
+    end do
+  end do
+  !$omp end parallel
+end subroutine
+!CHECK-LABEL:  @_QPsb2
+!CHECK:    %[[I_ADDR:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsb2Ei"}
+!CHECK:    %[[I_DECL:.*]]:2 = hlfir.declare %[[I_ADDR]] {uniq_name = "_QFsb2Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:    %[[J_ADDR:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFsb2Ej"}
+!CHECK:    %[[J_DECL:.*]]:2 = hlfir.declare %[[J_ADDR]] {uniq_name = "_QFsb2Ej"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:    %[[K_ADDR:.*]] = fir.alloca i32 {bindc_name = "k", uniq_name = "_QFsb2Ek"}
+!CHECK:    %[[K_DECL:.*]]:2 = hlfir.declare %[[K_ADDR]] {uniq_name = "_QFsb2Ek"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:    omp.parallel   {
+!CHECK:      %[[I_PVT_ADDR:.*]] = fir.alloca i32 {bindc_name = "i", pinned}
+!CHECK:      %[[I_PVT_DECL:.*]]:2 = hlfir.declare %[[I_PVT_ADDR]] {uniq_name = "_QFsb2Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:      %[[J_PVT_ADDR:.*]] = fir.alloca i32 {bindc_name = "j", pinned}
+!CHECK:      %[[J_PVT_DECL:.*]]:2 = hlfir.declare %[[J_PVT_ADDR]] {uniq_name = "_QFsb2Ej"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:      %[[FINAL_J_VAL:.*]]:2 = fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[J_VAL:.*]] = %{{.*}}) -> (index, i32) {
+!CHECK:        fir.store %arg1 to %9#1 : !fir.ref<i32>
+!CHECK:        fir.if %{{.*}} {
+!CHECK:          %[[FINAL_I_VAL:.*]]:2 = fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[I_VAL:.*]] = %{{.*}}) -> (index, i32) {
+!CHECK:            fir.store %[[I_VAL]] to %[[I_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:          }
+!CHECK:          fir.store %[[FINAL_I_VAL]]#1 to %[[I_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:        }
+!CHECK:        %[[FINAL_I_VAL:.*]]:2 = fir.do_loop %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[I_VAL:.*]] = %{{.*}}) -> (index, i32) {
+!CHECK:          fir.store %[[I_VAL]] to %[[I_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:        }
+!CHECK:        fir.store %[[FINAL_I_VAL]]#1 to %[[I_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:      }
+!CHECK:      fir.store %[[FINAL_J_VAL]]#1 to %[[J_PVT_DECL]]#1 : !fir.ref<i32>
+!CHECK:      omp.terminator
+!CHECK:    }
+!CHECK:    return
+!CHECK:  }