[flang][hlfir] Fixed where/elsewhere mask saving in case of conflicts.

[vzakhari]

Summary

The assignments inside where/elsewhere may affect variables participating
in the mask expression, but execution of the assignments must not affect
the established control mask(s) (F'18 10.2.3.2 p. 13).

The following example must print all 42's:

program test
  integer c(3)
  logical :: mask(3) = .true.
  where (mask)
     c = f()
  end where
  print *, c
contains
  integer function f()
    mask = .false.
    f = 42
  end function f
end program test

[tblah]

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

572

I think this applies to FORALL too. See 10.2.4.2.1: the evaluation of the mask-expr happens before execution of the FORALL body.

But it isn't spelt out as clearly as for WHERE/ELSEWHERE

[vzakhari]

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

572

At the same time, FORALL has a limitation that only pure function references can appear in the body. And I do not know if an evaluation can modify the mask-expr variables' values without a function reference (note that the assignment effects have been already taken into account in the code). I would like to keep it like this for the time being.

[tblah]

LGTM, thanks for this

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

public:
   /// the expression has write effects (which is possible outside of forall),
   /// create an action in the schedule to save the value in the schedule before
   /// evaluating the current assignment. For expression with write effect,
   /// saving them ensures they are evaluated only once. A region whose value
   /// was saved in a previous run is considered to have no side effects with the
   /// current assignment: the saved value will be used.
   void saveEvaluationIfConflict(mlir::Region &yieldRegion,
                                 bool leafRegionsMayOnlyRead,
-                                bool yieldIsImplicitRead = true);
+                                bool yieldIsImplicitRead = true,
+                                bool evaluationsMayConflict = false);
 
   /// Finish evaluating a group of independent regions. The current independent
   /// regions effects are added to the "parent" effect list since evaluating the
   /// next analyzed region would require evaluating the current independent
   /// regions.
   void finishIndependentEvaluationGroup() {
     parentEvaluationEffects.append(independentEvaluationEffects.begin(),
                                    independentEvaluationEffects.end());
private:
 
   /// Can the current assignment be schedule with the previous run. This is
   /// only possible if the assignment and all of its dependencies have no side
   /// effects conflicting with the previous run.
   bool canFuseAssignmentWithPreviousRun();
 
   /// Memory effects of the assignments being lowered.
   llvm::SmallVector<mlir::MemoryEffects::EffectInstance> assignEffects;
+  /// Memory effects of the evaluations implied by the assignments
+  /// being lowered. They do not include the implicit writes
+  /// to the LHS of the assignments.
+  llvm::SmallVector<mlir::MemoryEffects::EffectInstance> assignEvaluateEffects;
   /// Memory effects of the unsaved evaluation region that are controlling or
   /// masking the current assignments.
   llvm::SmallVector<mlir::MemoryEffects::EffectInstance>
       parentEvaluationEffects;
   /// Same as parentEvaluationEffects, but for the current "leaf group" being
   /// analyzed scheduled.
   llvm::SmallVector<mlir::MemoryEffects::EffectInstance>
       independentEvaluationEffects;
if (!regionAssign.getUserDefinedAssignment().empty()) {
     // directly.
     if (userDefAssignmentMayOnlyWriteToAssignedVariable)
       assignEffects.emplace_back(mlir::MemoryEffects::Read::get());
     else
       assignEffects.emplace_back(mlir::MemoryEffects::Write::get());
   }
 }
 
