[flang] Support allocate with source for polymorphic entities

[clementval]

Summary

Apply the source type spec to the descriptor for
polyrmophic entities.

[clementval]

Use alloc object rank instead of source rank

[jeanPerier]

Questions about some edge cases, the rest looks good.

flang/lib/Lower/Allocatable.cpp

578

Is it possible for the sourceExpr->GetType() to be null ?

Is it also possible that the source be an unlimited polymorphic if the allocatable object also is one ?

subroutine foo(any, any_source)
 class(*), allocatable :: any
 class(*), pointer :: any_source
 allocate(any, source = any_source)
end subroutine

(If so, this will probably require some custom runtime entry point, and it's fine to leave a TODO at first).

[clementval]

Simplify condition + add test for unlimited polymorphic allocation from
unlimited polyrmophic source.

[clementval]

flang/lib/Lower/Allocatable.cpp

578

I believe sourceExpr->GetType() is always defined.

It is possible to have an unlimited polymorphic source to allocate an unlimited polymorphic source. This falls in the TypeCategory::Derived category as unlimited polymorphic descriptor are defined as DerivedType in the frontend.

[jeanPerier]

flang/lib/Lower/Allocatable.cpp

578

But what if the unlimited polymorphic dynamic type is an intrinsic type and not a derived type (say any_source is pointing to an integer). Is it correct to call genInitDerived in that case ?

578

I believe sourceExpr->GetType() is always defined.

Ok, makes sense. I was not sure. However, I remember that in some rare semantics bugs, it happened that GetType() returned null (because something failed in expression resolution and was not correctly reported), so I would tend add asserts when using that (or using the DEREF() front end macro).

[clementval]

flang/lib/Lower/Allocatable.cpp

578

Yeah I was thinking about this after the update. This would need to go through a custom runtime entry point. I'll make a separate patch for this.

578

An assert makes sense.

[clementval]

• Add a TODO for unlimited polymoprhic source and allocate object
• Add an assert for GetType()

[PeteSteinfeld]

Other than some clang-format diffs, all builds and tests correctly and looks good.

flang/lib/Lower/Allocatable.cpp

570–575

I get clang-format diffs here.

627

For some reason this comment generates a clang-format diff.

634–638

More clang-format diffs here.

646–650

More clang-format diffs here.

676–680

More clang-format diffs.

[clementval]

flang/lib/Lower/Allocatable.cpp

570–575

I guess it was before the last update. Pre-commit check is clean right now.

[clementval]

Rebase after MOLD patch

flang/lib/Lower/Allocatable.cpp

private:
 
   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));
+    if (alloc.type.IsPolymorphic()) {
+      assert(sourceExpr->GetType() && "null type not expected");
+      if (alloc.type.IsUnlimitedPolymorphic() &&
+          sourceExpr->GetType()->IsUnlimitedPolymorphic())
+        TODO(loc, "allocate unlimited polymorphic entity from unlimited "
+                  "polymorphic source");
+
+      if (sourceExpr->GetType()->category() == TypeCategory::Derived) {
+        mlir::Type tdescType =
+            fir::TypeDescType::get(mlir::NoneType::get(builder.getContext()));
+        mlir::Value typeDescAddr = builder.create<fir::BoxTypeDescOp>(
+            loc, tdescType, fir::getBase(sourceExv));
+        genInitDerived(box, typeDescAddr, alloc.getSymbol().Rank());
+      } else {
+        genInitIntrinsic(box, sourceExpr->GetType()->category(),
+                         sourceExpr->GetType()->kind(),
+                         alloc.getSymbol().Rank());
+      }
+    }
     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 &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 PointerNullifyDerived or AllocatableInitDerived
