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

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
```

Reviewed By: tblah

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

Author: vzakhari

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

@@ -81,7 +81,8 @@ class Scheduler {
   /// 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
@@ -117,6 +118,10 @@ class Scheduler {
 
   /// 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>
@@ -260,6 +265,20 @@ static void gatherAssignEffects(
   }
 }
 
+/// 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.
 //===----------------------------------------------------------------------===//
@@ -344,11 +363,19 @@ anyWrite(llvm::ArrayRef<mlir::MemoryEffects::EffectInstance> effects) {
 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.
@@ -377,6 +404,14 @@ void Scheduler::saveEvaluationIfConflict(mlir::Region &yieldRegion,
     // 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());
@@ -444,6 +479,7 @@ void Scheduler::finishSchedulingAssignment(hlfir::RegionAssignOp assign) {
   schedule.back().memoryEffects.append(assignEffects.begin(),
                                        assignEffects.end());
   assignEffects.clear();
+  assignEvaluateEffects.clear();
   parentEvaluationEffects.clear();
   independentEvaluationEffects.clear();
   savedAnyRegionForCurrentAssignment = false;
@@ -533,9 +569,13 @@ hlfir::buildEvaluationSchedule(hlfir::OrderedAssignmentTreeOpInterface root,
       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.

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

@@ -89,6 +89,7 @@ pure real function 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
@@ -107,6 +108,7 @@ pure real function 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

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

@@ -103,6 +103,30 @@ subroutine internal
 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 ------------
@@ -115,9 +139,12 @@ subroutine internal
 !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
@@ -151,3 +178,20 @@ subroutine internal
 !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