+/// Gather the effects of evaluations implied by the given assignment.
+/// These are the effects of operations from LHS and RHS.
+static void gatherAssignEvaluationEffects(
+    hlfir::RegionAssignOp regionAssign,
+    bool userDefAssignmentMayOnlyWriteToAssignedVariable,
+    llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &assignEffects) {
+  gatherMemoryEffects(regionAssign.getLhsRegion(),
+                      userDefAssignmentMayOnlyWriteToAssignedVariable,
+                      assignEffects);
+  gatherMemoryEffects(regionAssign.getRhsRegion(),
+                      userDefAssignmentMayOnlyWriteToAssignedVariable,
+                      assignEffects);
+}
+
 //===----------------------------------------------------------------------===//
 // Scheduling Implementation : finding conflicting memory effects.
 //===----------------------------------------------------------------------===//
 
 /// Follow addressing and declare like operation to the storage source.
 /// This allows using FIR alias analysis that otherwise does not know
 /// about those operations. This is correct, but ignoring the designate
 /// and declare info may yield false positive regarding aliasing (e.g,

 
 //===----------------------------------------------------------------------===//
 // Scheduling Implementation : Scheduler class implementation
 //===----------------------------------------------------------------------===//
 
 void Scheduler::startSchedulingAssignment(hlfir::RegionAssignOp assign,
                                           bool leafRegionsMayOnlyRead) {
   gatherAssignEffects(assign, leafRegionsMayOnlyRead, assignEffects);
+  // Unconditionally collect effects of the evaluations of LHS and RHS
+  // in case they need to be analyzed for any parent that might be
+  // affected by conflicts of these evaluations.
+  // This collection migth be skipped, if there are no such parents,
+  // but for the time being we run it always.
+  gatherAssignEvaluationEffects(assign, leafRegionsMayOnlyRead,
+                                assignEvaluateEffects);
 }
 
 void Scheduler::saveEvaluationIfConflict(mlir::Region &yieldRegion,
                                          bool leafRegionsMayOnlyRead,
-                                         bool yieldIsImplicitRead) {
+                                         bool yieldIsImplicitRead,
+                                         bool evaluationsMayConflict) {
   // If the region evaluation was previously executed and saved, the saved
   // value will be used when evaluating the current assignment and this has
   // no effects in the current assignment evaluation.
   if (savedRegions.contains(&yieldRegion))
     return;
   llvm::SmallVector<mlir::MemoryEffects::EffectInstance> effects;
   gatherMemoryEffects(yieldRegion, leafRegionsMayOnlyRead, effects);
   // Yield has no effect as such, but in the context of order assignments.
if (!leafRegionsMayOnlyRead && anyWrite(effects)) {
     saveEvaluation(yieldRegion, effects, /*anyWrite=*/true);
   } else if (conflict(effects, assignEffects)) {
     // Region that conflicts with the current assignments must be fully
     // evaluated and saved before doing the assignment (Note that it may
     // have already have been evaluated without saving it before, but this
     // implies that it never conflicted with a prior assignment, so its value
     // should be the same.)
     saveEvaluation(yieldRegion, effects, /*anyWrite=*/false);
+  } else if (evaluationsMayConflict &&
+             conflict(effects, assignEvaluateEffects)) {
+    // If evaluations of the assignment may conflict with the yield
+    // evaluations, we have to save yield evaluation.
+    // 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));
   } else {
     // Can be executed while doing the assignment.
     independentEvaluationEffects.append(effects.begin(), effects.end());
   }
 }
 
 void Scheduler::saveEvaluation(
     mlir::Region &yieldRegion,
void Scheduler::finishSchedulingAssignment(hlfir::RegionAssignOp assign) {
   // parentEvaluationEffects that already in the previous run effects (since
   // assignments may share the same parents), otherwise, this can make the
   // conflict() calls more and more expensive.
   schedule.back().memoryEffects.append(parentEvaluationEffects.begin(),
                                        parentEvaluationEffects.end());
   schedule.back().memoryEffects.append(assignEffects.begin(),
                                        assignEffects.end());
   assignEffects.clear();
+  assignEvaluateEffects.clear();
   parentEvaluationEffects.clear();
   independentEvaluationEffects.clear();
   savedAnyRegionForCurrentAssignment = false;
   LLVM_DEBUG(logAssignmentEvaluation(llvm::dbgs(), schedule.size(), assign));
 }
 
 //===----------------------------------------------------------------------===//
 // Scheduling Implementation : driving the Scheduler in the assignment tree.
for (hlfir::RegionAssignOp assign : assignments) {
     // region that must be saved is saved in order.
     llvm::SmallVector<hlfir::OrderedAssignmentTreeOpInterface> parents;
     gatherAssignmentParents(assign, parents);
     for (hlfir::OrderedAssignmentTreeOpInterface parent :
          llvm::reverse(parents)) {
       scheduler.startIndependentEvaluationGroup();
       llvm::SmallVector<mlir::Region *, 4> yieldRegions;
       parent.getLeafRegions(yieldRegions);
+      // TODO: is this really limited to WHERE/ELSEWHERE?
+      bool evaluationsMayConflict = mlir::isa<hlfir::WhereOp>(parent) ||
+                                    mlir::isa<hlfir::ElseWhereOp>(parent);
       for (mlir::Region *yieldRegion : yieldRegions)
-        scheduler.saveEvaluationIfConflict(*yieldRegion,
-                                           leafRegionsMayOnlyRead);
+        scheduler.saveEvaluationIfConflict(*yieldRegion, leafRegionsMayOnlyRead,
+                                           /*yieldIsImplicitRead=*/true,
+                                           evaluationsMayConflict);
       scheduler.finishIndependentEvaluationGroup();
     }
     // Look for conflicts between the RHS/LHS evaluation and the assignments.
     // The LHS yield has no implicit read effect on the produced variable (the
     // variable is not read before the assignment).
     scheduler.startIndependentEvaluationGroup();
     scheduler.saveEvaluationIfConflict(assign.getRhsRegion(),
                                        leafRegionsMayOnlyRead);

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

pure real function foo(x, i)
       real, intent(in) :: x(10)
     end function
   end interface
   real :: x(10)
   forall(i=1:10) x(i) = foo(x, i)
 end subroutine
 !CHECK-LABEL: ------------ scheduling forall in _QPunknown_function_call ------------
 !CHECK-NEXT: unknown effect: {{.*}} fir.call @_QPfoo
+!CHECK-NEXT: unknown effect: {{.*}} fir.call @_QPfoo
 !CHECK-NEXT: conflict: R/W: <unknown> W:<block argument> of type '!fir.ref<!fir.array<10xf32>>' at index: 0
 !CHECK-NEXT: run 1 save    : forall/region_assign1/rhs
 !CHECK-NEXT: run 2 evaluate: forall/region_assign1
 
 subroutine unknown_function_call2(x)
   interface
     pure real function foo2(i)
       integer, value :: i
     end function
   end interface
   ! foo2 may read x since it is a target, even if it is pure,
   ! if the actual argument of x is a module variable accessible
   ! to foo via host association.
   real, target :: x(:)
   forall(i=1:10) x(i) = foo2(i)
 end subroutine
 !CHECK-LABEL: ------------ scheduling forall in _QPunknown_function_call2 ------------
 !CHECK-NEXT: unknown effect: {{.*}} fir.call @_QPfoo2(
+!CHECK-NEXT: unknown effect: {{.*}} fir.call @_QPfoo2(
 !CHECK-NEXT: conflict: R/W: <unknown> W:<block argument> of type '!fir.box<!fir.array<?xf32>>' at index: 0
 !CHECK-NEXT: run 1 save    : forall/region_assign1/rhs
 !CHECK-NEXT: run 2 evaluate: forall/region_assign1
 
 subroutine forall_mask_conflict(x)
   integer :: x(:)
   forall(i=1:10, x(11-i)>0) x(i) = 42
 end subroutine

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

subroutine internal
   ! are being assigned. But this will require more work in the
   ! conflict analysis to prove that the lowered DAG of `x(:, y(i, j))`
   ! are the same and that the access to this designator is done in the
   ! same ordered inside the mask and LHS.
   where (x(:, y(i, j)) == y(i, j))  x(:, y(i, j)) = 42
 end subroutine
 end subroutine
 
+subroutine where_construct_unknown_conflict(x, mask)
+  real :: x(:)
+  logical :: mask(:)
+  interface
+     real function f()
+     end function f
+  end interface
+  where (mask) x = f()
+end subroutine
+
+subroutine elsewhere_construct_unknown_conflict(x, y, mask1, mask2)
+  real :: x(:), y(:)
+  logical :: mask1(:), mask2(:)
+  interface
+     real function f()
+     end function f
+  end interface
+  where (mask1)
+     x = 1.0
+  elsewhere (mask2)
+     y = f()
+  end where
+end subroutine
+
 !CHECK-LABEL: ------------ scheduling where in _QPno_conflict ------------
 !CHECK-NEXT: run 1 evaluate: where/region_assign1
 !CHECK-LABEL: ------------ scheduling where in _QPfake_conflict ------------
 !CHECK-NEXT: conflict: R/W: <block argument> of type '!fir.box<!fir.array<?xf32>>' at index: 0 W:<block argument> of type '!fir.box<!fir.array<?xf32>>' at index: 0
 !CHECK-NEXT: run 1 save    : where/mask
 !CHECK-NEXT: run 2 evaluate: where/region_assign1
 !CHECK-LABEL: ------------ scheduling where in _QPonly_once ------------
 !CHECK-NEXT: unknown effect: %9 = fir.call @llvm.stacksave() fastmath<contract> : () -> !fir.ref<i8>
 !CHECK-NEXT: run 1 save  (w): where/mask
 !CHECK-NEXT: run 2 evaluate: where/region_assign1
 !CHECK-NEXT: run 3 evaluate: where/region_assign2
 !CHECK-LABEL: ------------ scheduling where in _QPrhs_lhs_conflict ------------
-!CHECK-NEXT: unknown effect: %2 = hlfir.transpose %0#0 : (!fir.box<!fir.array<?x?xf32>>) -> !hlfir.expr<?x?xf32>
-!CHECK-NEXT: run 1 save  (w): where/region_assign1/rhs
-!CHECK-NEXT: run 2 evaluate: where/region_assign1
+!CHECK-NEXT: unknown effect: %{{.*}} = hlfir.transpose %{{.*}} : (!fir.box<!fir.array<?x?xf32>>) -> !hlfir.expr<?x?xf32>
+!CHECK-NEXT: conflict: R/W: %6 = hlfir.designate %{{.*}} (%{{.*}}, %{{.*}})  : (!fir.box<!fir.array<?x?xf32>>, index, index) -> !fir.ref<f32> W:<unknown>
+!CHECK-NEXT: run 1 save    : where/mask
+!CHECK-NEXT: unknown effect: %{{.*}} = hlfir.transpose %{{.*}} : (!fir.box<!fir.array<?x?xf32>>) -> !hlfir.expr<?x?xf32>
+!CHECK-NEXT: run 2 save  (w): where/region_assign1/rhs
+!CHECK-NEXT: run 3 evaluate: where/region_assign1
 !CHECK-LABEL: ------------ scheduling where in _QPwhere_construct_no_conflict ------------
 !CHECK-NEXT: run 1 evaluate: where/region_assign1
 !CHECK-NEXT: run 2 evaluate: where/elsewhere1/region_assign1
 !CHECK-LABEL: ------------ scheduling where in _QPwhere_construct_conflict ------------
 !CHECK-NEXT: run 1 evaluate: where/region_assign1
 !CHECK-NEXT: conflict: R/W: <block argument> of type '!fir.box<!fir.array<?x?xf32>>' at index: 1 W:<block argument> of type '!fir.box<!fir.array<?x?xf32>>' at index: 1
 !CHECK-NEXT: run 2 save    : where/mask
 !CHECK-NEXT: run 3 evaluate: where/elsewhere1/region_assign1

 !CHECK-NEXT: run 1 save    : forall/where1/mask
 !CHECK-NEXT: conflict: R/W: <block argument> of type '!fir.box<!fir.array<?x?xf32>>' at index: 0 W:<block argument> of type '!fir.box<!fir.array<?x?xf32>>' at index: 0
 !CHECK-NEXT: run 1 save    : forall/where1/region_assign1/rhs
 !CHECK-NEXT: run 2 evaluate: forall/where1/region_assign1
 !CHECK-LABEL: ------------ scheduling where in _QFno_need_to_make_lhs_tempPinternal ------------
 !CHECK-NEXT: conflict: R/W: %7 = fir.load %6 : !fir.llvm_ptr<!fir.ref<i32>> W:%13 = fir.load %12 : !fir.ref<!fir.box<!fir.array<?x?xi32>>>
 !CHECK-NEXT: run 1 save    : where/mask
 !CHECK-NEXT: run 2 evaluate: where/region_assign1
+!CHECK-NEXT: ------------ scheduling where in _QPwhere_construct_unknown_conflict ------------
+!CHECK-NEXT: unknown effect: %{{.*}} = fir.call @_QPf() fastmath<contract> : () -> f32
+!CHECK-NEXT: conflict: R/W: %{{.*}} = hlfir.declare %{{.*}} {uniq_name = "_QFwhere_construct_unknown_conflictEmask"} : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> (!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: 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 ------------
+!CHECK-NEXT: run 1 evaluate: where/region_assign1
+!CHECK-NEXT: unknown effect: %{{.*}} = fir.call @_QPf() fastmath<contract> : () -> f32
+!CHECK-NEXT: conflict: R/W: %{{.*}} = hlfir.declare %{{.*}} {uniq_name = "_QFelsewhere_construct_unknown_conflictEmask1"} : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> (!fir.box<!fir.array<?x!fir.logical<4>>>, !fir.box<!fir.array<?x!fir.logical<4>>>) W:<unknown>
+!CHECK-NEXT: run 2 save    : where/mask
+!CHECK-NEXT: conflict: R/W: %{{.*}} = hlfir.declare %{{.*}} {uniq_name = "_QFelsewhere_construct_unknown_conflictEmask2"} : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> (!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: run 3 save  (w): where/elsewhere1/region_assign1/rhs
+!CHECK-NEXT: run 4 evaluate: where/elsewhere1/region_assign1