+  /// to set the dynamic type information.
+  void genInitDerived(const fir::MutableBoxValue &box, mlir::Value typeDescAddr,
+                      int rank, int corank = 0) {
+    mlir::func::FuncOp callee =
+        box.isPointer()
+            ? fir::runtime::getRuntimeFunc<mkRTKey(PointerNullifyDerived)>(
+                  loc, builder)
+            : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableInitDerived)>(
+                  loc, builder);
+
+    llvm::ArrayRef<mlir::Type> inputTypes =
+        callee.getFunctionType().getInputs();
+    llvm::SmallVector<mlir::Value> args;
+    args.push_back(builder.createConvert(loc, inputTypes[0], box.getAddr()));
+    args.push_back(builder.createConvert(loc, inputTypes[1], typeDescAddr));
+    mlir::Value rankValue =
+        builder.createIntegerConstant(loc, inputTypes[2], rank);
+    mlir::Value corankValue =
+        builder.createIntegerConstant(loc, inputTypes[3], corank);
+    args.push_back(rankValue);
+    args.push_back(corankValue);
+    builder.create<fir::CallOp>(loc, callee, args);
+  }
+
+  /// Generate call to PointerNullifyIntrinsic or AllocatableInitIntrinsic to
+  /// set the dynamic type information for a polymorphic entity from an
+  /// intrinsic type spec.
+  void genInitIntrinsic(const fir::MutableBoxValue &box,
+                        const TypeCategory category, int64_t kind, int rank,
+                        int corank = 0) {
+    mlir::func::FuncOp callee =
+        box.isPointer()
+            ? fir::runtime::getRuntimeFunc<mkRTKey(PointerNullifyIntrinsic)>(
+                  loc, builder)
+            : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableInitIntrinsic)>(
+                  loc, builder);
+
+    llvm::ArrayRef<mlir::Type> inputTypes =
+        callee.getFunctionType().getInputs();
+    llvm::SmallVector<mlir::Value> args;
+    args.push_back(builder.createConvert(loc, inputTypes[0], box.getAddr()));
+    mlir::Value categoryValue = builder.createIntegerConstant(
+        loc, inputTypes[1], static_cast<int32_t>(category));
+    mlir::Value kindValue =
+        builder.createIntegerConstant(loc, inputTypes[2], kind);
+    mlir::Value rankValue =
+        builder.createIntegerConstant(loc, inputTypes[3], rank);
+    mlir::Value corankValue =
+        builder.createIntegerConstant(loc, inputTypes[4], corank);
+    args.push_back(categoryValue);
+    args.push_back(kindValue);
+    args.push_back(rankValue);
+    args.push_back(corankValue);
+    builder.create<fir::CallOp>(loc, callee, args);
+  }
+
   /// 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();
 
     // No type spec provided in allocate statement so the declared type spec is
     // used.
     if (!typeSpec)
       typeSpec = &alloc.type;
     assert(typeSpec && "type spec missing for polymorphic allocation");
 
     // Set up the descriptor for allocation for intrinsic type spec on
     // unlimited polymorphic entity.
     if (typeSpec->AsIntrinsic() &&
         fir::isUnlimitedPolymorphicType(fir::getBase(box).getType())) {
-      mlir::func::FuncOp callee =
-          box.isPointer()
-              ? fir::runtime::getRuntimeFunc<mkRTKey(PointerNullifyIntrinsic)>(
-                    loc, builder)
-              : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableInitIntrinsic)>(
-                    loc, builder);
-
-      llvm::ArrayRef<mlir::Type> inputTypes =
-          callee.getFunctionType().getInputs();
-      llvm::SmallVector<mlir::Value> args;
-      args.push_back(builder.createConvert(loc, inputTypes[0], box.getAddr()));
-      mlir::Value category = builder.createIntegerConstant(
-          loc, inputTypes[1],
-          static_cast<int32_t>(typeSpec->AsIntrinsic()->category()));
-      mlir::Value kind = builder.createIntegerConstant(
-          loc, inputTypes[2],
-          Fortran::evaluate::ToInt64(typeSpec->AsIntrinsic()->kind()).value());
-      mlir::Value rank = builder.createIntegerConstant(
-          loc, inputTypes[3], alloc.getSymbol().Rank());
-      mlir::Value corank = builder.createIntegerConstant(loc, inputTypes[4], 0);
-      args.push_back(category);
-      args.push_back(kind);
-      args.push_back(rank);
-      args.push_back(corank);
-      builder.create<fir::CallOp>(loc, callee, args);
+      genInitIntrinsic(
+          box, typeSpec->AsIntrinsic()->category(),
+          Fortran::evaluate::ToInt64(typeSpec->AsIntrinsic()->kind()).value(),
+          alloc.getSymbol().Rank());
       return;
     }
 
     // Do not generate calls for non derived-type type spec.
     if (!typeSpec->AsDerived())
       return;
 
     auto typeDescAddr = Fortran::lower::getTypeDescAddr(
         builder, loc, typeSpec->derivedTypeSpec());
