[flang][hlfir] do not propagate polymorphic temporary as allocatables (#142609)

Polymorphic temporary are currently propagated as
fir.ref<fir.class<fir.heap<>>> because their allocation may be delayed
to the hlfir.assign copy (using realloc).

This patch moves away from this and directly allocate the temp and
propagate it as a fir.class.

The polymorphic temporaries creating is also simplified by avoiding the
need to call the runtime to setup the descriptor altogether (the runtime
is still call for the allocation currently because alloca/allocmem do
not support polymorphism).

Author: jeanPerier

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

@@ -271,14 +271,14 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   /// Sample genDeclare callback for createArrayTemp() below.
   /// It creates fir.declare operation using the given operands.
   /// \p memref is the base of the allocated temporary,
-  /// which may be !fir.ref<!fir.array<>> or !fir.ref<!fir.box/class<>>.
+  /// which may be !fir.ref<!fir.array<>> or !fir.box/class<>.
   static mlir::Value genTempDeclareOp(fir::FirOpBuilder &builder,
                                       mlir::Location loc, mlir::Value memref,
                                       llvm::StringRef name, mlir::Value shape,
                                       llvm::ArrayRef<mlir::Value> typeParams,
                                       fir::FortranVariableFlagsAttr attrs);
 
-  /// Create a temporary array with the given \p arrayType,
+  /// Create a temporary with the given \p baseType,
   /// \p shape, \p extents and \p typeParams. An optional
   /// \p polymorphicMold specifies the entity which dynamic type
   /// has to be used for the allocation.
@@ -291,16 +291,26 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   /// If \p useStack is true, the function will try to do the allocation
   /// in stack memory (which is not always possible currently).
   /// The first return value is the base of the temporary object,
-  /// which may be !fir.ref<!fir.array<>> or !fir.ref<!fir.box/class<>>.
+  /// which may be !fir.ref<!fir.array<>> or !fir.box/class<>.
   /// The second return value is true, if the actual allocation
   /// was done in heap memory.
+  std::pair<mlir::Value, bool> createAndDeclareTemp(
+      mlir::Location loc, mlir::Type baseType, mlir::Value shape,
+      llvm::ArrayRef<mlir::Value> extents,
+      llvm::ArrayRef<mlir::Value> typeParams,
+      const std::function<decltype(genTempDeclareOp)> &genDeclare,
+      mlir::Value polymorphicMold, bool useStack, llvm::StringRef tmpName);
+  /// Create and declare an array temporary.
   std::pair<mlir::Value, bool>
   createArrayTemp(mlir::Location loc, fir::SequenceType arrayType,
                   mlir::Value shape, llvm::ArrayRef<mlir::Value> extents,
                   llvm::ArrayRef<mlir::Value> typeParams,
                   const std::function<decltype(genTempDeclareOp)> &genDeclare,
                   mlir::Value polymorphicMold, bool useStack = false,
-                  llvm::StringRef tmpName = ".tmp.array");
+                  llvm::StringRef tmpName = ".tmp.array") {
+    return createAndDeclareTemp(loc, arrayType, shape, extents, typeParams,
+                                genDeclare, polymorphicMold, useStack, tmpName);
+  }
 
   /// Create an LLVM stack save intrinsic op. Returns the saved stack pointer.
   /// The stack address space is fetched from the data layout of the current

flang/include/flang/Optimizer/Builder/MutableBox.h

