[flang][stack-arrays] Collect analysis results for OMP ws loops (llvm…

…#103590) (#145)

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

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);

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