[flang-commits] [flang] 698b42c - [flang][stack-arrays] Collect analysis results for OMP ws loops (#103590)

[via flang-commits]

Author: Kareem Ergawy
Date: 2024-08-16T09:27:13+02:00
New Revision: 698b42ccff69fde6509c6ad0baf262c257c039bb

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

LOG: [flang][stack-arrays] Collect analysis results for OMP ws loops (#103590)

We missed collecting the analysis results for regions terminated with
`omp.yield`. This result in missing some opportunities for malloc
optimizations inside omp regions.

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/StackArrays.cpp
    flang/test/Transforms/stack-arrays-hlfir.f90

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index bdc2d9cd9c6c43..a5a95138ac1281 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -457,10 +457,12 @@ StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
     if (lattice)
       (void)point.join(*lattice);
   };
+
   func->walk([&](mlir::func::ReturnOp child) { joinOperationLattice(child); });
   func->walk([&](fir::UnreachableOp child) { joinOperationLattice(child); });
   func->walk(
       [&](mlir::omp::TerminatorOp child) { joinOperationLattice(child); });
+  func->walk([&](mlir::omp::YieldOp child) { joinOperationLattice(child); });
 
   llvm::DenseSet<mlir::Value> freedValues;
   point.appendFreedValues(freedValues);

diff  --git a/flang/test/Transforms/stack-arrays-hlfir.f90 b/flang/test/Transforms/stack-arrays-hlfir.f90
index 50261b3078466d..296edbaf35417e 100644
--- a/flang/test/Transforms/stack-arrays-hlfir.f90
+++ b/flang/test/Transforms/stack-arrays-hlfir.f90
@@ -53,3 +53,31 @@ end subroutine omp_temp_array
 ! CHECK-NEXT:    }
 ! CHECK:         return
 ! CHECK-NEXT:  }
+
+subroutine omp_target_wsloop
+  implicit none
+  integer (8) :: lV, i
+  integer (8), dimension (2) :: iaVS
+
+  lV = 202
+
+  !$omp target teams distribute
+  do i = 1, 10
+    iaVS = [lV, lV]
+  end do
+  !$omp end target teams distribute
+end subroutine omp_target_wsloop
+! CHECK-LABEL: func.func @_QPomp_target_wsloop{{.*}} {
+! CHECK:         omp.target {{.*}} {
+! CHECK-NOT:       fir.allocmem
+! CHECK-NOT:       fir.freemem
+! CHECK:           fir.alloca !fir.array<2xi64>
+! CHECK:         omp.teams {
+! CHECK:         omp.distribute {
+! CHECK:         omp.loop_nest {{.*}} {
+! CHECK-NOT:       fir.allocmem
+! CHECK-NOT:       fir.freemem
+! CHECK:           omp.yield
+! CHECK-NEXT:    }
+! CHECK:         return
+! CHECK-NEXT:  }