[flang] Lower allocation with MOLD

[clementval]

Summary

Lower allocate statement with MOLD= to calls to the Fortran
runtime. PointerApplyMold and AllocatableApplyMold are called
depending on the object to be allocated.

[jeanPerier]

Apart from the error condition case handling, looks good to me.

flang/lib/Lower/Allocatable.cpp

224

I do not get why this is returning a value, it seems AllocatableApplyMold returns void [1].

https://github.com/llvm/llvm-project/blob/01e8e50ce397dcb08e7dfadab93a53ae31174c67/flang/include/flang/Runtime/allocatable.h#L49

579

I think AllocatableApplyMold will set the status to unallocated regardless of the previous allocation. We should somehow preserve this info so that the program can crash/recover as needed if the allocatable is not deallocated to start with (9.7.1.3 point 1, first bullet point).

Maybe AllocatableApplyMold should actually be the one doing this check like with some other entry points.

[clementval]

Return from AllocatableApplyMold if the descriptor is
allocated so the error can be caught by AllocatableAllocate
to respect the first bullet of 9.7.1.3 point 1.

[jeanPerier]

clang-format issue in flang/lib/Lower/Allocatable.cpp somewhere, looks great otherwise.

[clementval]

clang-format

[PeteSteinfeld]

I saw one clang-format diff. Otherwise, all builds and tests correctly and looks good.

flang/lib/Lower/Allocatable.cpp

216–217

I get a clang-format diff here.

flang/include/flang/Runtime/allocatable.h

 // (If there's no STAT=, the error will be caught later anyway, but
 // this API allows the error to be caught before descriptor is modified.)
 // Return 0 on success (deallocated state), else the STAT= value.
 int RTNAME(AllocatableCheckAllocated)(Descriptor &,
     const Descriptor *errMsg = nullptr, const char *sourceFile = nullptr,
     int sourceLine = 0);
 
 // For MOLD= allocation; sets bounds, cobounds, and length type
-// parameters from another descriptor.  The destination descriptor must
+// parameters from another descriptor. The destination descriptor must
 // be initialized and deallocated.
 void RTNAME(AllocatableApplyMold)(Descriptor &, const Descriptor &mold);
 
 // Explicitly sets the bounds and length type parameters of an initialized
 // deallocated allocatable.
 void RTNAME(AllocatableSetBounds)(
     Descriptor &, int zeroBasedDim, SubscriptValue lower, SubscriptValue upper);
 

flang/lib/Lower/Allocatable.cpp

