[flang] Initial support of allocate statement with source

Support allocate statement with source in runtime version. The source
expression is evaluated only once for each allocate statement. When the
source expression has shape-spec, uses it for bounds. Otherwise, get
the bounds from the source expression. Get the length if the source
expression has deferred length parameter.

Reviewed By: clementval, jeanPerier

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

Author: PeixinQiao

flang/include/flang/Runtime/assign.h

@@ -6,15 +6,15 @@
 //
 //===----------------------------------------------------------------------===//
 
-// External and internal APIs for data assignment (both intrinsic assignment
-// and TBP defined generic ASSIGNMENT(=)).  Should be called by lowering
-// for any assignments possibly needing special handling.  Intrinsic assignment
-// to non-allocatable variables whose types are intrinsic need not come through
-// here (though they may do so).  Assignments to allocatables, and assignments
-// whose types may be polymorphic or are monomorphic and of derived types with
-// finalization, allocatable components, or components with type-bound defined
-// assignments, in the original type or the types of its non-pointer components
-// (recursively) must arrive here.
+// External APIs for data assignment (both intrinsic assignment and TBP defined
+// generic ASSIGNMENT(=)).  Should be called by lowering for any assignments
+// possibly needing special handling.  Intrinsic assignment to non-allocatable
+// variables whose types are intrinsic need not come through here (though they
+// may do so).  Assignments to allocatables, and assignments whose types may be
+// polymorphic or are monomorphic and of derived types with finalization,
+// allocatable components, or components with type-bound defined assignments, in
+// the original type or the types of its non-pointer components (recursively)
+// must arrive here.
 //
 // Non-type-bound generic INTERFACE ASSIGNMENT(=) is resolved in semantics and
 // need not be handled here in the runtime; ditto for type conversions on
