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

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.

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

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

@@ -85,6 +85,10 @@ public:
   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;
@@ -109,6 +113,7 @@ public:
     return copy.getBase();
   }
   void destroy(mlir::Location loc, fir::FirOpBuilder &builder);
+  bool canBeFetchedAfterPush() const { return true; }
 
 public:
   /// Temporary storage for the copy.
@@ -130,6 +135,7 @@ public:
   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
@@ -165,6 +171,12 @@ public:
   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;

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

@@ -979,6 +979,17 @@ computeLoopNestIterationNumber(mlir::Location loc, fir::FirOpBuilder &builder,
   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;
@@ -996,18 +1007,18 @@ void OrderedAssignmentRewriter::generateSaveEntity(
     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);
@@ -1042,8 +1053,13 @@ void OrderedAssignmentRewriter::generateSaveEntity(
   }
 
   // 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");

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

@@ -36,7 +36,7 @@ func.func @test_elsewhere_impure_mask(%x: !fir.ref<!fir.array<10xi32>>, %y: !fir
 // 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 {

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

@@ -237,7 +237,7 @@ func.func @test_where_mask_save(%arg0: !fir.box<!fir.array<?xi32>>) {
 // 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]] {
@@ -292,9 +292,9 @@ func.func @test_where_rhs_save(%x: !fir.ref<!fir.array<10xi32>>, %mask: !fir.ref
 // 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>>

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

@@ -0,0 +1,104 @@
+// 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

@@ -16,7 +16,7 @@ func.func @test_simple_scalar(%i: !fir.ref<i32>, %l: !fir.ref<!fir.logical<4>>)
 // 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
 
@@ -54,7 +54,7 @@ func.func @test_elemental_overlap(%i: !fir.ref<!fir.array<10xi32>>) {
 // 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
@@ -102,7 +102,7 @@ func.func @test_array_overlap(%i: !fir.ref<!fir.array<10xi32>>) {
 // 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>>>) -> ()