[flang] handle allocatable components when creating array temps

[jeanPerier]

Summary

When creating an array temporary in the array copy pass, care must be
taken with allocatable components. The element components needs to be
given a clean unallocated status before being used in the assignments.
This is because assignment of allocatable components makes deep copy,
and may cause deallocation of the previous value if it was allocated.
Hence the previous allocation status cannot be let undefined.

On top of that, when cleaning-up the temp, all allocatable components
that may have been allocated must be deallocated.

This patch implements this by centralizing the code making and cleaning
array temps in ArrayValueCopy.cpp, and by calling Initialize and Destroy
runtime entry points when they are allocatable components.

[clementval]

LGTM

flang/lib/Optimizer/Transforms/ArrayValueCopy.cpp

 //===-- ArrayValueCopy.cpp ------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "flang/Lower/Todo.h"
 #include "flang/Optimizer/Builder/Array.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Factory.h"
+#include "flang/Optimizer/Builder/Runtime/Derived.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Support/FIRContext.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/Debug.h"
for (auto [s, x] : llvm::zip(seqTy.getShape(), extents))
     if (s == SequenceType::getUnknownExtent())
       nce.emplace_back(x);
   if (extents.size() > seqTy.getShape().size())
     for (auto x : extents.drop_front(seqTy.getShape().size()))
       nce.emplace_back(x);
   return nce;
 }
 
