[flang] Lower general forall statement

This patch lowers general forall statements. The forall
are lowered to nested loops.

This patch is part of the upstreaming effort from fir-dev branch.

Depends on D121385

Reviewed By: PeteSteinfeld, schweitz

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

Co-authored-by: V Donaldson <vdonaldson@nvidia.com>
Co-authored-by: Jean Perier <jperier@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

flang/include/flang/Lower/Allocatable.h

@@ -24,15 +24,22 @@ class Location;
 
 namespace fir {
 class MutableBoxValue;
-} // namespace fir
+}
 
 namespace Fortran::parser {
 struct AllocateStmt;
 struct DeallocateStmt;
 } // namespace Fortran::parser
 
+namespace Fortran::evaluate {
+template <typename T>
+class Expr;
+struct SomeType;
+} // namespace Fortran::evaluate
+
 namespace Fortran::lower {
 class AbstractConverter;
+class StatementContext;
 
 namespace pft {
 struct Variable;
@@ -48,13 +55,23 @@ void genDeallocateStmt(Fortran::lower::AbstractConverter &,
 
 /// Create a MutableBoxValue for an allocatable or pointer entity.
 /// If the variables is a local variable that is not a dummy, it will be
-/// initialized to unallocated/disassociated status.
+/// initialized to unallocated/diassociated status.
 fir::MutableBoxValue createMutableBox(Fortran::lower::AbstractConverter &,
                                       mlir::Location,
                                       const Fortran::lower::pft::Variable &var,
                                       mlir::Value boxAddr,
                                       mlir::ValueRange nonDeferredParams);
 
+/// Update a MutableBoxValue to describe the entity designated by the expression
+/// \p source. This version takes care of \p source lowering.
+/// If \lbounds is not empty, it is used to defined the MutableBoxValue
+/// lower bounds, otherwise, the lower bounds from \p source are used.
+void associateMutableBox(
+    Fortran::lower::AbstractConverter &, mlir::Location,
+    const fir::MutableBoxValue &,
+    const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &source,
+    mlir::ValueRange lbounds, Fortran::lower::StatementContext &);
+
 } // namespace Fortran::lower
 
 #endif // FORTRAN_LOWER_ALLOCATABLE_H

flang/include/flang/Lower/ConvertExpr.h

@@ -100,7 +100,10 @@ fir::ExtendedValue createSomeArrayBox(AbstractConverter &converter,
 /// The returned value is null otherwise.
 mlir::Value createSubroutineCall(AbstractConverter &converter,
                                  const evaluate::ProcedureRef &call,
-                                 SymMap &symMap, StatementContext &stmtCtx);
+                                 ExplicitIterSpace &explicitIterSpace,
+                                 ImplicitIterSpace &implicitIterSpace,
+                                 SymMap &symMap, StatementContext &stmtCtx,
+                                 bool isUserDefAssignment);
 
 /// Create the address of the box.
 /// \p expr must be the designator of an allocatable/pointer entity.

flang/include/flang/Optimizer/Builder/BoxValue.h

@@ -24,6 +24,8 @@
 #include <utility>
 
 namespace fir {
+class FirOpBuilder;
+
 class CharBoxValue;
 class ArrayBoxValue;
 class CharArrayBoxValue;
@@ -402,6 +404,12 @@ bool isArray(const ExtendedValue &exv);
 /// Get the type parameters for `exv`.
 llvm::SmallVector<mlir::Value> getTypeParams(const ExtendedValue &exv);
 
+/// Get exactly one extent for any array-like extended value, \p exv. If \p exv
+/// is not an array or has rank less then \p dim, the result will be a nullptr.
+mlir::Value getExtentAtDimension(const ExtendedValue &exv,
+                                 FirOpBuilder &builder, mlir::Location loc,
+                                 unsigned dim);
+
 /// An extended value is a box of values pertaining to a discrete entity. It is
 /// used in lowering to track all the runtime values related to an entity. For
 /// example, an entity may have an address in memory that contains its value(s)

flang/include/flang/Optimizer/Builder/Runtime/Inquiry.h

@@ -0,0 +1,46 @@
+//===-- Inquiry.h - generate inquiry runtime API calls ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_RUNTIME_INQUIRY_H
+#define FORTRAN_OPTIMIZER_BUILDER_RUNTIME_INQUIRY_H
+
+namespace mlir {
+class Value;
+class Location;
+} // namespace mlir
+
+namespace fir {
+class FirOpBuilder;
+}
+
+namespace fir::runtime {
+
+/// Generate call to general `LboundDim` runtime routine.  Calls to LBOUND
+/// without a DIM argument get transformed into descriptor inquiries so they're
+/// not handled in the runtime.
+mlir::Value genLboundDim(fir::FirOpBuilder &builder, mlir::Location loc,
+                         mlir::Value array, mlir::Value dim);
+
+/// Generate call to general `Ubound` runtime routine.  Calls to UBOUND
+/// with a DIM argument get transformed into an expression equivalent to
+/// SIZE() + LBOUND() - 1, so they don't have an intrinsic in the runtime.
+void genUbound(fir::FirOpBuilder &builder, mlir::Location loc,
+               mlir::Value resultBox, mlir::Value array, mlir::Value kind);
+
+/// Generate call to `Size` runtime routine. This routine is a specialized
+/// version when the DIM argument is not specified by the user.
+mlir::Value genSize(fir::FirOpBuilder &builder, mlir::Location loc,
+                    mlir::Value array);
+
+/// Generate call to general `SizeDim` runtime routine.  This version is for
+/// when the user specifies a DIM argument.
+mlir::Value genSizeDim(fir::FirOpBuilder &builder, mlir::Location loc,
+                       mlir::Value array, mlir::Value dim);
+
+} // namespace fir::runtime
+#endif // FORTRAN_OPTIMIZER_BUILDER_RUNTIME_INQUIRY_H

flang/lib/Lower/Allocatable.cpp

@@ -666,3 +666,33 @@ fir::MutableBoxValue Fortran::lower::createMutableBox(
     fir::factory::disassociateMutableBox(builder, loc, box);
   return box;
 }
+
+//===----------------------------------------------------------------------===//
+// MutableBoxValue reading interface implementation
+//===----------------------------------------------------------------------===//
+
+static bool
+isArraySectionWithoutVectorSubscript(const Fortran::lower::SomeExpr &expr) {
+  return expr.Rank() > 0 && Fortran::evaluate::IsVariable(expr) &&
+         !Fortran::evaluate::UnwrapWholeSymbolDataRef(expr) &&
+         !Fortran::evaluate::HasVectorSubscript(expr);
+}
+
+void Fortran::lower::associateMutableBox(
+    Fortran::lower::AbstractConverter &converter, mlir::Location loc,
+    const fir::MutableBoxValue &box, const Fortran::lower::SomeExpr &source,
+    mlir::ValueRange lbounds, Fortran::lower::StatementContext &stmtCtx) {
+  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+  if (Fortran::evaluate::UnwrapExpr<Fortran::evaluate::NullPointer>(source)) {
+    fir::factory::disassociateMutableBox(builder, loc, box);
+    return;
+  }
+  // The right hand side must not be evaluated in a temp.
+  // Array sections can be described by fir.box without making a temp.
+  // Otherwise, do not generate a fir.box to avoid having to later use a
+  // fir.rebox to implement the pointer association.
+  fir::ExtendedValue rhs = isArraySectionWithoutVectorSubscript(source)
+                               ? converter.genExprBox(source, stmtCtx, loc)
+                               : converter.genExprAddr(source, stmtCtx);
+  fir::factory::associateMutableBox(builder, loc, box, rhs, lbounds);
+}