diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp --- a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp +++ b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp @@ -320,10 +320,18 @@ } }; -static bool allOtherUsesAreDestroys(mlir::Value value, - mlir::Operation *currentUse) { +/// The current hlfir.associate lowering does not handle multiple uses of a +/// non-trivial expression value because it generates the cleanup for the +/// expression bufferization at hlfir.end_associate. If there was more than one +/// hlfir.end_associate, it would be cleaned up multiple times, perhaps before +/// one of the other uses. +/// hlfir.shape_of doesn't actually use the bufferized expression so it doesn't +/// care when the expression buffer is deallocated. +static bool allOtherUsesAreSafeForAssociate(mlir::Value value, + mlir::Operation *currentUse) { for (mlir::Operation *useOp : value.getUsers()) - if (!mlir::isa(useOp) && useOp != currentUse) + if (!mlir::isa(useOp) && + !mlir::isa(useOp) && useOp != currentUse) return false; return true; } @@ -389,8 +397,9 @@ // If this is the last use of the expression value and this is an hlfir.expr // that was bufferized, re-use the storage. // Otherwise, create a temp and assign the storage to it. - if (!isTrivialValue && allOtherUsesAreDestroys(associate.getSource(), - associate.getOperation())) { + if (!isTrivialValue && + allOtherUsesAreSafeForAssociate(associate.getSource(), + associate.getOperation())) { // Re-use hlfir.expr buffer if this is the only use of the hlfir.expr // outside of the hlfir.destroy. Take on the cleaning-up responsibility // for the related hlfir.end_associate, and erase the hlfir.destroy (if