[flang][hlfir] do not consider local temps as conflicting in assignment (#113330)

Last patch required to avoid creating a temporary for the LHS when dealing with `x([a,b]) = y`. The code dealing with "ordered assignments" (where, forall, user and vector subscripted assignments) is saving the evaluated RHS/LHS and masks if they have write effects because this write effects should not be evaluated when they affect entities that may be written to in other contexts after the evaluation and before the re-evaluation. But when dealing with write to storage allocated in the region for the expression being evluated, there is no problem to re-evaluate the write: it has no effect outside of the expression evaluation that owns the allocation. In the case of `x([a,b]) = y`, the temporary is created for the vector subscript. Raising the HLFIR abstraction for simple array constructors may be a good idea, but local temps are created in other contexts, so this fix is more generic.

Author: jeanPerier

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

@@ -347,12 +347,23 @@ conflict(llvm::ArrayRef<mlir::MemoryEffects::EffectInstance> effectsA,
  anyRAWorWAW(effectsB, effectsA, aliasAnalysis);
 }
 
-/// Could there be any write effects in "effects"?
+/// Could there be any write effects in "effects" affecting memory storages
+/// that are not local to the current region.
 static bool
-anyWrite(llvm::ArrayRef<mlir::MemoryEffects::EffectInstance> effects) {
+anyNonLocalWrite(llvm::ArrayRef<mlir::MemoryEffects::EffectInstance> effects,
+ mlir::Region &region) {
  return llvm::any_of(
- effects, [](const mlir::MemoryEffects::EffectInstance &effect) {
- return mlir::isa<mlir::MemoryEffects::Write>(effect.getEffect());
+ effects, [&region](const mlir::MemoryEffects::EffectInstance &effect) {
+ if (mlir::isa<mlir::MemoryEffects::Write>(effect.getEffect())) {
+ if (mlir::Value v = effect.getValue()) {
+ v = getStorageSource(v);
+ if (v.getDefiningOp<fir::AllocaOp>() ||
+ v.getDefiningOp<fir::AllocMemOp>())
+ return !region.isAncestor(v.getParentRegion());
+ }
+ return true;
+ }
+ return false;
  });
 }
 
@@ -393,9 +404,13 @@ void Scheduler::saveEvaluationIfConflict(mlir::Region &yieldRegion,
  if (entity && hlfir::isFortranVariableType(entity->get().getType()))
  effects.emplace_back(mlir::MemoryEffects::Read::get(), entity);
  }
- if (!leafRegionsMayOnlyRead && anyWrite(effects)) {
- // Region with write effect must be executed only once: save it the first
- // time it is encountered.
+ if (!leafRegionsMayOnlyRead && anyNonLocalWrite(effects, yieldRegion)) {
+ // Region with write effect must be executed only once (unless all writes
+ // affect storages allocated inside the region): save it the first time it
+ // is encountered.
+ LLVM_DEBUG(llvm::dbgs()
+ << "saving eval because write effect prevents re-evaluation"
+ << "\n";);
  saveEvaluation(yieldRegion, effects, /*anyWrite=*/true);
  } else if (conflict(effects, assignEffects)) {
  // Region that conflicts with the current assignments must be fully
@@ -411,7 +426,8 @@ void Scheduler::saveEvaluationIfConflict(mlir::Region &yieldRegion,
  // For example, a WHERE mask might be written by the masked assignment
  // evaluations, and it has to be saved in this case:
  // where (mask) r = f() ! function f modifies mask
- saveEvaluation(yieldRegion, effects, anyWrite(effects));
+ saveEvaluation(yieldRegion, effects,
+ anyNonLocalWrite(effects, yieldRegion));
  } else {
  // Can be executed while doing the assignment.
  independentEvaluationEffects.append(effects.begin(), effects.end());

flang/test/HLFIR/order_assignments/vector-subscripts-scheduling.fir

@@ -0,0 +1,31 @@
+// Test local alloca and store inside hlfir.region_assign do not trigger the
+// creation of a temporary for the LHS.
+
+// RUN: fir-opt -o - -lower-hlfir-ordered-assignments --debug-only=flang-ordered-assignment -flang-dbg-order-assignment-schedule-only %s 2>&1 | FileCheck %s
+// REQUIRES: asserts
+
+func.func @simple(%arg0: !fir.ref<!fir.array<100x10xf32>> , %arg1: !fir.ref<!fir.array<10xi64>> , %arg2: !fir.ref<!fir.array<10xf32>>, %i: i64, %f: f32) {
+ %c10 = arith.constant 10 : index
+ %c100 = arith.constant 100 : index
+ %0 = fir.shape %c100, %c10 : (index, index) -> !fir.shape<2>
+ %1:2 = hlfir.declare %arg0(%0) {uniq_name = "_QFsimpleEx"} : (!fir.ref<!fir.array<100x10xf32>>, !fir.shape<2>) -> (!fir.ref<!fir.array<100x10xf32>>, !fir.ref<!fir.array<100x10xf32>>)
+ %2 = fir.shape %c10 : (index) -> !fir.shape<1>
+ %3:2 = hlfir.declare %arg1(%2) {uniq_name = "y"} : (!fir.ref<!fir.array<10xi64>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi64>>, !fir.ref<!fir.array<10xi64>>)
+ hlfir.region_assign {
+ hlfir.yield %f : f32
+ } to {
+ %local_temp = fir.alloca i64
+ fir.store %i to %local_temp : !fir.ref<i64>
+ %icopy = fir.load %local_temp : !fir.ref<i64>
+ hlfir.elemental_addr %2 : !fir.shape<1> {
+ ^bb0(%arg3: index):
+ %5 = hlfir.designate %3#0 (%arg3) : (!fir.ref<!fir.array<10xi64>>, index) -> !fir.ref<i64>
+ %6 = fir.load %5 : !fir.ref<i64>
+ %7 = hlfir.designate %1#0 (%icopy, %6) : (!fir.ref<!fir.array<100x10xf32>>, i64, i64) -> !fir.ref<f32>
+ hlfir.yield %7 : !fir.ref<f32>
+ }
+ }
+ return
+}
+
+// CHECK: run 1 evaluate: region_assign

flang/test/HLFIR/order_assignments/where-scheduling.f90

@@ -135,6 +135,7 @@ end function f
 !CHECK-NEXT: run 2 evaluate: where/region_assign1
 !CHECK-LABEL: ------------ scheduling where in _QPonly_once ------------
 !CHECK-NEXT: unknown effect: %{{[0-9]+}} = llvm.intr.stacksave : !llvm.ptr
+!CHECK-NEXT: saving eval because write effect prevents re-evaluation
 !CHECK-NEXT: run 1 save (w): where/mask
 !CHECK-NEXT: run 2 evaluate: where/region_assign1
 !CHECK-NEXT: run 3 evaluate: where/region_assign2
@@ -180,6 +181,7 @@ end function f
 !CHECK-NEXT: conflict: R/W: %{{.*}} = hlfir.declare %{{.*}} {uniq_name = "_QFwhere_construct_unknown_conflictEmask"} : (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.dscope) -> (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.box<!fir.array<?x!fir.logical<4>>>) W:<unknown>
 !CHECK-NEXT: run 1 save : where/mask
 !CHECK-NEXT: unknown effect: %{{.*}} = fir.call @_QPf() fastmath<contract> : () -> f32
+!CHECK-NEXT: saving eval because write effect prevents re-evaluation
 !CHECK-NEXT: run 2 save (w): where/region_assign1/rhs
 !CHECK-NEXT: run 3 evaluate: where/region_assign1
 !CHECK-NEXT: ------------ scheduling where in _QPelsewhere_construct_unknown_conflict ------------
@@ -190,5 +192,6 @@ end function f
 !CHECK-NEXT: conflict: R/W: %{{.*}} = hlfir.declare %{{.*}} {uniq_name = "_QFelsewhere_construct_unknown_conflictEmask2"} : (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.dscope) -> (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.box<!fir.array<?x!fir.logical<4>>>) W:<unknown>
 !CHECK-NEXT: run 2 save : where/elsewhere1/mask
 !CHECK-NEXT: unknown effect: %{{.*}} = fir.call @_QPf() fastmath<contract> : () -> f32
+!CHECK-NEXT: saving eval because write effect prevents re-evaluation
 !CHECK-NEXT: run 3 save (w): where/elsewhere1/region_assign1/rhs
 !CHECK-NEXT: run 4 evaluate: where/elsewhere1/region_assign1