Skip to content

Commit

Permalink
[flang] handle allocatable components when creating array temps
Browse files Browse the repository at this point in the history
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.

Differential Revision: https://reviews.llvm.org/D121892
  • Loading branch information
jeanPerier committed Mar 17, 2022
1 parent 0e694e1 commit 3ed899c
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 12 deletions.
64 changes: 52 additions & 12 deletions flang/lib/Optimizer/Transforms/ArrayValueCopy.cpp
Expand Up @@ -12,6 +12,7 @@
#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"
Expand Down Expand Up @@ -1008,6 +1009,50 @@ findNonconstantExtents(mlir::Type memrefTy,
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
Expand Down Expand Up @@ -1039,11 +1084,8 @@ class ArrayUpdateConversionBase : public mlir::OpRewritePattern<ArrayOp> {
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);
Expand All @@ -1061,7 +1103,7 @@ class ArrayUpdateConversionBase : public mlir::OpRewritePattern<ArrayOp> {
// Copy out.
genArrayCopy</*copyIn=*/false>(store.getLoc(), rewriter, store.getMemref(),
allocmem, shapeOp, store.getSlice(), load);
rewriter.create<FreeMemOp>(loc, allocmem);
genTempCleanUp(rewriter);
return coor;
}

Expand Down Expand Up @@ -1091,11 +1133,9 @@ class ArrayUpdateConversionBase : public mlir::OpRewritePattern<ArrayOp> {
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);
Expand All @@ -1113,7 +1153,7 @@ class ArrayUpdateConversionBase : public mlir::OpRewritePattern<ArrayOp> {
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
Expand Down
55 changes: 55 additions & 0 deletions flang/test/Fir/array-value-copy-3.fir
@@ -0,0 +1,55 @@
// 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]]

0 comments on commit 3ed899c

Please sign in to comment.