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

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Kareem Ergawy (ergawy)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/103590.diff


2 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/StackArrays.cpp (+2) 
- (modified) flang/test/Transforms/stack-arrays-hlfir.f90 (+28) 


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

``````````

</details>


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