From e6238ab52590a4f31bc8c6e806c7947e9ef04b57 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier@nvidia.com>
Date: Tue, 20 Sep 2022 10:39:39 +0200
Subject: [PATCH] [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
---
 flang/lib/Lower/Bridge.cpp                         | 12 ++-
 flang/test/Lower/where-allocatable-assignments.f90 | 94 ++++++++++++++++++++++
 2 files changed, 102 insertions(+), 4 deletions(-)
 create mode 100644 flang/test/Lower/where-allocatable-assignments.f90

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index a8f020a..a6b8caa 100644
--- 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
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
-- 
2.7.4