+/// Allocate temporary storage for an ArrayLoadOp \load and initialize any
+/// allocatable direct components of the array elements with an unallocated
+/// status. Returns the temporary address as well as a callback to generate the
+/// temporary clean-up once it has been used. The clean-up will take care of
+/// deallocating all the element allocatable components that may have been
+/// allocated while using the temporary.
+static std::pair<mlir::Value,
+                 std::function<void(mlir::PatternRewriter &rewriter)>>
+allocateArrayTemp(mlir::Location loc, mlir::PatternRewriter &rewriter,
+                  ArrayLoadOp load, llvm::ArrayRef<mlir::Value> extents,
+                  mlir::Value shape) {
+  mlir::Type baseType = load.getMemref().getType();
+  llvm::SmallVector<mlir::Value> nonconstantExtents =
+      findNonconstantExtents(baseType, extents);
+  llvm::SmallVector<mlir::Value> typeParams =
+      genArrayLoadTypeParameters(loc, rewriter, load);
+  mlir::Value allocmem = rewriter.create<AllocMemOp>(
+      loc, dyn_cast_ptrOrBoxEleTy(baseType), typeParams, nonconstantExtents);
+  mlir::Type eleType =
+      fir::unwrapSequenceType(fir::unwrapPassByRefType(baseType));
+  if (fir::isRecordWithAllocatableMember(eleType)) {
+    // The allocatable component descriptors need to be set to a clean
+    // deallocated status before anything is done with them.
+    mlir::Value box = rewriter.create<fir::EmboxOp>(
+        loc, fir::BoxType::get(baseType), allocmem, shape,
+        /*slice=*/mlir::Value{}, typeParams);
+    auto module = load->getParentOfType<mlir::ModuleOp>();
+    FirOpBuilder builder(rewriter, getKindMapping(module));
+    runtime::genDerivedTypeInitialize(builder, loc, box);
+    // Any allocatable component that may have been allocated must be
+    // deallocated during the clean-up.
+    auto cleanup = [=](mlir::PatternRewriter &r) {
+      FirOpBuilder builder(r, getKindMapping(module));
+      runtime::genDerivedTypeDestroy(builder, loc, box);
+      r.create<FreeMemOp>(loc, allocmem);
+    };
+    return {allocmem, cleanup};
+  }
+  auto cleanup = [=](mlir::PatternRewriter &r) {
+    r.create<FreeMemOp>(loc, allocmem);
+  };
+  return {allocmem, cleanup};
+}
+
 namespace {
 /// Conversion of fir.array_update and fir.array_modify Ops.
 /// If there is a conflict for the update, then we need to perform a
 /// copy-in/copy-out to preserve the original values of the array. If there is
 /// no conflict, then it is save to eschew making any copies.
 template <typename ArrayOp>
 class ArrayUpdateConversionBase : public mlir::OpRewritePattern<ArrayOp> {
 public:
mlir::Value referenceToClone(mlir::Location loc,
     auto load = mlir::cast<ArrayLoadOp>(loadOp);
     auto eleTy = access.getType();
     rewriter.setInsertionPoint(loadOp);
     // Copy in.
     llvm::SmallVector<mlir::Value> extents;
     bool copyUsingSlice = false;
     auto shapeOp = getOrReadExtentsAndShapeOp(loc, rewriter, load, extents,
                                               copyUsingSlice);
-    llvm::SmallVector<mlir::Value> nonconstantExtents =
-        findNonconstantExtents(load.getMemref().getType(), extents);
-    auto allocmem = rewriter.create<AllocMemOp>(
-        loc, dyn_cast_ptrOrBoxEleTy(load.getMemref().getType()),
-        genArrayLoadTypeParameters(loc, rewriter, load), nonconstantExtents);
+    auto [allocmem, genTempCleanUp] =
+        allocateArrayTemp(loc, rewriter, load, extents, shapeOp);
     genArrayCopy</*copyIn=*/true>(load.getLoc(), rewriter, allocmem,
                                   load.getMemref(), shapeOp, load.getSlice(),
                                   load);
     // Generate the reference for the access.
     rewriter.setInsertionPoint(op);
     auto coor =
         genCoorOp(rewriter, loc, getEleTy(load.getType()), eleTy, allocmem,
                   shapeOp, copyUsingSlice ? mlir::Value{} : load.getSlice(),
                   access.getIndices(), load.getTypeparams(),
                   access->hasAttr(factory::attrFortranArrayOffsets()));
     // Copy out.
     auto *storeOp = useMap.lookup(loadOp);
     auto store = mlir::cast<ArrayMergeStoreOp>(storeOp);
     rewriter.setInsertionPoint(storeOp);
     // Copy out.
     genArrayCopy</*copyIn=*/false>(store.getLoc(), rewriter, store.getMemref(),
                                    allocmem, shapeOp, store.getSlice(), load);
-    rewriter.create<FreeMemOp>(loc, allocmem);
+    genTempCleanUp(rewriter);
     return coor;
   }
 
   /// Copy the RHS element into the LHS and insert copy-in/copy-out between a
   /// temp and the LHS if the analysis found potential overlaps between the RHS
   /// and LHS arrays. The element copy generator must be provided in \p
   /// assignElement. \p update must be the ArrayUpdateOp or the ArrayModifyOp.
   /// Returns the address of the LHS element inside the loop and the LHS
if (analysis.hasPotentialConflict(loadOp)) {
       // the lhs array. This yields Fortran's copy-in copy-out array semantics.
       LLVM_DEBUG(llvm::outs() << "Yes, conflict was found\n");
       rewriter.setInsertionPoint(loadOp);
       // Copy in.
       llvm::SmallVector<mlir::Value> extents;
       bool copyUsingSlice = false;
       auto shapeOp = getOrReadExtentsAndShapeOp(loc, rewriter, load, extents,
                                                 copyUsingSlice);
-      llvm::SmallVector<mlir::Value> nonconstantExtents =
-          findNonconstantExtents(load.getMemref().getType(), extents);
-      auto allocmem = rewriter.create<AllocMemOp>(
-          loc, dyn_cast_ptrOrBoxEleTy(load.getMemref().getType()),
-          genArrayLoadTypeParameters(loc, rewriter, load), nonconstantExtents);
+      auto [allocmem, genTempCleanUp] =
+          allocateArrayTemp(loc, rewriter, load, extents, shapeOp);
+
       genArrayCopy</*copyIn=*/true>(load.getLoc(), rewriter, allocmem,
                                     load.getMemref(), shapeOp, load.getSlice(),
                                     load);
       rewriter.setInsertionPoint(op);
       auto coor = genCoorOp(
           rewriter, loc, getEleTy(load.getType()), lhsEltRefType, allocmem,
           shapeOp, copyUsingSlice ? mlir::Value{} : load.getSlice(),
           update.getIndices(), load.getTypeparams(),
           update->hasAttr(factory::attrFortranArrayOffsets()));
       assignElement(coor);
       auto *storeOp = useMap.lookup(loadOp);
       auto store = mlir::cast<ArrayMergeStoreOp>(storeOp);
       rewriter.setInsertionPoint(storeOp);
       // Copy out.
       genArrayCopy</*copyIn=*/false>(store.getLoc(), rewriter,
                                      store.getMemref(), allocmem, shapeOp,
                                      store.getSlice(), load);
-      rewriter.create<FreeMemOp>(loc, allocmem);
+      genTempCleanUp(rewriter);
       return {coor, load.getResult()};
     }
     // Otherwise, when there is no conflict (a possible loop-carried
     // dependence), the lhs array can be updated in place.
     LLVM_DEBUG(llvm::outs() << "No, conflict wasn't found\n");
     rewriter.setInsertionPoint(op);
     auto coorTy = getEleTy(load.getType());
     auto coor = genCoorOp(rewriter, loc, coorTy, lhsEltRefType,

flang/test/Fir/array-value-copy-3.fir

+// Test overlapping assignment of derived type arrays with allocatable components.
+// This requires initializing the allocatable components to an unallocated status
+// before they can be used in component assignments, and to deallocate the components
+// that may have been allocated in the end.
+
+// RUN: fir-opt --array-value-copy %s | FileCheck %s
+
+
+!t_with_alloc_comp = type !fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>
+func private @custom_assign(!fir.ref<!t_with_alloc_comp>, !fir.ref<!t_with_alloc_comp>)
+func @test_overlap_with_alloc_components(%arg0: !fir.ref<!fir.array<10x!t_with_alloc_comp>>) {
+  %0 = fir.alloca !fir.box<!t_with_alloc_comp>
+  %c10 = arith.constant 10 : index
+  %c9 = arith.constant 9 : index
+  %c1 = arith.constant 1 : index
+  %c-1 = arith.constant -1 : index
+  %c0 = arith.constant 0 : index
+  %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %6 = fir.slice %c10, %c1, %c-1 : (index, index, index) -> !fir.slice<1>
+  %2 = fir.array_load %arg0(%1) : (!fir.ref<!fir.array<10x!t_with_alloc_comp>>, !fir.shape<1>) -> !fir.array<10x!t_with_alloc_comp>
+  %7 = fir.array_load %arg0(%1) [%6] : (!fir.ref<!fir.array<10x!t_with_alloc_comp>>, !fir.shape<1>, !fir.slice<1>) -> !fir.array<10x!t_with_alloc_comp>
+  %9 = fir.do_loop %arg1 = %c0 to %c9 step %c1 unordered iter_args(%arg2 = %2) -> (!fir.array<10x!t_with_alloc_comp>) {
+    %10 = fir.array_access %7, %arg1 : (!fir.array<10x!t_with_alloc_comp>, index) -> !fir.ref<!t_with_alloc_comp>
+    %11 = fir.array_access %arg2, %arg1 : (!fir.array<10x!t_with_alloc_comp>, index) -> !fir.ref<!t_with_alloc_comp>
+    fir.call @custom_assign(%11, %10) : (!fir.ref<!t_with_alloc_comp>, !fir.ref<!t_with_alloc_comp>) -> none
+    %19 = fir.array_amend %arg2, %11 : (!fir.array<10x!t_with_alloc_comp>, !fir.ref<!t_with_alloc_comp>) -> !fir.array<10x!t_with_alloc_comp>
+    fir.result %19 : !fir.array<10x!t_with_alloc_comp>
+  }
+  fir.array_merge_store %2, %9 to %arg0 : !fir.array<10x!t_with_alloc_comp>, !fir.array<10x!t_with_alloc_comp>, !fir.ref<!fir.array<10x!t_with_alloc_comp>>
+  return
+}
+
+// CHECK-LABEL: func @test_overlap_with_alloc_components(
+// CHECK-SAME:  %[[VAL_0:.*]]: !fir.ref<!fir.array<10x!fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>) {
+// CHECK:   %[[VAL_4:.*]] = arith.constant 10 : index
+// CHECK:   %[[VAL_6:.*]] = arith.constant 1 : index
+// CHECK:   %[[VAL_7:.*]] = arith.constant -1 : index
+// CHECK:   %[[VAL_9:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
+// CHECK:   %[[VAL_10:.*]] = fir.slice %[[VAL_4]], %[[VAL_6]], %[[VAL_7]] : (index, index, index) -> !fir.slice<1>
+// CHECK:   %[[VAL_11:.*]] = fir.allocmem !fir.array<10x!fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>
+// CHECK:   %[[VAL_12:.*]] = fir.embox %[[VAL_11]](%[[VAL_9]]) : (!fir.heap<!fir.array<10x!fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>, !fir.shape<1>) -> !fir.box<!fir.ref<!fir.array<10x!fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>>
+// CHECK:   %[[VAL_16:.*]] = fir.convert %[[VAL_12]] : (!fir.box<!fir.ref<!fir.array<10x!fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>>) -> !fir.box<none>
+// CHECK:   fir.call @_FortranAInitialize(%[[VAL_16]], %{{.*}}, %{{.*}}) : (!fir.box<none>, !fir.ref<i8>, i32) -> none
+// CHECK:   fir.do_loop {{.*}} {
+// CHECK:     fir.call @_FortranAAssign
+// CHECK:   }
+// CHECK:   fir.do_loop {{.*}} {
+// CHECK:     fir.call @custom_assign
+// CHECK:   }
+// CHECK:   fir.do_loop %{{.*}} {
+// CHECK:     fir.call @_FortranAAssign
+// CHECK:   }
+// CHECK:   %[[VAL_72:.*]] = fir.convert %[[VAL_12]] : (!fir.box<!fir.ref<!fir.array<10x!fir.type<t{i:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>>) -> !fir.box<none>
+// CHECK:   %[[VAL_73:.*]] = fir.call @_FortranADestroy(%[[VAL_72]]) : (!fir.box<none>) -> none
+// CHECK:   fir.freemem %[[VAL_11]]