[flang][hlfir] Do not reuse hlfir.expr mask when saving RHS.

[jeanPerier]

Summary

In WHERE and masked FORALL assignment, both the mask and the
RHS may need to be saved in some temporary storage before evaluating
the assignment.

The code was trying to "optimize" that case when evaluating the RHS
by not fetching the mask temporary that was just created, but in simple
cases of WHERE construct where the evaluated mask is an hlfir.expr,
this caused the hlfir.expr to be both used in an hlfir.associate and
later in an hlfir.apply to create the fir.if to mask the RHS evaluation.
This double usage prevents codegen from inlining the hlfir.expr at the
hlfir.apply, and from "moving" the hlfir.expr storage into the temp
during hlfir.associate bufferization. So this is pessimizing the code:
this would lead to created two mask array temporary storages

This was caught by the unexpectedly high number of "not yet implemented:
hlfir.associate of hlfir.expr with more than one use" that were firing.

Use the mask temporary instead (the hlfir.associate result) when possible.
Some temporary (the "inlined stack") do not support fetching and pushing
in the same run (a single counter is used to keep track of the fetching
and pushing position). Add a canBeFetchedAfterPush() for safety,
but this limitation is anyway not relevant for hlfir.expr since the
inlined stack is only used to save "trivial" scalars.

Also update the temporary storage name to only indicate "forall" if
the top level construct is a FORALL. This is not a very precise name,
but it should at least give a correct context to indicate in the IR
why some temporary array storage was created.

[vzakhari]

Looks good to me. Thank you!

[tblah]

LGTM, thanks!

flang/include/flang/Optimizer/Builder/TemporaryStorage.h

public:
   void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
 
   /// Move the temporary storage into a rank one array expression value
   /// (hlfir.expr<?xT>). The temporary should not be used anymore after this
   /// call.
   hlfir::Entity moveStackAsArrayExpr(mlir::Location loc,
                                      fir::FirOpBuilder &builder);
 
+  ///  "fetch" cannot be called right after "pushValue" because the counter is
+  ///  both used for pushing and fetching.
+  bool canBeFetchedAfterPush() const { return false; }
+
 private:
   /// Allocate the temporary on the heap.
   const bool allocateOnHeap;
   /// Counter to keep track of the insertion or fetching position.
   Counter counter;
   /// Temporary storage.
   mlir::Value temp;
 };
 
 /// Structure to hold the value of a single entity.
 class SimpleCopy {
 public:
   SimpleCopy(mlir::Location loc, fir::FirOpBuilder &builder,
              hlfir::Entity source, llvm::StringRef tempName);
 
   void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                  mlir::Value value) {
     assert(false && "must not be called: value already set");
   }
   void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder){};
   mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
     return copy.getBase();
   }
   void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
+  bool canBeFetchedAfterPush() const { return true; }
 
 public:
   /// Temporary storage for the copy.
   hlfir::AssociateOp copy;
 };
 
 /// Data structure to stack any kind of values with the same static type and
 /// rank. Each value may have different type parameters, bounds, and dynamic
 /// type. Fetching value N will return a value with the same dynamic type,
 /// bounds, and type parameters as the Nth value that was pushed. It is
 /// implemented using runtime.
 class AnyValueStack {
 public:
   AnyValueStack(mlir::Location loc, fir::FirOpBuilder &builder,
                 mlir::Type valueStaticType);
 
   void pushValue(mlir::Location loc, fir::FirOpBuilder &builder,
                  mlir::Value value);
   void resetFetchPosition(mlir::Location loc, fir::FirOpBuilder &builder);
   mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder);
   void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
+  bool canBeFetchedAfterPush() const { return true; }
 
 private:
   /// Keep the original value type. Values may be stored by the runtime
   /// with a different type (i1 cannot be passed by descriptor).
   mlir::Type valueStaticType;
   /// Runtime cookie created by the runtime. It is a pointer to an opaque
   /// runtime data structure that manages the stack.
   mlir::Value opaquePtr;
