[flang] Deallocate WHERE masks after all assignments

Allocatable assignments were triggering lowering to clean-up
any WHERE mask temporaries, even if some assignments where left
in the WHERE construct.

This is because allocatable assignments lowering was being passed the
wrong statement context. Fix this by selecting the where/forall statement
context instead of a local one when there is one.

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

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index a8f020a..a6b8caa 100644 (file)
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -2188,9 +2188,15 @@ private:
   /// pointer assignment.)
   void genArrayAssignment(
       const Fortran::evaluate::Assignment &assign,
-      Fortran::lower::StatementContext &stmtCtx,
+      Fortran::lower::StatementContext &localStmtCtx,
       llvm::Optional<llvm::SmallVector<mlir::Value>> lbounds = llvm::None,
       llvm::Optional<llvm::SmallVector<mlir::Value>> ubounds = llvm::None) {
+
+    Fortran::lower::StatementContext &stmtCtx =
+        explicitIterationSpace()
+            ? explicitIterSpace.stmtContext()
+            : (implicitIterationSpace() ? implicitIterSpace.stmtContext()
+                                        : localStmtCtx);
     if (Fortran::lower::isWholeAllocatable(assign.lhs)) {
       // Assignment to allocatables may require the lhs to be
       // deallocated/reallocated. See Fortran 2018 10.2.1.3 p3
@@ -2230,9 +2236,7 @@ private:
     // implied by the lhs array expression.
     Fortran::lower::createAnyMaskedArrayAssignment(
         *this, assign.lhs, assign.rhs, explicitIterSpace, implicitIterSpace,
-        localSymbols,
-        explicitIterationSpace() ? explicitIterSpace.stmtContext()
-                                 : implicitIterSpace.stmtContext());
+        localSymbols, stmtCtx);
   }
 
 #if !defined(NDEBUG)

diff --git a/flang/test/Lower/where-allocatable-assignments.f90 b/flang/test/Lower/where-allocatable-assignments.f90
new file mode 100644 (file)
index 0000000..ce68acc
--- /dev/null
+++ b/flang/test/Lower/where-allocatable-assignments.f90
@@ -0,0 +1,94 @@
+! Test that WHERE mask clean-up occurs at the right time when the
+! WHERE contains whole allocatable assignments.
+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+
+module mtest
+contains
+
+! CHECK-LABEL: func.func @_QMmtestPfoo(
+! CHECK-SAME:       %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "a"},
+! CHECK-SAME:       %[[VAL_1:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {fir.bindc_name = "b"}) {
+subroutine foo(a, b)
+  integer :: a(:)
+  integer, allocatable :: b(:)
+! CHECK:  %[[VAL_2:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+! CHECK:  %[[VAL_3:.*]] = arith.constant 0 : index
+! CHECK:  %[[VAL_4:.*]]:3 = fir.box_dims %[[VAL_2]], %[[VAL_3]] : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> (index, index, index)
+          ! WHERE mask temp allocation
+! CHECK:  %[[VAL_9:.*]] = fir.allocmem !fir.array<?x!fir.logical<4>>, %[[VAL_4]]#1 {uniq_name = ".array.expr"}
+! CHECK:  %[[VAL_15:.*]] = fir.do_loop {{.*}} {
+!           ! WHERE mask element computation
+! CHECK:  }
+! CHECK:  fir.array_merge_store %{{.*}}, %[[VAL_15]] to %[[VAL_9]] : !fir.array<?x!fir.logical<4>>, !fir.array<?x!fir.logical<4>>, !fir.heap<!fir.array<?x!fir.logical<4>>>
+
+          ! First assignment to a whole allocatable (in WHERE)
+! CHECK:  fir.if {{.*}} {
+! CHECK:    fir.if {{.*}} {
+            ! assignment into new storage (`b` allocated with bad shape)
+! CHECK:      fir.allocmem
+! CHECK:      fir.do_loop {{.*}} {
+! CHECK:        fir.array_coor %[[VAL_9]]
+! CHECK:        fir.if %{{.*}} {
+                  ! WHERE
+! CHECK:          fir.array_update {{.*}}
+! CHECK:        } else {
+! CHECK:        }
+! CHECK:      }
+! CHECK:    } else {
+              ! assignment into old storage (`b` allocated with the same shape)
+! CHECK:      fir.do_loop {{.*}} {
+! CHECK:        fir.array_coor %[[VAL_9]]
+! CHECK:        fir.if %{{.*}} {
+                  ! WHERE
+! CHECK:          fir.array_update {{.*}}
+! CHECK:        } else {
+! CHECK:        }
+! CHECK:      }
+! CHECK:    }
+! CHECK:  } else {
+            ! assignment into new storage (`b` unallocated)
+! CHECK:    fir.allocmem
+! CHECK:    fir.do_loop %{{.*}} {
+! CHECK:      fir.array_coor %[[VAL_9]]
+! CHECK:      fir.if %{{.*}} {
+                ! WHERE
+! CHECK:        fir.array_update {{.*}}
+! CHECK:      } else {
+! CHECK:      }
+! CHECK:    }
+! CHECK:  }
+! CHECK:  fir.if {{.*}} {
+! CHECK:    fir.if {{.*}} {
+              ! deallocation of `b` old allocatable data store
+! CHECK:    }
+            ! update of `b` descriptor
+! CHECK:  }
+          ! Second assignment (in ELSEWHERE)
+! CHECK:  fir.do_loop {{.*}} {
+! CHECK:    fir.array_coor %[[VAL_9]]{{.*}} : (!fir.heap<!fir.array<?x!fir.logical<4>>>, !fir.shape<1>, index) -> !fir.ref<!fir.logical<4>>
+! CHECK:    fir.if {{.*}} {
+! CHECK:    } else {
+              ! elsewhere
+! CHECK:      fir.array_update
+! CHECK:    }
+! CHECK:  }
+          ! WHERE temp clean-up
+! CHECK:  fir.freemem %[[VAL_9]] : !fir.heap<!fir.array<?x!fir.logical<4>>>
+! CHECK-NEXT:  return
+  where (b > 0)
+    b = a
+  elsewhere
+    b(:) = 0
+  end where
+end
+end module
+
+  use mtest
+  integer, allocatable :: a(:), b(:)
+  allocate(a(10),b(10))
+  a = 5
+  b = 1
+  call foo(a, b)
+  print*, b
+  deallocate(a,b)
+end