[flang-commits] [flang] e6238ab - [flang] Deallocate WHERE masks after all assignments

[Jean Perier via flang-commits]

Author: Jean Perier
Date: 2022-09-20T10:40:30+02:00
New Revision: e6238ab52590a4f31bc8c6e806c7947e9ef04b57

URL: https://github.com/llvm/llvm-project/commit/e6238ab52590a4f31bc8c6e806c7947e9ef04b57
DIFF: https://github.com/llvm/llvm-project/commit/e6238ab52590a4f31bc8c6e806c7947e9ef04b57.diff

LOG: [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

Added: 
    flang/test/Lower/where-allocatable-assignments.f90

Modified: 
    flang/lib/Lower/Bridge.cpp

Removed: 
    


################################################################################
diff  --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index a8f020ab04f7..a6b8caa10ffc 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -2188,9 +2188,15 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   /// 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 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     // 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 000000000000..ce68acc9b129
--- /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