public:
   }
   mlir::Value fetch(mlir::Location loc, fir::FirOpBuilder &builder) {
     return std::visit([&](auto &temp) { return temp.fetch(loc, builder); },
                       impl);
   }
   void destroy(mlir::Location loc, fir::FirOpBuilder &builder) {
     std::visit([&](auto &temp) { temp.destroy(loc, builder); }, impl);
   }
+  /// Can "fetch" be called to get the last value pushed with
+  /// "pushValue"?
+  bool canBeFetchedAfterPush() const {
+    return std::visit([&](auto &temp) { return temp.canBeFetchedAfterPush(); },
+                      impl);
+  }
 
 private:
   std::variant<HomogeneousScalarStack, SimpleCopy, AnyValueStack> impl;
 };
 } // namespace fir::factory
 #endif // FORTRAN_OPTIMIZER_BUILDER_TEMPORARYSTORAGE_H

flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp

if (!loopExtent)
       loopExtent = extent;
     else
       loopExtent = builder.create<mlir::arith::MulIOp>(loc, loopExtent, extent);
   }
   assert(loopExtent && "loopNest must not be empty");
   return loopExtent;
 }
 
+/// Return a name for temporary storage that indicates in which context
+/// the temporary storage was created.
+static llvm::StringRef
+getTempName(hlfir::OrderedAssignmentTreeOpInterface root) {
+  if (mlir::isa<hlfir::ForallOp>(root.getOperation()))
+    return ".tmp.forall";
+  if (mlir::isa<hlfir::WhereOp>(root.getOperation()))
+    return ".tmp.where";
+  return ".tmp.assign";
+}
+
 void OrderedAssignmentRewriter::generateSaveEntity(
     hlfir::SaveEntity savedEntity, bool willUseSavedEntityInSameRun) {
   mlir::Region &region = *savedEntity.yieldRegion;
   mlir::Location loc = region.getParentOp()->getLoc();
 
   if (!mlir::isa<hlfir::YieldOp>(region.back().back()))
     TODO(loc, "creating temporary storage for vector subscripted LHS");
 
   // Evaluate the region inside the loop nest (if any).
   auto [clonedValue, oldYield] = generateYieldedEntity(region);
   hlfir::Entity entity{clonedValue};
   if (isLeftHandSide(region)) // Need to save the address, not the values.
     TODO(loc, "creating temporary storage for LHS");
   else
     entity = hlfir::loadTrivialScalar(loc, builder, entity);
   mlir::Type entityType = entity.getType();
 
-  static constexpr char tempName[] = ".tmp.forall";
+  llvm::StringRef tempName = getTempName(root);
+  fir::factory::TemporaryStorage *temp = nullptr;
   if (constructStack.empty()) {
     // Value evaluated outside of any loops (this may be the first MASK of a
     // WHERE construct, or an LHS/RHS temp of hlfir.region_assign outside of
     // WHERE/FORALL).
-    insertSavedEntity(region,
-                      fir::factory::SimpleCopy(loc, builder, entity, tempName));
+    temp = insertSavedEntity(
+        region, fir::factory::SimpleCopy(loc, builder, entity, tempName));
   } else {
     // Need to create a temporary for values computed inside loops.
     // Create temporary storage outside of the loop nest given the entity
     // type (and the loop context).
-    fir::factory::TemporaryStorage *temp;
     llvm::SmallVector<fir::DoLoopOp> loopNest;
     bool loopShapeCanBePreComputed =
         currentLoopNestIterationNumberCanBeComputed(loopNest);
     doBeforeLoopNest([&] {
       /// For simple scalars inside loops whose total iteration number can be
       /// pre-computed, create a rank-1 array outside of the loops. It will be
       /// assigned/fetched inside the loops like a normal Fortran array given
       /// the iteration count.
doBeforeLoopNest([&] {
       }
     });
     // Inside the loop nest (and any fir.if if there are active masks), copy
     // the value to the temp and do clean-ups for the value if any.
     temp->pushValue(loc, builder, entity);
   }
 
   // Delay the clean-up if the entity will be used in the same run (i.e., the
-  // parent construct will be visited and needs to be lowered).
-  if (willUseSavedEntityInSameRun) {
+  // parent construct will be visited and needs to be lowered). When possible,
+  // this is not done for hlfir.expr because this use would prevent the
+  // hlfir.expr storage from being moved when creating the temporary in
+  // bufferization, and that would lead to an extra copy.
+  if (willUseSavedEntityInSameRun &&
+      (!temp->canBeFetchedAfterPush() ||
+       !mlir::isa<hlfir::ExprType>(entity.getType()))) {
     auto inserted =
         savedInCurrentRunBeforeUse.try_emplace(&region, entity, oldYield);
     assert(inserted.second && "entity must have been emplaced");
     (void)inserted;
   } else {
     generateCleanupIfAny(oldYield);
   }
 }

flang/test/HLFIR/order_assignments/impure-where.fir

hlfir.where {
     }
   }
   return
 }
 // CHECK-LABEL:   func.func @test_elsewhere_impure_mask(
 // CHECK:           %[[VAL_12:.*]] = fir.call @impure() : () -> !fir.heap<!fir.array<10x!fir.logical<4>>>
 // CHECK:           %[[VAL_21:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>, %[[extent:[^ ]*]]
 // CHECK:           %[[VAL_22:.*]] = fir.shape %[[extent]] : (index) -> !fir.shape<1>