-    mlir::func::FuncOp callee =
-        box.isPointer()
-            ? fir::runtime::getRuntimeFunc<mkRTKey(PointerNullifyDerived)>(
-                  loc, builder)
-            : fir::runtime::getRuntimeFunc<mkRTKey(AllocatableInitDerived)>(
-                  loc, builder);
-
-    llvm::ArrayRef<mlir::Type> inputTypes =
-        callee.getFunctionType().getInputs();
-    llvm::SmallVector<mlir::Value> args;
-    args.push_back(builder.createConvert(loc, inputTypes[0], box.getAddr()));
-    args.push_back(builder.createConvert(loc, inputTypes[1], typeDescAddr));
-    mlir::Value rank = builder.createIntegerConstant(loc, inputTypes[2],
-                                                     alloc.getSymbol().Rank());
-    mlir::Value corank = builder.createIntegerConstant(loc, inputTypes[3], 0);
-    args.push_back(rank);
-    args.push_back(corank);
-    builder.create<fir::CallOp>(loc, callee, args);
+    genInitDerived(box, typeDescAddr, alloc.getSymbol().Rank());
   }
 
   /// Returns a pointer to the DeclTypeSpec if a type-spec is provided in the
   /// allocate statement. Returns a null pointer otherwise.
   const Fortran::semantics::DeclTypeSpec *getIfAllocateStmtTypeSpec() const {
     if (const auto &typeSpec =
             std::get<std::optional<Fortran::parser::TypeSpec>>(stmt.t))
       return typeSpec->declTypeSpec;

flang/test/Lower/allocatable-polymorphic.f90

 ! 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
 
+  subroutine test_allocate_with_source()
+    type(p2) :: x(10)
+    class(p1), pointer :: p(:)
+    integer(4) :: i(20)
+    class(*), pointer :: up(:)
+
+    allocate(p, source=x)
+    allocate(up, source=i)
+  end subroutine
+
+! CHECK-LABEL: func.func @_QMpolyPtest_allocate_with_source() {
+! CHECK: %[[I:.*]] = fir.alloca !fir.array<20xi32> {bindc_name = "i", uniq_name = "_QMpolyFtest_allocate_with_sourceEi"}
+! 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_sourceEp"}
+! CHECK: %[[UP:.*]] = fir.alloca !fir.class<!fir.ptr<!fir.array<?xnone>>> {bindc_name = "up", uniq_name = "_QMpolyFtest_allocate_with_sourceEup"}
+! CHECK: %[[X:.*]] = fir.alloca !fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>> {bindc_name = "x", uniq_name = "_QMpolyFtest_allocate_with_sourceEx"}
+
+! 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: %[[TYPE_DESC_X:.*]] = fir.box_tdesc %[[EMBOX_X]] : (!fir.box<!fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>>>) -> !fir.tdesc<none>
+! CHECK: %[[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: %[[TYPE_DESC_NONE:.*]] = fir.convert %[[TYPE_DESC_X]] : (!fir.tdesc<none>) -> !fir.ref<none>
+! CHECK: %[[RANK:.*]] = arith.constant 1 : i32
+! CHECK: %[[CORANK:.*]] = arith.constant 0 : i32
+! CHECK: %{{.*}} = fir.call @_FortranAPointerNullifyDerived(%[[BOX_NONE]], %[[TYPE_DESC_NONE]], %[[RANK]], %[[CORANK]]) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.ref<none>, i32, i32) -> none
+! CHECK: %{{.*}} = fir.call @_FortranAPointerSetBounds
+! CHECK: %[[BOX_NONE_P:.*]] = 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: %[[BOX_NONE_X:.*]] = fir.convert %[[EMBOX_X]] : (!fir.box<!fir.array<10x!fir.type<_QMpolyTp2{a:i32,b:i32,c:i32}>>>) -> !fir.box<none>
+! CHECK: %{{.*}} = fir.call @_FortranAPointerAllocateSource(%[[BOX_NONE_P]], %[[BOX_NONE_X]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+
+! CHECK: %[[UP_BOX_NONE:.*]] = fir.convert %[[UP]] : (!fir.ref<!fir.class<!fir.ptr<!fir.array<?xnone>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %[[CAT:.*]] = arith.constant 0 : i32
+! CHECK: %[[KIND:.*]] = arith.constant 4 : i32
+! CHECK: %[[RANK:.*]] = arith.constant 1 : i32
+! CHECK: %[[CORANK:.*]] = arith.constant 0 : i32
+! CHECK: %{{.*}} = fir.call @_FortranAPointerNullifyIntrinsic(%[[UP_BOX_NONE]], %[[CAT]], %[[KIND]], %[[RANK]], %[[CORANK]]) {{.*}} : (!fir.ref<!fir.box<none>>, i32, i32, i32, i32) -> none
+
 end module
 
 
 program test_alloc
   use poly
 
   call test_allocatable()
   call test_pointer()