@@ -181,6 +181,18 @@ mlir::Value genIsNotAllocatedOrAssociatedTest(fir::FirOpBuilder &builder,
 mlir::Value genNullBoxStorage(fir::FirOpBuilder &builder, mlir::Location loc,
                               mlir::Type boxTy);
 
+/// Generate an unallocated box of the given \p boxTy with the
+/// bounds, type parameters, and dynamic type set according to the
+/// parameters.
+/// \p shape may be null for scalars, and \p polymorphicMold may be null for
+/// statically typed entities. This box can then be directly passed to the
+/// runtime for allocation.
+mlir::Value getAndEstablishBoxStorage(fir::FirOpBuilder &builder,
+                                      mlir::Location loc,
+                                      fir::BaseBoxType boxTy, mlir::Value shape,
+                                      llvm::ArrayRef<mlir::Value> typeParams,
+                                      mlir::Value polymorphicMold);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_MUTABLEBOX_H

flang/include/flang/Optimizer/Dialect/FIRType.h

@@ -47,6 +47,9 @@ class BaseBoxType : public mlir::Type {
   /// Returns the element type of this box type.
   mlir::Type getEleTy() const;
 
+  /// Get the raw address type of the memory described by the box.
+  mlir::Type getBaseAddressType() const;
+
   /// Unwrap element type from fir.heap, fir.ptr and fir.array.
   mlir::Type unwrapInnerType() const;
 
@@ -56,6 +59,9 @@ class BaseBoxType : public mlir::Type {
   /// Is this a box for a pointer?
   bool isPointer() const;
 
+  /// Is this a box describing volatile memory?
+  bool isVolatile() const;
+
   /// Return the same type, except for the shape, that is taken the shape
   /// of shapeMold.
   BaseBoxType getBoxTypeWithNewShape(mlir::Type shapeMold) const;

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -366,55 +366,29 @@ mlir::Value fir::FirOpBuilder::createHeapTemporary(
                                  name, dynamicLength, dynamicShape, attrs);
 }
 
-std::pair<mlir::Value, bool> fir::FirOpBuilder::createArrayTemp(
-    mlir::Location loc, fir::SequenceType arrayType, mlir::Value shape,
+std::pair<mlir::Value, bool> fir::FirOpBuilder::createAndDeclareTemp(
+    mlir::Location loc, mlir::Type baseType, mlir::Value shape,
     llvm::ArrayRef<mlir::Value> extents, llvm::ArrayRef<mlir::Value> typeParams,
     const std::function<decltype(FirOpBuilder::genTempDeclareOp)> &genDeclare,
     mlir::Value polymorphicMold, bool useStack, llvm::StringRef tmpName) {
   if (polymorphicMold) {
     // Create *allocated* polymorphic temporary using the dynamic type
-    // of the mold and the provided shape/extents. The created temporary
-    // array will be written element per element, that is why it has to be
-    // allocated.
-    mlir::Type boxHeapType = fir::HeapType::get(arrayType);
-    mlir::Value alloc = fir::factory::genNullBoxStorage(
-        *this, loc, fir::ClassType::get(boxHeapType));
-    fir::FortranVariableFlagsAttr declAttrs =
-        fir::FortranVariableFlagsAttr::get(
-            getContext(), fir::FortranVariableFlagsEnum::allocatable);
-
-    mlir::Value base = genDeclare(*this, loc, alloc, tmpName,
-                                  /*shape=*/nullptr, typeParams, declAttrs);
-
-    int rank = extents.size();
-    fir::runtime::genAllocatableApplyMold(*this, loc, alloc, polymorphicMold,
-                                          rank);
-    if (!extents.empty()) {
-      mlir::Type idxTy = getIndexType();
-      mlir::Value one = createIntegerConstant(loc, idxTy, 1);
-      unsigned dim = 0;
-      for (mlir::Value extent : extents) {
-        mlir::Value dimIndex = createIntegerConstant(loc, idxTy, dim++);
-        fir::runtime::genAllocatableSetBounds(*this, loc, alloc, dimIndex, one,
-                                              extent);
-      }
-    }
-    if (!typeParams.empty()) {
-      // We should call AllocatableSetDerivedLength() here.
-      // TODO: does the mold provide the length parameters or
-      // the operation itself or should they be in sync?
-      TODO(loc, "polymorphic type with length parameters");
-    }
-    fir::runtime::genAllocatableAllocate(*this, loc, alloc);
-
+    // of the mold and the provided shape/extents.
+    auto boxType = fir::ClassType::get(fir::HeapType::get(baseType));
+    mlir::Value boxAddress = fir::factory::getAndEstablishBoxStorage(
+        *this, loc, boxType, shape, typeParams, polymorphicMold);
+    fir::runtime::genAllocatableAllocate(*this, loc, boxAddress);
+    mlir::Value box = create<fir::LoadOp>(loc, boxAddress);
+    mlir::Value base =
+        genDeclare(*this, loc, box, tmpName, /*shape=*/mlir::Value{},
+                   typeParams, fir::FortranVariableFlagsAttr{});
     return {base, /*isHeapAllocation=*/true};
   }
   mlir::Value allocmem;
   if (useStack)
-    allocmem = createTemporary(loc, arrayType, tmpName, extents, typeParams);
+    allocmem = createTemporary(loc, baseType, tmpName, extents, typeParams);
   else
-    allocmem =
-        createHeapTemporary(loc, arrayType, tmpName, extents, typeParams);
+    allocmem = createHeapTemporary(loc, baseType, tmpName, extents, typeParams);
   mlir::Value base = genDeclare(*this, loc, allocmem, tmpName, shape,
                                 typeParams, fir::FortranVariableFlagsAttr{});
   return {base, !useStack};

flang/lib/Optimizer/Builder/HLFIRTools.cpp

@@ -1331,63 +1331,36 @@ bool hlfir::elementalOpMustProduceTemp(hlfir::ElementalOp elemental) {
 std::pair<hlfir::Entity, mlir::Value>
 hlfir::createTempFromMold(mlir::Location loc, fir::FirOpBuilder &builder,
                           hlfir::Entity mold) {
+  assert(!mold.isAssumedRank() &&
+         "cannot create temporary from assumed-rank mold");
   llvm::SmallVector<mlir::Value> lenParams;
   hlfir::genLengthParameters(loc, builder, mold, lenParams);
   llvm::StringRef tmpName{".tmp"};
-  mlir::Value alloc;
-  mlir::Value isHeapAlloc;
-  mlir::Value shape{};
-  fir::FortranVariableFlagsAttr declAttrs;
 
-  if (mold.isPolymorphic()) {
-    // Create unallocated polymorphic temporary using the dynamic type
-    // of the mold. The static type of the temporary matches
-    // the static type of the mold, but then the dynamic type
-    // of the mold is applied to the temporary's descriptor.
-
-    if (mold.isArray())
-      hlfir::genShape(loc, builder, mold);
-
-    // Create polymorphic allocatable box on the stack.
-    mlir::Type boxHeapType = fir::HeapType::get(fir::unwrapRefType(
-        mlir::cast<fir::BaseBoxType>(mold.getType()).getEleTy()));
-    // The box must be initialized, because AllocatableApplyMold
-    // may read its contents (e.g. for checking whether it is allocated).
-    alloc = fir::factory::genNullBoxStorage(builder, loc,
-                                            fir::ClassType::get(boxHeapType));
-    // The temporary is unallocated even after AllocatableApplyMold below.
-    // If the temporary is used as assignment LHS it will be automatically
-    // allocated on the heap, as long as we use Assign family
-    // runtime functions. So set MustFree to true.
-    isHeapAlloc = builder.createBool(loc, true);
-    declAttrs = fir::FortranVariableFlagsAttr::get(
-        builder.getContext(), fir::FortranVariableFlagsEnum::allocatable);
-  } else if (mold.isArray()) {
-    mlir::Type sequenceType =
-        hlfir::getFortranElementOrSequenceType(mold.getType());
+  mlir::Value shape{};
+  llvm::SmallVector<mlir::Value> extents;
+  if (mold.isArray()) {
     shape = hlfir::genShape(loc, builder, mold);
-    auto extents = hlfir::getIndexExtents(loc, builder, shape);
-    alloc = builder.createHeapTemporary(loc, sequenceType, tmpName, extents,
-                                        lenParams);
-    isHeapAlloc = builder.createBool(loc, true);
-  } else {
-    alloc = builder.createTemporary(loc, mold.getFortranElementType(), tmpName,
-                                    /*shape=*/std::nullopt, lenParams);
-    isHeapAlloc = builder.createBool(loc, false);
-  }
-  auto declareOp =
-      builder.create<hlfir::DeclareOp>(loc, alloc, tmpName, shape, lenParams,
-                                       /*dummy_scope=*/nullptr, declAttrs);
-  if (mold.isPolymorphic()) {
-    int rank = mold.getRank();
-    // TODO: should probably read rank from the mold.
-    if (rank < 0)
-      TODO(loc, "create temporary for assumed rank polymorphic");
-    fir::runtime::genAllocatableApplyMold(builder, loc, alloc,
-                                          mold.getFirBase(), rank);
+    extents = hlfir::getExplicitExtentsFromShape(shape, builder);
   }
 
-  return {hlfir::Entity{declareOp.getBase()}, isHeapAlloc};
+  bool useStack = !mold.isArray() && !mold.isPolymorphic();
+  auto genTempDeclareOp =
+      [](fir::FirOpBuilder &builder, mlir::Location loc, mlir::Value memref,
+         llvm::StringRef name, mlir::Value shape,
+         llvm::ArrayRef<mlir::Value> typeParams,
+         fir::FortranVariableFlagsAttr attrs) -> mlir::Value {
+    auto declareOp =
+        builder.create<hlfir::DeclareOp>(loc, memref, name, shape, typeParams,
+                                         /*dummy_scope=*/nullptr, attrs);
+    return declareOp.getBase();
+  };
+
+  auto [base, isHeapAlloc] = builder.createAndDeclareTemp(
+      loc, mold.getElementOrSequenceType(), shape, extents, lenParams,
+      genTempDeclareOp, mold.isPolymorphic() ? mold.getBase() : nullptr,
+      useStack, tmpName);
+  return {hlfir::Entity{base}, builder.createBool(loc, isHeapAlloc)};
 }
 
 hlfir::Entity hlfir::createStackTempFromMold(mlir::Location loc,

flang/lib/Optimizer/Builder/MutableBox.cpp

@@ -349,10 +349,7 @@ mlir::Value fir::factory::createUnallocatedBox(
   const bool isAssumedRank = baseBoxType.isAssumedRank();
   if (isAssumedRank)
     baseBoxType = baseBoxType.getBoxTypeWithNewShape(/*rank=*/0);
-  auto baseAddrType = baseBoxType.getEleTy();
-  if (!fir::isa_ref_type(baseAddrType))
-    baseAddrType =
-        builder.getRefType(baseAddrType, fir::isa_volatile_type(baseBoxType));
+  auto baseAddrType = baseBoxType.getBaseAddressType();
   auto type = fir::unwrapRefType(baseAddrType);
   auto eleTy = fir::unwrapSequenceType(type);
   if (auto recTy = mlir::dyn_cast<fir::RecordType>(eleTy))
@@ -982,3 +979,20 @@ mlir::Value fir::factory::genNullBoxStorage(fir::FirOpBuilder &builder,
   builder.create<fir::StoreOp>(loc, nullBox, boxStorage);
   return boxStorage;
 }
+
+mlir::Value fir::factory::getAndEstablishBoxStorage(
+    fir::FirOpBuilder &builder, mlir::Location loc, fir::BaseBoxType boxTy,
+    mlir::Value shape, llvm::ArrayRef<mlir::Value> typeParams,
+    mlir::Value polymorphicMold) {
+  mlir::Value boxStorage = builder.createTemporary(loc, boxTy);
+  mlir::Value nullAddr =
+      builder.createNullConstant(loc, boxTy.getBaseAddressType());
+  mlir::Value box =
+      builder.create<fir::EmboxOp>(loc, boxTy, nullAddr, shape,
+                                   /*emptySlice=*/mlir::Value{},
+                                   fir::factory::elideLengthsAlreadyInType(
+                                       boxTy.unwrapInnerType(), typeParams),
+                                   polymorphicMold);
+  builder.create<fir::StoreOp>(loc, box, boxStorage);
+  return boxStorage;
+}

flang/lib/Optimizer/CodeGen/LowerRepackArrays.cpp

@@ -190,12 +190,10 @@ mlir::Value PackArrayConversion::allocateTempBuffer(
   if (useStack && canAllocateTempOnStack(origBox))
     assert(!isHeapAllocation && "temp must have been allocated on the stack");
 
-  if (isHeapAllocation)
-    if (auto baseType = mlir::dyn_cast<fir::ReferenceType>(base.getType()))
-      if (mlir::isa<fir::BaseBoxType>(baseType.getEleTy()))
-        return builder.create<fir::LoadOp>(loc, base);
-
   mlir::Type ptrType = base.getType();
+  if (llvm::isa<fir::BaseBoxType>(ptrType))
+    return base;
+
   mlir::Type tempBoxType = fir::BoxType::get(mlir::isa<fir::HeapType>(ptrType)
                                                  ? ptrType
                                                  : fir::unwrapRefType(ptrType));

flang/lib/Optimizer/Dialect/FIRType.cpp

@@ -1407,6 +1407,13 @@ mlir::Type BaseBoxType::getEleTy() const {
           [](auto type) { return type.getEleTy(); });
 }
 
+mlir::Type BaseBoxType::getBaseAddressType() const {
+  mlir::Type eleTy = getEleTy();
+  if (fir::isa_ref_type(eleTy))
+    return eleTy;
+  return fir::ReferenceType::get(eleTy, isVolatile());
+}
+
 mlir::Type BaseBoxType::unwrapInnerType() const {
   return fir::unwrapInnerType(getEleTy());
 }
@@ -1492,6 +1499,7 @@ bool fir::BaseBoxType::isPointer() const {
   return llvm::isa<fir::PointerType>(getEleTy());
 }
 
+bool BaseBoxType::isVolatile() const { return fir::isa_volatile_type(*this); }
 //===----------------------------------------------------------------------===//
 // FIROpsDialect
 //===----------------------------------------------------------------------===//

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

@@ -1586,7 +1586,7 @@ void hlfir::AssociateOp::build(mlir::OpBuilder &builder,
   mlir::Type firVarType;
   auto sourceExprType = mlir::dyn_cast<hlfir::ExprType>(source.getType());
   if (sourceExprType && sourceExprType.isPolymorphic())
-    firVarType = fir::ClassType::get(fir::HeapType::get(dataType));
+    firVarType = fir::ClassType::get(dataType);
   else
     firVarType = fir::ReferenceType::get(dataType);