-// CHECK:           %[[VAL_23:.*]]:2 = hlfir.declare %[[VAL_21]](%{{.*}}) {uniq_name = ".tmp.forall"}
+// CHECK:           %[[VAL_23:.*]]:2 = hlfir.declare %[[VAL_21]](%{{.*}}) {uniq_name = ".tmp.where"}
 // CHECK:           fir.do_loop
 // CHECK:             fir.if {{.*}} {
 // CHECK:             } else {
 // CHECK:               %[[VAL_28:.*]] = hlfir.designate %[[VAL_12]] (%{{.*}})
 // CHECK:               %[[VAL_29:.*]] = fir.load %[[VAL_28]] : !fir.ref<!fir.logical<4>>
 // CHECK:               %[[VAL_32:.*]] = hlfir.designate %[[VAL_23]]#0 (%{{.*}})  : (!fir.box<!fir.array<?x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
 // CHECK:               hlfir.assign %[[VAL_29]] to %[[VAL_32]] : !fir.logical<4>, !fir.ref<!fir.logical<4>>
 // CHECK:             }

flang/test/HLFIR/order_assignments/inlined-stack-temp.fir

 // CHECK:           %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] : (!fir.shape<1>) -> !hlfir.expr<?x!fir.logical<4>> {
 // CHECK:           ^bb0(%[[VAL_7:.*]]: index):
 // CHECK:             %[[VAL_8:.*]] = hlfir.designate %[[VAL_3]]#0 (%[[VAL_7]])  : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
 // CHECK:             %[[VAL_9:.*]] = fir.load %[[VAL_8]] : !fir.ref<i32>
 // CHECK:             %[[VAL_10:.*]] = arith.cmpi sgt, %[[VAL_9]], %[[VAL_2]] : i32
 // CHECK:             %[[VAL_11:.*]] = fir.convert %[[VAL_10]] : (i1) -> !fir.logical<4>
 // CHECK:             hlfir.yield_element %[[VAL_11]] : !fir.logical<4>
 // CHECK:           }