@@ -27,14 +27,6 @@
 
 namespace Fortran::runtime {
 class Descriptor;
-class Terminator;
-
-// Assigns one object to another via intrinsic assignment (F'2018 10.2.1.3) or
-// type-bound (only!) defined assignment (10.2.1.4), as appropriate.  Performs
-// finalization, scalar expansion, & allocatable (re)allocation as needed.
-// Does not perform intrinsic assignment implicit type conversion.  Both
-// descriptors must be initialized.  Recurses as needed to handle components.
-void Assign(Descriptor &, const Descriptor &, Terminator &);
 
 extern "C" {
 // API for lowering assignment

flang/lib/Lower/Allocatable.cpp

@@ -183,6 +183,29 @@ static mlir::Value genRuntimeAllocate(fir::FirOpBuilder &builder,
   return builder.create<fir::CallOp>(loc, callee, operands).getResult(0);
 }
 
+/// Generate a sequence of runtime calls to allocate memory and assign with the
+/// \p source.
+static mlir::Value genRuntimeAllocateSource(fir::FirOpBuilder &builder,
+                                            mlir::Location loc,
+                                            const fir::MutableBoxValue &box,
+                                            fir::ExtendedValue source,
+                                            ErrorManager &errorManager) {
+  mlir::func::FuncOp callee =
+      box.isPointer()
+          ? fir::runtime::getRuntimeFunc<mkRTKey(PointerAllocateSource)>(
+                loc, builder)
+          : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableAllocateSource)>(
+                loc, builder);
+  llvm::SmallVector<mlir::Value> args{
+      box.getAddr(),           fir::getBase(source),
+      errorManager.hasStat,    errorManager.errMsgAddr,
+      errorManager.sourceFile, errorManager.sourceLine};
+  llvm::SmallVector<mlir::Value> operands;
+  for (auto [fst, snd] : llvm::zip(args, callee.getFunctionType().getInputs()))
+    operands.emplace_back(builder.createConvert(loc, snd, fst));
+  return builder.create<fir::CallOp>(loc, callee, operands).getResult(0);
+}
+
 /// Generate a runtime call to deallocate memory.
 static mlir::Value genRuntimeDeallocate(fir::FirOpBuilder &builder,
                                         mlir::Location loc,
@@ -255,8 +278,11 @@ class AllocateStmtHelper {
     visitAllocateOptions();
     lowerAllocateLengthParameters();
     errorManager.init(converter, loc, statExpr, errMsgExpr);
-    if (sourceExpr || moldExpr)
-      TODO(loc, "lower MOLD/SOURCE expr in allocate");
+    Fortran::lower::StatementContext stmtCtx;
+    if (sourceExpr)
+      sourceExv = converter.genExprBox(loc, *sourceExpr, stmtCtx);
+    if (moldExpr)
+      TODO(loc, "lower MOLD expr in allocate");
     mlir::OpBuilder::InsertPoint insertPt = builder.saveInsertionPoint();
     for (const auto &allocation :
          std::get<std::list<Fortran::parser::Allocation>>(stmt.t))
@@ -393,45 +419,13 @@ class AllocateStmtHelper {
     }
     // Generate a sequence of runtime calls.
     errorManager.genStatCheck(builder, loc);
-    if (box.isPointer()) {
-      // For pointers, the descriptor may still be uninitialized (see Fortran
-      // 2018 19.5.2.2). The allocation runtime needs to be given a descriptor
-      // with initialized rank, types and attributes. Initialize the descriptor
-      // here to ensure these constraints are fulfilled.
-      mlir::Value nullPointer = fir::factory::createUnallocatedBox(
-          builder, loc, box.getBoxTy(), box.nonDeferredLenParams());
-      builder.create<fir::StoreOp>(loc, nullPointer, box.getAddr());
-    } else {
-      assert(box.isAllocatable() && "must be an allocatable");
-      // For allocatables, sync the MutableBoxValue and descriptor before the
-      // calls in case it is tracked locally by a set of variables.
-      fir::factory::getMutableIRBox(builder, loc, box);
-    }
+    genAllocateObjectInit(box);
     if (alloc.hasCoarraySpec())
       TODO(loc, "coarray allocation");
     if (alloc.type.IsPolymorphic())
       genSetType(alloc, box, loc);
     genSetDeferredLengthParameters(alloc, box);
-    // Set bounds for arrays
-    mlir::Type idxTy = builder.getIndexType();
-    mlir::Type i32Ty = builder.getIntegerType(32);
-    Fortran::lower::StatementContext stmtCtx;
-    for (const auto &iter : llvm::enumerate(alloc.getShapeSpecs())) {
-      mlir::Value lb;
-      const auto &bounds = iter.value().t;
-      if (const std::optional<Fortran::parser::BoundExpr> &lbExpr =
-              std::get<0>(bounds))
-        lb = fir::getBase(converter.genExprValue(
-            loc, Fortran::semantics::GetExpr(*lbExpr), stmtCtx));
-      else
-        lb = builder.createIntegerConstant(loc, idxTy, 1);
-      mlir::Value ub = fir::getBase(converter.genExprValue(
-          loc, Fortran::semantics::GetExpr(std::get<1>(bounds)), stmtCtx));
-      mlir::Value dimIndex =
-          builder.createIntegerConstant(loc, i32Ty, iter.index());
-      // Runtime call
-      genRuntimeSetBounds(builder, loc, box, dimIndex, lb, ub);
-    }
+    genAllocateObjectBounds(alloc, box);
     mlir::Value stat = genRuntimeAllocate(builder, loc, box, errorManager);
     fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
     errorManager.assignStat(builder, loc, stat);
@@ -478,8 +472,87 @@ class AllocateStmtHelper {
       TODO(loc, "derived type length parameters in allocate");
   }
 
-  void genSourceAllocation(const Allocation &, const fir::MutableBoxValue &) {
-    TODO(loc, "SOURCE allocation");
+  void genAllocateObjectInit(const fir::MutableBoxValue &box) {
+    if (box.isPointer()) {
+      // For pointers, the descriptor may still be uninitialized (see Fortran
+      // 2018 19.5.2.2). The allocation runtime needs to be given a descriptor
+      // with initialized rank, types and attributes. Initialize the descriptor
+      // here to ensure these constraints are fulfilled.
+      mlir::Value nullPointer = fir::factory::createUnallocatedBox(
+          builder, loc, box.getBoxTy(), box.nonDeferredLenParams());
+      builder.create<fir::StoreOp>(loc, nullPointer, box.getAddr());
+    } else {
+      assert(box.isAllocatable() && "must be an allocatable");
+      // For allocatables, sync the MutableBoxValue and descriptor before the
+      // calls in case it is tracked locally by a set of variables.
+      fir::factory::getMutableIRBox(builder, loc, box);
+    }
+  }
+
+  void genAllocateObjectBounds(const Allocation &alloc,
+                               const fir::MutableBoxValue &box) {
+    // Set bounds for arrays
+    mlir::Type idxTy = builder.getIndexType();
+    mlir::Type i32Ty = builder.getIntegerType(32);
+    Fortran::lower::StatementContext stmtCtx;
+    for (const auto &iter : llvm::enumerate(alloc.getShapeSpecs())) {
+      mlir::Value lb;
+      const auto &bounds = iter.value().t;
+      if (const std::optional<Fortran::parser::BoundExpr> &lbExpr =
+              std::get<0>(bounds))
+        lb = fir::getBase(converter.genExprValue(
+            loc, Fortran::semantics::GetExpr(*lbExpr), stmtCtx));
+      else
+        lb = builder.createIntegerConstant(loc, idxTy, 1);
+      mlir::Value ub = fir::getBase(converter.genExprValue(
+          loc, Fortran::semantics::GetExpr(std::get<1>(bounds)), stmtCtx));
+      mlir::Value dimIndex =
+          builder.createIntegerConstant(loc, i32Ty, iter.index());
+      // Runtime call
+      genRuntimeSetBounds(builder, loc, box, dimIndex, lb, ub);
+    }
+    if (sourceExpr && sourceExpr->Rank() > 0 &&
+        alloc.getShapeSpecs().size() == 0) {
+      // If the alloc object does not have shape list, get the bounds from the
+      // source expression.
+      mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+      const auto *sourceBox = sourceExv.getBoxOf<fir::BoxValue>();
+      assert(sourceBox && "source expression should be lowered to one box");
+      for (int i = 0; i < sourceExpr->Rank(); ++i) {
+        auto dimVal = builder.createIntegerConstant(loc, idxTy, i);
+        auto dimInfo = builder.create<fir::BoxDimsOp>(
+            loc, idxTy, idxTy, idxTy, sourceBox->getAddr(), dimVal);
+        mlir::Value lb =
+            fir::factory::readLowerBound(builder, loc, sourceExv, i, one);
+        mlir::Value extent = dimInfo.getResult(1);
+        mlir::Value ub = builder.create<mlir::arith::SubIOp>(
+            loc, builder.create<mlir::arith::AddIOp>(loc, extent, lb), one);
+        mlir::Value dimIndex = builder.createIntegerConstant(loc, i32Ty, i);
+        genRuntimeSetBounds(builder, loc, box, dimIndex, lb, ub);
+      }
+    }
+  }
+
+  void genSourceAllocation(const Allocation &alloc,
+                           const fir::MutableBoxValue &box) {
+    // Generate a sequence of runtime calls.
+    errorManager.genStatCheck(builder, loc);
+    genAllocateObjectInit(box);
+    if (alloc.hasCoarraySpec())
+      TODO(loc, "coarray allocation");
+    if (alloc.type.IsPolymorphic())
+      TODO(loc, "polymorphic allocation with SOURCE specifier");
+    // Set length of the allocate object if it has. Otherwise, get the length
+    // from source for the deferred length parameter.
+    if (lenParams.empty() && box.isCharacter() &&
+        !box.hasNonDeferredLenParams())
+      lenParams.push_back(fir::factory::readCharLen(builder, loc, sourceExv));
+    genSetDeferredLengthParameters(alloc, box);
+    genAllocateObjectBounds(alloc, box);
+    mlir::Value stat =
+        genRuntimeAllocateSource(builder, loc, box, sourceExv, errorManager);
+    fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
+    errorManager.assignStat(builder, loc, stat);
   }
   void genMoldAllocation(const Allocation &, const fir::MutableBoxValue &) {
     TODO(loc, "MOLD allocation");
@@ -576,6 +649,8 @@ class AllocateStmtHelper {
   // value of the length parameters that were specified inside.
   llvm::SmallVector<mlir::Value> lenParams;
   ErrorManager errorManager;
+  // 9.7.1.2(7) The source-expr is evaluated exactly once for each AllocateStmt.
+  fir::ExtendedValue sourceExv;
 
   mlir::Location loc;
 };

flang/runtime/allocatable.cpp

@@ -7,11 +7,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Runtime/allocatable.h"
+#include "assign.h"
 #include "derived.h"
 #include "stat.h"
 #include "terminator.h"
 #include "type-info.h"
-#include "flang/Runtime/assign.h"
 
 namespace Fortran::runtime {
 extern "C" {
@@ -88,6 +88,22 @@ int RTNAME(AllocatableAllocate)(Descriptor &descriptor, bool hasStat,
   return stat;
 }
 
+int RTNAME(AllocatableAllocateSource)(Descriptor &alloc,
+    const Descriptor &source, bool hasStat, const Descriptor *errMsg,
+    const char *sourceFile, int sourceLine) {
+  if (alloc.Elements() == 0) {
+    return StatOk;
+  }
+  int stat{RTNAME(AllocatableAllocate)(
+      alloc, hasStat, errMsg, sourceFile, sourceLine)};
+  if (stat == StatOk) {
+    Terminator terminator{sourceFile, sourceLine};
+    // 9.7.1.2(7)
+    Assign(alloc, source, terminator, /*skipRealloc=*/true);
+  }
+  return stat;
+}
+
 int RTNAME(AllocatableDeallocate)(Descriptor &descriptor, bool hasStat,
     const Descriptor *errMsg, const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};
@@ -125,6 +141,6 @@ void RTNAME(AllocatableDeallocateNoFinal)(
   }
 }
 
-// TODO: AllocatableCheckLengthParameter, AllocatableAllocateSource
+// TODO: AllocatableCheckLengthParameter
 }
 } // namespace Fortran::runtime

flang/runtime/assign.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Runtime/assign.h"
+#include "assign.h"
 #include "derived.h"
 #include "stat.h"
 #include "terminator.h"
@@ -59,7 +60,8 @@ static void DoElementalDefinedAssignment(const Descriptor &to,
   }
 }
 
-void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator) {
+void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator,
+    bool skipRealloc) {
   DescriptorAddendum *toAddendum{to.Addendum()};
   const typeInfo::DerivedType *toDerived{
       toAddendum ? toAddendum->derivedType() : nullptr};
@@ -69,7 +71,7 @@ void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator) {
   bool wasJustAllocated{false};
   if (to.IsAllocatable()) {
     std::size_t lenParms{fromDerived ? fromDerived->LenParameters() : 0};
-    if (to.IsAllocated()) {
+    if (to.IsAllocated() && !skipRealloc) {
       // Top-level assignments to allocatable variables (*not* components)
       // may first deallocate existing content if there's about to be a
       // change in type or shape; see F'2018 10.2.1.3(3).
@@ -196,7 +198,7 @@ void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator) {
             comp.CreatePointerDescriptor(toCompDesc, to, terminator, toAt);
             comp.CreatePointerDescriptor(
                 fromCompDesc, from, terminator, fromAt);
-            Assign(toCompDesc, fromCompDesc, terminator);
+            Assign(toCompDesc, fromCompDesc, terminator, /*skipRealloc=*/false);
           }
         } else { // Component has intrinsic type; simply copy raw bytes
           std::size_t componentByteSize{comp.SizeInBytes(to)};
@@ -241,7 +243,7 @@ void Assign(Descriptor &to, const Descriptor &from, Terminator &terminator) {
               continue; // F'2018 10.2.1.3(13)(2)
             }
           }
-          Assign(*toDesc, *fromDesc, terminator);
+          Assign(*toDesc, *fromDesc, terminator, /*skipRealloc=*/false);
         }
         break;
       }

flang/runtime/assign.h

@@ -0,0 +1,30 @@
+//===-- runtime/assign.h-----------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+// Internal APIs for data assignment (both intrinsic assignment and TBP defined
+// generic ASSIGNMENT(=)).
+
+#ifndef FORTRAN_RUNTIME_ASSIGN_INTERNAL_H_
+#define FORTRAN_RUNTIME_ASSIGN_INTERNAL_H_
+
+namespace Fortran::runtime {
+class Descriptor;
+class Terminator;
+
+// Assigns one object to another via intrinsic assignment (F'2018 10.2.1.3) or
+// type-bound (only!) defined assignment (10.2.1.4), as appropriate.  Performs
+// finalization, scalar expansion, & allocatable (re)allocation as needed.
+// Does not perform intrinsic assignment implicit type conversion.  Both
+// descriptors must be initialized.  Recurses as needed to handle components.
+// Do not perform allocatable reallocation if \p skipRealloc is true, which is
+// used for allocate statement with source specifier.
+void Assign(
+    Descriptor &, const Descriptor &, Terminator &, bool skipRealloc = false);
+
+} // namespace Fortran::runtime
+#endif // FORTRAN_RUNTIME_ASSIGN_INTERNAL_H_

flang/runtime/pointer.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "flang/Runtime/pointer.h"
+#include "assign.h"
 #include "derived.h"
 #include "stat.h"
 #include "terminator.h"
@@ -132,6 +133,22 @@ int RTNAME(PointerAllocate)(Descriptor &pointer, bool hasStat,
   return stat;
 }
 
+int RTNAME(PointerAllocateSource)(Descriptor &pointer, const Descriptor &source,
+    bool hasStat, const Descriptor *errMsg, const char *sourceFile,
+    int sourceLine) {
+  if (pointer.Elements() == 0) {
+    return StatOk;
+  }
+  int stat{RTNAME(PointerAllocate)(
+      pointer, hasStat, errMsg, sourceFile, sourceLine)};
+  if (stat == StatOk) {
+    Terminator terminator{sourceFile, sourceLine};
+    // 9.7.1.2(7)
+    Assign(pointer, source, terminator, /*skipRealloc=*/true);
+  }
+  return stat;
+}
+
 int RTNAME(PointerDeallocate)(Descriptor &pointer, bool hasStat,
     const Descriptor *errMsg, const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};
@@ -187,7 +204,7 @@ bool RTNAME(PointerIsAssociatedWith)(
   return true;
 }
 
-// TODO: PointerCheckLengthParameter, PointerAllocateSource
+// TODO: PointerCheckLengthParameter
 
 } // extern "C"
 } // namespace Fortran::runtime