llvm::SmallVector<mlir::Value> args{
       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 runtime call to apply mold to the descriptor.
+static void genRuntimeAllocateApplyMold(fir::FirOpBuilder &builder,
+                                        mlir::Location loc,
+                                        const fir::MutableBoxValue &box,
+                                        fir::ExtendedValue mold) {
+  mlir::func::FuncOp callee =
+      box.isPointer()
+          ? fir::runtime::getRuntimeFunc<mkRTKey(PointerApplyMold)>(loc,
+                                                                    builder)
+          : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableApplyMold)>(
+                loc, builder);
+  llvm::SmallVector<mlir::Value> args{box.getAddr(), fir::getBase(mold)};
+  llvm::SmallVector<mlir::Value> operands;
+  for (auto [fst, snd] : llvm::zip(args, callee.getFunctionType().getInputs()))
+    operands.emplace_back(builder.createConvert(loc, snd, fst));
+  builder.create<fir::CallOp>(loc, callee, operands);
+}
+
 /// Generate a runtime call to deallocate memory.
 static mlir::Value genRuntimeDeallocate(fir::FirOpBuilder &builder,
                                         mlir::Location loc,
                                         const fir::MutableBoxValue &box,
                                         ErrorManager &errorManager,
                                         mlir::Value declaredTypeDesc = {}) {
   // Ensure fir.box is up-to-date before passing it to deallocate runtime.
   mlir::Value boxAddress = fir::factory::getMutableIRBox(builder, loc, box);
public:
   void lower() {
     visitAllocateOptions();
     lowerAllocateLengthParameters();
     errorManager.init(converter, loc, statExpr, errMsgExpr);
     Fortran::lower::StatementContext stmtCtx;
     if (sourceExpr)
       sourceExv = converter.genExprBox(loc, *sourceExpr, stmtCtx);
     if (moldExpr)
-      TODO(loc, "lower MOLD expr in allocate");
+      moldExv = converter.genExprBox(loc, *moldExpr, stmtCtx);
     mlir::OpBuilder::InsertPoint insertPt = builder.saveInsertionPoint();
     for (const auto &allocation :
          std::get<std::list<Fortran::parser::Allocation>>(stmt.t))
       lowerAllocation(unwrapAllocation(allocation));
     builder.restoreInsertionPoint(insertPt);
   }
 
 private:
if (lenParams.empty() && box.isCharacter() &&
       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");
+
+  void genMoldAllocation(const Allocation &alloc,
+                         const fir::MutableBoxValue &box) {
+    genRuntimeAllocateApplyMold(builder, loc, box, moldExv);
+    errorManager.genStatCheck(builder, loc);
+    mlir::Value stat = genRuntimeAllocate(builder, loc, box, errorManager);
+    fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
+    errorManager.assignStat(builder, loc, stat);
   }
 
   /// Generate call to the AllocatableInitDerived to set up the type descriptor
   /// and other part of the descriptor for derived type.
   void genSetType(const Allocation &alloc, const fir::MutableBoxValue &box,
                   mlir::Location loc) {
     const Fortran::semantics::DeclTypeSpec *typeSpec =
         getIfAllocateStmtTypeSpec();
private:
   const Fortran::lower::SomeExpr *statExpr{nullptr};
   const Fortran::lower::SomeExpr *errMsgExpr{nullptr};
   // If the allocate has a type spec, lenParams contains the
   // 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;
+  fir::ExtendedValue moldExv;
 
   mlir::Location loc;
 };
 } // namespace
 
 void Fortran::lower::genAllocateStmt(
     Fortran::lower::AbstractConverter &converter,
     const Fortran::parser::AllocateStmt &stmt, mlir::Location loc) {

flang/runtime/allocatable.cpp

void RTNAME(AllocatableSetDerivedLength)(
     Descriptor &descriptor, int which, SubscriptValue x) {
   DescriptorAddendum *addendum{descriptor.Addendum()};
   INTERNAL_CHECK(addendum != nullptr);
   addendum->SetLenParameterValue(which, x);
 }
 
 void RTNAME(AllocatableApplyMold)(
     Descriptor &descriptor, const Descriptor &mold) {
+  if (descriptor.IsAllocated()) {
+    // 9.7.1.3 Return so the error can be emitted by AllocatableAllocate.
+    return;
+  }
   descriptor = mold;
   descriptor.set_base_addr(nullptr);
   descriptor.raw().attribute = CFI_attribute_allocatable;
 }
 
 int RTNAME(AllocatableAllocate)(Descriptor &descriptor, bool hasStat,
     const Descriptor *errMsg, const char *sourceFile, int sourceLine) {
   Terminator terminator{sourceFile, sourceLine};

flang/test/Lower/allocatable-polymorphic.f90

subroutine test_deallocate()
     deallocate(p)
   end subroutine
 
   subroutine test_type_with_polymorphic_pointer_component()
     type(with_alloc), pointer :: a
     allocate(a)
     allocate(a%element)
   end subroutine
-end module
 
 ! CHECK-LABEL: func.func @_QMpolyPtest_type_with_polymorphic_pointer_component()
 ! CHECK: %[[TYPE_PTR:.*]] = fir.alloca !fir.ptr<!fir.type<_QMpolyTwith_alloc{element:!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>}>> {uniq_name = "_QMpolyFtest_type_with_polymorphic_pointer_componentEa.addr"}
 ! CHECK: %[[TYPE_PTR_LOAD:.*]] = fir.load %[[TYPE_PTR]] : !fir.ref<!fir.ptr<!fir.type<_QMpolyTwith_alloc{element:!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>}>>>
 ! CHECK: %[[ELEMENT:.*]] = fir.field_index element, !fir.type<_QMpolyTwith_alloc{element:!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>}>
 ! CHECK: %[[ELEMENT_DESC:.*]] = fir.coordinate_of %[[TYPE_PTR_LOAD]], %[[ELEMENT]] : (!fir.ptr<!fir.type<_QMpolyTwith_alloc{element:!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>}>>, !fir.field) -> !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>
 ! CHECK: %[[ZERO:.*]] = fir.zero_bits !fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>
 ! CHECK: %[[ZERO_DESC:.*]] = fir.embox %[[ZERO]] : (!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>) -> !fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>
 ! CHECK: fir.store %[[ZERO_DESC]] to %[[ELEMENT_DESC]] : !fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>
 ! CHECK: %[[TYPE_DESC_P1:.*]] = fir.address_of(@_QMpolyE.dt.p1) : !fir.ref<!fir.type<{{.*}}>>
 ! CHECK: %[[ELEMENT_DESC_CAST:.*]] = fir.convert %[[ELEMENT_DESC]] : (!fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %[[TYPE_DESC_P1_CAST:.*]] = fir.convert %[[TYPE_DESC_P1]] : (!fir.ref<!fir.type<{{.*}}>>) -> !fir.ref<none> 
 ! CHECK: %[[RANK:.*]] = arith.constant 0 : i32
 ! CHECK: %[[CORANK:.*]] = arith.constant 0 : i32
 ! CHECK: %{{.*}} = fir.call @_FortranAPointerNullifyDerived(%[[ELEMENT_DESC_CAST]], %[[TYPE_DESC_P1_CAST]], %[[RANK]], %[[CORANK]]) {{.*}}: (!fir.ref<!fir.box<none>>, !fir.ref<none>, i32, i32) -> none
 ! CHECK: %[[ELEMENT_DESC_CAST:.*]] = fir.convert %[[ELEMENT_DESC]] : (!fir.ref<!fir.class<!fir.ptr<!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %{{.*}} = fir.call @_FortranAPointerAllocate(%[[ELEMENT_DESC_CAST]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}}: (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
 
+  subroutine test_allocate_with_mold()
+    type(p2) :: x(10)
+    class(p1), pointer :: p(:)
+    integer(4) :: i(20)
+    class(*), pointer :: up(:)
+
+    allocate(p, mold=x)
+    allocate(up, mold=i)
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMpolyPtest_allocate_with_mold() {
+! CHECK: %[[I:.*]] = fir.alloca !fir.array<20xi32> {bindc_name = "i", uniq_name = "_QMpolyFtest_allocate_with_moldEi"}
+! CHECK: %[[P:.*]] = fir.alloca !fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>> {bindc_name = "p", uniq_name = "_QMpolyFtest_allocate_with_moldEp"}
+! CHECK: %[[UP:.*]] = fir.alloca !fir.class<!fir.ptr<!fir.array<?xnone>>> {bindc_name = "up", uniq_name = "_QMpolyFtest_allocate_with_moldEup"}
+! CHECK: %[[X:.*]] = fir.alloca !fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>> {bindc_name = "x", uniq_name = "_QMpolyFtest_allocate_with_moldEx"}
+! CHECK: %[[EMBOX_X:.*]] = fir.embox %[[X]](%{{.*}}) : (!fir.ref<!fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>>>, !fir.shape<1>) -> !fir.box<!fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>>>
+! CHECK: %[[P_BOX_NONE:.*]] = fir.convert %[[P]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %[[X_BOX_NONE:.*]] = fir.convert %[[EMBOX_X]] : (!fir.box<!fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>>>) -> !fir.box<none>
+! CHECK: %{{.*}} = fir.call @_FortranAPointerApplyMold(%[[P_BOX_NONE]], %[[X_BOX_NONE]]) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>) -> none
+! CHECK: %[[P_BOX_NONE:.*]] = fir.convert %[[P]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?x!fir.type<_QMpolyTp1{a:i32,b:i32}>>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAPointerAllocate(%[[P_BOX_NONE]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+! CHECK: %[[EMBOX_I:.*]] = fir.embox %[[I]](%{{.*}}) : (!fir.ref<!fir.array<20xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<20xi32>>
+! CHECK: %[[UP_BOX_NONE:.*]] = fir.convert %[[UP]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?xnone>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %[[I_BOX_NONE:.*]] = fir.convert %[[EMBOX_I]] : (!fir.box<!fir.array<20xi32>>) -> !fir.box<none>
+! CHECK: %{{.*}} = fir.call @_FortranAPointerApplyMold(%[[UP_BOX_NONE]], %[[I_BOX_NONE]]) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>) -> none
+! CHECK: %[[UP_BOX_NONE:.*]] = fir.convert %[[UP]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?xnone>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAPointerAllocate(%[[UP_BOX_NONE]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+
+end module
+
+
 program test_alloc
   use poly
 
   call test_allocatable()
   call test_pointer()
 end
 
 ! Check code generation of allocate runtime calls for polymorphic entities. This

flang/test/Lower/allocatable-runtime.f90

subroutine char_explicit_dyn(n, l1, l2)
   ! CHECK-NOT: AllocatableInitCharacter
   ! CHECK: AllocatableAllocate
   ! CHECK-NOT: AllocatableInitCharacter
   ! CHECK: AllocatableAllocate
   deallocate(scalar, zarray)
   ! CHECK: AllocatableDeallocate
   ! CHECK: AllocatableDeallocate
 end subroutine
+
+subroutine mold_allocation()
+  integer :: m(10)
+  integer, allocatable :: a(:)
+
+  allocate(a, mold=m)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPmold_allocation() {
+! CHECK: %[[A:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>> {bindc_name = "a", uniq_name = "_QFmold_allocationEa"}
+! CHECK: %[[M:.*]] = fir.alloca !fir.array<10xi32> {bindc_name = "m", uniq_name = "_QFmold_allocationEm"}
+! CHECK: %[[EMBOX_M:.*]] = fir.embox %[[M]](%{{.*}}) : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> !fir.box<!fir.array<10xi32>>
+! CHECK: %[[A_BOX_NONE:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %[[M_BOX_NONE:.*]] = fir.convert %[[EMBOX_M]] : (!fir.box<!fir.array<10xi32>>) -> !fir.box<none>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableApplyMold(%[[A_BOX_NONE]], %[[M_BOX_NONE]]) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>) -> none
+! CHECK: %[[A_BOX_NONE:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableAllocate(%[[A_BOX_NONE]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32