-// CHECK:           %[[VAL_12:.*]]:3 = hlfir.associate %[[VAL_13:.*]](%[[VAL_5]]) {uniq_name = ".tmp.forall"} : (!hlfir.expr<?x!fir.logical<4>>, !fir.shape<1>) -> (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.ref<!fir.array<?x!fir.logical<4>>>, i1)
+// CHECK:           %[[VAL_12:.*]]:3 = hlfir.associate %[[VAL_13:.*]](%[[VAL_5]]) {uniq_name = ".tmp.where"} : (!hlfir.expr<?x!fir.logical<4>>, !fir.shape<1>) -> (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.ref<!fir.array<?x!fir.logical<4>>>, i1)
 // CHECK:           hlfir.destroy %[[VAL_13]] : !hlfir.expr<?x!fir.logical<4>>
 // CHECK:           %[[VAL_14:.*]] = arith.constant 1 : index
 // CHECK:           fir.do_loop %[[VAL_15:.*]] = %[[VAL_14]] to %[[VAL_4]]#1 step %[[VAL_14]] {
 // CHECK:             %[[VAL_16:.*]] = hlfir.designate %[[VAL_12]]#0 (%[[VAL_15]])  : (!fir.box<!fir.array<?x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
 // CHECK:             %[[VAL_17:.*]] = fir.load %[[VAL_16]] : !fir.ref<!fir.logical<4>>
 // CHECK:             %[[VAL_18:.*]] = fir.convert %[[VAL_17]] : (!fir.logical<4>) -> i1
 // CHECK:             fir.if %[[VAL_18]] {
 // CHECK:               %[[VAL_19:.*]] = hlfir.designate %[[VAL_3]]#0 (%[[VAL_15]])  : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>

 // CHECK:           %[[VAL_12:.*]] = arith.subi %[[VAL_7]], %[[VAL_9]] : index
 // CHECK:           %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_9]] : index
 // CHECK:           %[[VAL_14:.*]] = arith.divsi %[[VAL_13]], %[[VAL_9]] : index
 // CHECK:           %[[VAL_15:.*]] = arith.cmpi sgt, %[[VAL_14]], %[[VAL_11]] : index
 // CHECK:           %[[VAL_16:.*]] = arith.select %[[VAL_15]], %[[VAL_14]], %[[VAL_11]] : index
 // CHECK:           %[[VAL_17:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_18:.*]] = arith.constant 1 : index
 // CHECK:           fir.store %[[VAL_17]] to %[[VAL_2]] : !fir.ref<index>
-// CHECK:           %[[VAL_19:.*]] = fir.allocmem !fir.array<?xi32>, %[[VAL_16]] {bindc_name = ".tmp.forall", uniq_name = ""}
+// CHECK:           %[[VAL_19:.*]] = fir.allocmem !fir.array<?xi32>, %[[VAL_16]] {bindc_name = ".tmp.where", uniq_name = ""}
 // CHECK:           %[[VAL_20:.*]] = fir.shape %[[VAL_16]] : (index) -> !fir.shape<1>
-// CHECK:           %[[VAL_21:.*]]:2 = hlfir.declare %[[VAL_19]](%[[VAL_20]]) {uniq_name = ".tmp.forall"} : (!fir.heap<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.heap<!fir.array<?xi32>>)
+// CHECK:           %[[VAL_21:.*]]:2 = hlfir.declare %[[VAL_19]](%[[VAL_20]]) {uniq_name = ".tmp.where"} : (!fir.heap<!fir.array<?xi32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.heap<!fir.array<?xi32>>)
 // CHECK:           fir.do_loop %[[VAL_22:.*]] = %[[VAL_9]] to %[[VAL_7]] step %[[VAL_9]] {
 // CHECK:             %[[VAL_23:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_22]])  : (!fir.ref<!fir.array<10x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
 // CHECK:             %[[VAL_24:.*]] = fir.load %[[VAL_23]] : !fir.ref<!fir.logical<4>>
 // CHECK:             %[[VAL_25:.*]] = fir.convert %[[VAL_24]] : (!fir.logical<4>) -> i1
 // CHECK:             fir.if %[[VAL_25]] {
 // CHECK:               %[[VAL_26:.*]] = hlfir.designate %[[VAL_10]] (%[[VAL_22]])  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
 // CHECK:               %[[VAL_27:.*]] = fir.load %[[VAL_26]] : !fir.ref<i32>
 // CHECK:               %[[VAL_28:.*]] = fir.load %[[VAL_2]] : !fir.ref<index>

flang/test/HLFIR/order_assignments/saving-mask-and-rhs.fir

