[flang][hlfir] Avoid expr buffer reuse when end_associate may cycle.

If end_associate may execute more times than the expr value producer,
then it cannot take ownership of the expr buffer. Otherwise, it may
result in double-free errors.
Note that the LIT test exposes a different issue with fir.alloca
inside the do-loop produced for hlfir.elemental. This may cause
out-of-stack conditions in valid Fortran programs that are not expected
to run out of stack. I will create an issue for this.

Reviewed By: tblah

Differential Revision: https://reviews.llvm.org/D155778

2 files changed, 72 insertions, 1 deletions

diff --git a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
index 86fa2f4..bb110d4 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp
@@ -379,9 +379,26 @@ struct GetLengthOpConversion
 /// 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.
+/// Note that we have to be careful about expressions used by a single
+/// hlfir.end_associate that may be executed more times than the producer
+/// of the expression value. This may also cause multiple clean-ups
+/// for the same memory (e.g. cause double-free errors). For example,
+/// hlfir.end_associate inside hlfir.elemental may cause such issues
+/// for expressions produced outside of hlfir.elemental.
 static bool allOtherUsesAreSafeForAssociate(mlir::Value value,
                                             mlir::Operation *currentUse,
                                             mlir::Operation *endAssociate) {
+  // If value producer is from a different region than
+  // hlfir.associate/end_associate, then conservatively assume
+  // that the hlfir.end_associate may execute more times than
+  // the value producer.
+  // TODO: this may be improved for operations that cannot
+  // result in multiple executions (e.g. ifOp).
+  if (value.getParentRegion() != currentUse->getParentRegion() ||
+      (endAssociate &&
+       value.getParentRegion() != endAssociate->getParentRegion()))
+    return false;
+
   for (mlir::Operation *useOp : value.getUsers())
     if (!mlir::isa<hlfir::DestroyOp>(useOp) && useOp != currentUse) {
       // hlfir.shape_of and hlfir.get_length will not disrupt cleanup so it is
@@ -504,7 +521,7 @@ struct AssociateOpConversion
     }
     // non-trivial value with more than one use. We will have to make a copy and
     // use that
-    hlfir::Entity source = hlfir::Entity{adaptor.getSource()};
+    hlfir::Entity source = hlfir::Entity{bufferizedExpr};
     auto [temp, cleanup] = createTempFromMold(loc, builder, source);
     builder.create<hlfir::AssignOp>(loc, source, temp, temp.isAllocatable(),
                                     /*keep_lhs_length_if_realloc=*/false,
diff --git a/flang/test/HLFIR/associate-codegen.fir b/flang/test/HLFIR/associate-codegen.fir
index 4bd67bd..493a8bb 100644
--- a/flang/test/HLFIR/associate-codegen.fir
+++ b/flang/test/HLFIR/associate-codegen.fir
@@ -366,3 +366,57 @@ func.func private @take_i4(!fir.ref<i32>)
 func.func private @take_r4(!fir.ref<f32>)
 func.func private @take_l4(!fir.ref<!fir.logical<4>>)
 func.func private @take_c(!fir.boxchar<1>)
+
+// Test the hlfir.associate/hlfir.end_associate does not take ownership over
+// a single-use hlfir.expr if hlfir.end_associate might be executed more times
+// than the producer of hlfir.expr. This might cause double-free effects.
+func.func @_QPtest_multiple_expr_uses_inside_elemental() {
+  %true = arith.constant true
+  %18 = fir.undefined !fir.heap<!fir.char<1,?>>
+  %17 = fir.undefined index
+  %19:2 = hlfir.declare %18 typeparams %17 {uniq_name = ".tmp.intrinsic_result"} : (!fir.heap<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.heap<!fir.char<1,?>>)
+  %20 = hlfir.as_expr %19#0 move %true : (!fir.boxchar<1>, i1) -> !hlfir.expr<!fir.char<1,?>>
+  %21 = fir.undefined index
+  %22 = fir.shape %21 : (index) -> !fir.shape<1>
+  %23 = hlfir.elemental %22 unordered : (!fir.shape<1>) -> !hlfir.expr<?x!fir.logical<4>> {
+  ^bb0(%arg2: index):
+    %35:3 = hlfir.associate %20 typeparams %17 {uniq_name = "adapt.valuebyref"} : (!hlfir.expr<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>, i1)
+    hlfir.end_associate %35#1, %35#2 : !fir.ref<!fir.char<1,?>>, i1
+    %ci1 = arith.constant 1 : i1
+    %42 = fir.convert %ci1 : (i1) -> !fir.logical<4>
+    hlfir.yield_element %42 : !fir.logical<4>
+  }
+  return
+}
+// CHECK-LABEL:   func.func @_QPtest_multiple_expr_uses_inside_elemental() {
+// CHECK:           %[[VAL_2:.*]] = arith.constant true
+// CHECK:           %[[VAL_3:.*]] = fir.undefined !fir.heap<!fir.char<1,?>>
+// CHECK:           %[[VAL_4:.*]] = fir.undefined index
+// CHECK:           %[[VAL_5:.*]]:2 = hlfir.declare %[[VAL_3]] typeparams %[[VAL_4]] {uniq_name = ".tmp.intrinsic_result"} : (!fir.heap<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.heap<!fir.char<1,?>>)
+// CHECK:           %[[VAL_6:.*]] = fir.undefined tuple<!fir.boxchar<1>, i1>
+// CHECK:           %[[VAL_7:.*]] = fir.insert_value %[[VAL_6]], %[[VAL_2]], [1 : index] : (tuple<!fir.boxchar<1>, i1>, i1) -> tuple<!fir.boxchar<1>, i1>
+// CHECK:           %[[VAL_8:.*]] = fir.insert_value %[[VAL_7]], %[[VAL_5]]#0, [0 : index] : (tuple<!fir.boxchar<1>, i1>, !fir.boxchar<1>) -> tuple<!fir.boxchar<1>, i1>
+// CHECK:           %[[VAL_9:.*]] = fir.undefined index
+// CHECK:           %[[VAL_10:.*]] = fir.shape %[[VAL_9]] : (index) -> !fir.shape<1>
+// CHECK:           %[[VAL_11:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>, %[[VAL_9]] {bindc_name = ".tmp.array", uniq_name = ""}
+// CHECK:           %[[VAL_12:.*]]:2 = hlfir.declare %[[VAL_11]](%[[VAL_10]]) {uniq_name = ".tmp.array"} : (!fir.heap<!fir.array<?x!fir.logical<4>>>, !fir.shape<1>) -> (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.heap<!fir.array<?x!fir.logical<4>>>)
+// CHECK:           %[[VAL_13:.*]] = arith.constant true
+// CHECK:           %[[VAL_14:.*]] = arith.constant 1 : index
+// CHECK:           fir.do_loop %[[VAL_15:.*]] = %[[VAL_14]] to %[[VAL_9]] step %[[VAL_14]] unordered {
+// CHECK:             %[[VAL_16:.*]] = fir.alloca !fir.char<1,?>(%[[VAL_4]] : index) {bindc_name = ".tmp"}
+// CHECK:             %[[VAL_17:.*]] = arith.constant false
+// CHECK:             %[[VAL_18:.*]]:2 = hlfir.declare %[[VAL_16]] typeparams %[[VAL_4]] {uniq_name = ".tmp"} : (!fir.ref<!fir.char<1,?>>, index) -> (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
+// CHECK:             hlfir.assign %[[VAL_5]]#0 to %[[VAL_18]]#0 temporary_lhs : !fir.boxchar<1>, !fir.boxchar<1>
+// CHECK:             %[[VAL_19:.*]] = fir.undefined tuple<!fir.boxchar<1>, i1>
+// CHECK:             %[[VAL_20:.*]] = fir.insert_value %[[VAL_19]], %[[VAL_17]], [1 : index] : (tuple<!fir.boxchar<1>, i1>, i1) -> tuple<!fir.boxchar<1>, i1>
+// CHECK:             %[[VAL_21:.*]] = fir.insert_value %[[VAL_20]], %[[VAL_18]]#0, [0 : index] : (tuple<!fir.boxchar<1>, i1>, !fir.boxchar<1>) -> tuple<!fir.boxchar<1>, i1>
+// CHECK:             %[[VAL_22:.*]] = arith.constant true
+// CHECK:             %[[VAL_23:.*]] = fir.convert %[[VAL_22]] : (i1) -> !fir.logical<4>
+// CHECK:             %[[VAL_24:.*]] = hlfir.designate %[[VAL_12]]#0 (%[[VAL_15]])  : (!fir.box<!fir.array<?x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
+// CHECK:             hlfir.assign %[[VAL_23]] to %[[VAL_24]] temporary_lhs : !fir.logical<4>, !fir.ref<!fir.logical<4>>
+// CHECK:           }
+// CHECK:           %[[VAL_25:.*]] = fir.undefined tuple<!fir.box<!fir.array<?x!fir.logical<4>>>, i1>
+// CHECK:           %[[VAL_26:.*]] = fir.insert_value %[[VAL_25]], %[[VAL_13]], [1 : index] : (tuple<!fir.box<!fir.array<?x!fir.logical<4>>>, i1>, i1) -> tuple<!fir.box<!fir.array<?x!fir.logical<4>>>, i1>
+// CHECK:           %[[VAL_27:.*]] = fir.insert_value %[[VAL_26]], %[[VAL_12]]#0, [0 : index] : (tuple<!fir.box<!fir.array<?x!fir.logical<4>>>, i1>, !fir.box<!fir.array<?x!fir.logical<4>>>) -> tuple<!fir.box<!fir.array<?x!fir.logical<4>>>, i1>
+// CHECK:           return
+// CHECK:         }