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

[Kareem Ergawy via flang-commits]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/103590

>From 032ef439d452d60f4021f741ee6f4a5de05113e4 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Tue, 13 Aug 2024 23:32:00 -0500
Subject: [PATCH] [flang][stack-arrays] Collect analysis results for OMP ws
 loops

We missed collecting the analysis results for regions terminated with
`omp.yield`. This result in missing some opportunities for malloc
optimizations inside omp regions.
---
 .../lib/Optimizer/Transforms/StackArrays.cpp  |  2 ++
 flang/test/Transforms/stack-arrays-hlfir.f90  | 28 +++++++++++++++++++
 2 files changed, 30 insertions(+)

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:  }