+// Test when the saved mask is used in the same run to save
+// another value (like the RHS).
+// RUN: fir-opt %s --lower-hlfir-ordered-assignments | FileCheck %s
+
+func.func @saving_mask_and_rhs(%arg0: !fir.ref<!fir.array<10xi32>>) {
+  %c-1 = arith.constant -1 : index
+  %c1 = arith.constant 1 : index
+  %c0_i32 = arith.constant 0 : i32
+  %c10 = arith.constant 10 : index
+  %0 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %1:2 = hlfir.declare %arg0(%0) {uniq_name = "x"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
+  hlfir.where {
+    %2 = hlfir.elemental %0 : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> {
+    ^bb0(%arg1: index):
+      %3 = hlfir.designate %1#0 (%arg1)  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+      %4 = fir.load %3 : !fir.ref<i32>
+      %5 = arith.cmpi sgt, %4, %c0_i32 : i32
+      %6 = fir.convert %5 : (i1) -> !fir.logical<4>
+      hlfir.yield_element %6 : !fir.logical<4>
+    }
+    hlfir.yield %2 : !hlfir.expr<10x!fir.logical<4>> cleanup {
+      hlfir.destroy %2 : !hlfir.expr<10x!fir.logical<4>>
+    }
+  } do {
+    hlfir.region_assign {
+      hlfir.yield %1#0 : !fir.ref<!fir.array<10xi32>>
+    } to {
+      %2 = hlfir.designate %1#0 (%c10:%c1:%c-1)  shape %0 : (!fir.ref<!fir.array<10xi32>>, index, index, index, !fir.shape<1>) -> !fir.box<!fir.array<10xi32>>
+      hlfir.yield %2 : !fir.box<!fir.array<10xi32>>
+    }
+  }
+  return
+}
+
+
+// Creating mask temporary.
+
+// CHECK-LABEL:   func.func @saving_mask_and_rhs(
+// CHECK:  %[[VAL_8:.*]] = hlfir.elemental {{.*}} !hlfir.expr<10x!fir.logical<4>>
+// CHECK:  %[[VAL_14:.*]]:3 = hlfir.associate %[[VAL_8]]({{.*}}) {uniq_name = ".tmp.where"}
+
+// Creating RHS temporary using the mask temporary (and not the hlfir.elemental)
+
+// CHECK:  %[[VAL_25:.*]] = fir.allocmem !fir.array<?xi32>, %{{.*}} {bindc_name = ".tmp.where", uniq_name = ""}
+// CHECK:  %[[VAL_27:.*]]:2 = hlfir.declare %[[VAL_25]]({{.*}}) {uniq_name = ".tmp.where"}
+// CHECK:  fir.do_loop
+// CHECK:    %[[VAL_29:.*]] = hlfir.designate %[[VAL_14]]#0 ({{.*}})
+// CHECK:    %[[VAL_30:.*]] = fir.load %[[VAL_29]] : !fir.ref<!fir.logical<4>>
+// CHECK:    %[[VAL_31:.*]] = fir.convert %[[VAL_30]] : (!fir.logical<4>) -> i1
+// CHECK:    fir.if %[[VAL_31]] {
+// CHECK:      %[[VAL_36:.*]] = hlfir.designate %[[VAL_27]]#0 ({{.*}})
+// CHECK:      hlfir.assign %{{.*}} to %[[VAL_36]] : i32, !fir.ref<i32>
+// CHECK:    }
+// CHECK:  }
+
+func.func @forall_mask_and_rhs(%arg0: !fir.ref<!fir.array<10xi32>>) {
+  %c0_i32 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %c11 = arith.constant 11 : index
+  %0 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %1:2 = hlfir.declare %arg0(%0) {uniq_name = "x"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
+  hlfir.forall lb {
+    hlfir.yield %c1 : index
+  } ub {
+    hlfir.yield %c10 : index
+  }  (%arg1: index) {
+    hlfir.forall_mask {
+      %2 = arith.subi %c11, %arg1 : index
+      %3 = hlfir.designate %1#0 (%2)  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+      %4 = fir.load %3 : !fir.ref<i32>
+      %5 = arith.cmpi sgt, %4, %c0_i32 : i32
+      hlfir.yield %5 : i1
+    } do {
+      hlfir.region_assign {
+        %2 = arith.subi %c11, %arg1 : index
+        %3 = hlfir.designate %1#0 (%2)  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+        %4 = fir.load %3 : !fir.ref<i32>
+        hlfir.yield %4 : i32
+      } to {
+        %2 = hlfir.designate %1#0 (%arg1)  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
+        hlfir.yield %2 : !fir.ref<i32>
+      }
+    }
+  }
+  return
+}
+
+// The mask and rhs are saved in the same loop, the mask value is a scalar i1
+// and it can be used directly to mask the rhs (instead of loading the temp).
+
+// CHECK-LABEL:   func.func @forall_mask_and_rhs(
+// CHECK:           %[[VAL_18:.*]] = fir.allocmem !fir.array<?xi1>, %{{.*}} {bindc_name = ".tmp.forall", uniq_name = ""}
+// CHECK:           %[[VAL_20:.*]]:2 = hlfir.declare %[[VAL_18]](%{{.*}}) {uniq_name = ".tmp.forall"}
+// CHECK:           %[[VAL_29:.*]] = fir.allocmem !fir.array<?xi32>, %{{.*}} {bindc_name = ".tmp.forall", uniq_name = ""}
+// CHECK:           %[[VAL_31:.*]]:2 = hlfir.declare %[[VAL_29]](%{{.*}}) {uniq_name = ".tmp.forall"}
+// CHECK:             %[[VAL_36:.*]] = arith.cmpi sgt, %{{.*}}, %{{.*}} : i32
+// CHECK:             %[[VAL_39:.*]] = hlfir.designate %[[VAL_20]]#0 (%{{.*}})
+// CHECK:             hlfir.assign %[[VAL_36]] to %[[VAL_39]] : i1, !fir.ref<i1>
+// CHECK:             fir.if %[[VAL_36]] {
+// CHECK:               %[[VAL_45:.*]] = hlfir.designate %[[VAL_31]]#0 (%{{.*}})
+// CHECK:               hlfir.assign %{{.*}} to %[[VAL_45]] : i32, !fir.ref<i32>
+// CHECK:             }
+// CHECK:           }

flang/test/HLFIR/order_assignments/user-defined-assignment.fir

hlfir.region_assign {
     fir.call @logical_to_numeric(%arg1, %arg0) : (!fir.ref<i32>, !fir.ref<!fir.logical<4>>) -> ()
   }
   return
 }
 // CHECK-LABEL:   func.func @test_simple_scalar(
 // CHECK-SAME:    %[[VAL_0:.*]]: !fir.ref<i32>,
 // CHECK-SAME:    %[[VAL_1:.*]]: !fir.ref<!fir.logical<4>>) {
 // CHECK:  %[[VAL_2:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.logical<4>>
-// CHECK:  %[[VAL_3:.*]]:3 = hlfir.associate %[[VAL_2]] {uniq_name = ".tmp.forall"} : (!fir.logical<4>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>, i1)
+// CHECK:  %[[VAL_3:.*]]:3 = hlfir.associate %[[VAL_2]] {uniq_name = ".tmp.assign"} : (!fir.logical<4>) -> (!fir.ref<!fir.logical<4>>, !fir.ref<!fir.logical<4>>, i1)
 // CHECK:  fir.call @logical_to_numeric(%[[VAL_0]], %[[VAL_3]]#0) : (!fir.ref<i32>, !fir.ref<!fir.logical<4>>) -> ()
 // CHECK:  hlfir.end_associate %[[VAL_3]]#1, %[[VAL_3]]#2 : !fir.ref<!fir.logical<4>>, i1
 
 func.func @test_elemental_overlap(%i: !fir.ref<!fir.array<10xi32>>) {
   %c0_i32 = arith.constant 0 : i32
   %c10 = arith.constant 10 : index
   %shape  = fir.shape %c10 : (index) -> !fir.shape<1>
   hlfir.region_assign {

 // CHECK:  %[[VAL_4:.*]] = hlfir.elemental %[[VAL_3]] : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> {
 // CHECK:  ^bb0(%[[VAL_5:.*]]: index):
 // CHECK:    %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]])  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
 // CHECK:    %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref<i32>
 // CHECK:    %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_7]], %[[VAL_1]] : i32
 // CHECK:    %[[VAL_9:.*]] = fir.convert %[[VAL_8]] : (i1) -> !fir.logical<4>
 // CHECK:    hlfir.yield_element %[[VAL_9]] : !fir.logical<4>
 // CHECK:  }
-// CHECK:  %[[VAL_10:.*]]:3 = hlfir.associate %[[VAL_11:.*]](%[[VAL_3]]) {uniq_name = ".tmp.forall"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>, i1)
+// CHECK:  %[[VAL_10:.*]]:3 = hlfir.associate %[[VAL_11:.*]](%[[VAL_3]]) {uniq_name = ".tmp.assign"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>, i1)
 // CHECK:  %[[VAL_12:.*]] = arith.constant 10 : index
 // CHECK:  %[[VAL_13:.*]] = fir.shape %[[VAL_12]] : (index) -> !fir.shape<1>
 // CHECK:  %[[VAL_14:.*]] = arith.constant 1 : index
 // CHECK:  fir.do_loop %[[VAL_15:.*]] = %[[VAL_14]] to %[[VAL_12]] step %[[VAL_14]] {
 // CHECK:    %[[VAL_16:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_15]])  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
 // CHECK:    %[[VAL_17:.*]] = hlfir.designate %[[VAL_10]]#0 (%[[VAL_15]])  : (!fir.ref<!fir.array<10x!fir.logical<4>>>, index) -> !fir.ref<!fir.logical<4>>
 // CHECK:    %[[VAL_18:.*]] = fir.load %[[VAL_17]] : !fir.ref<!fir.logical<4>>
 // CHECK:    fir.call @logical_value_to_numeric(%[[VAL_16]], %[[VAL_18]]) : (!fir.ref<i32>, !fir.logical<4>) -> ()

 // CHECK:  %[[VAL_4:.*]] = hlfir.elemental %[[VAL_3]] : (!fir.shape<1>) -> !hlfir.expr<10x!fir.logical<4>> {
 // CHECK:  ^bb0(%[[VAL_5:.*]]: index):
 // CHECK:    %[[VAL_6:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_5]])  : (!fir.ref<!fir.array<10xi32>>, index) -> !fir.ref<i32>
 // CHECK:    %[[VAL_7:.*]] = fir.load %[[VAL_6]] : !fir.ref<i32>
 // CHECK:    %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_7]], %[[VAL_1]] : i32
 // CHECK:    %[[VAL_9:.*]] = fir.convert %[[VAL_8]] : (i1) -> !fir.logical<4>
 // CHECK:    hlfir.yield_element %[[VAL_9]] : !fir.logical<4>
 // CHECK:  }
-// CHECK:  %[[VAL_10:.*]]:3 = hlfir.associate %[[VAL_11:.*]](%[[VAL_3]]) {uniq_name = ".tmp.forall"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>, i1)
+// CHECK:  %[[VAL_10:.*]]:3 = hlfir.associate %[[VAL_11:.*]](%[[VAL_3]]) {uniq_name = ".tmp.assign"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>, i1)
 // CHECK:  %[[VAL_12:.*]] = hlfir.as_expr %[[VAL_10]]#0 : (!fir.ref<!fir.array<10x!fir.logical<4>>>) -> !hlfir.expr<10x!fir.logical<4>>
 // CHECK:  %[[VAL_13:.*]]:3 = hlfir.associate %[[VAL_12]](%[[VAL_3]]) {uniq_name = "adapt.valuebyref"} : (!hlfir.expr<10x!fir.logical<4>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10x!fir.logical<4>>>, !fir.ref<!fir.array<10x!fir.logical<4>>>, i1)
 // CHECK:  fir.call @logical_array_to_numeric(%[[VAL_0]], %[[VAL_13]]#0) : (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10x!fir.logical<4>>>) -> ()
 // CHECK:  hlfir.end_associate %[[VAL_13]]#1, %[[VAL_13]]#2 : !fir.ref<!fir.array<10x!fir.logical<4>>>, i1
 // CHECK:  hlfir.destroy %[[VAL_12]] : !hlfir.expr<10x!fir.logical<4>>
 // CHECK:  hlfir.end_associate %[[VAL_10]]#1, %[[VAL_10]]#2 : !fir.ref<!fir.array<10x!fir.logical<4>>>, i1
 
 func.func @test_scalar_forall_overlap(%i: !fir.ref<!fir.array<10xi32>>) {