[flang][stack-arrays] Extend pass to work on declare ops and within omp regions #98810

ergawy · 2024-07-14T12:05:24Z

Extends the stack-arrays pass to support fir.declare ops. Before that, we did not recognize malloc-free pairs for which fir.declare is used to declare the allocated entity. This is because the free op was invoked on the result of the fir.declare op and did not directly use the allocated memory SSA value.

This also extends the pass to collect the analysis results within OpenMP regions.

llvmbot · 2024-07-14T12:05:55Z

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

Author: Kareem Ergawy (ergawy)

Changes

…mp regions

Extends the stack-arrays pass to support fir.declare ops. Before that, we did not recognize malloc-free pairs for which fir.declare is used to declare the allocated entity. This is because the free op was invoked on the result of the fir.declare op and did not directly use the allocated memory SSA value.

This also extends the pass to collect the analysis results within OpenMP regions.

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

3 Files Affected:

(modified) flang/lib/Optimizer/Transforms/StackArrays.cpp (+21-2)
(added) flang/test/Transforms/stack-arrays-hlfir.f90 (+55)
(modified) flang/test/Transforms/stack-arrays.fir (+2-5)

diff --git a/flang/lib/Optimizer/Transforms/StackArrays.cpp b/flang/lib/Optimizer/Transforms/StackArrays.cpp
index e8fa70ebc39d8..f695a7fa2f8fc 100644
--- a/flang/lib/Optimizer/Transforms/StackArrays.cpp
+++ b/flang/lib/Optimizer/Transforms/StackArrays.cpp
@@ -287,7 +287,7 @@ mlir::ChangeResult LatticePoint::join(const AbstractDenseLattice &lattice) {
 
 void LatticePoint::print(llvm::raw_ostream &os) const {
   for (const auto &[value, state] : stateMap) {
-    os << value << ": ";
+    os << "\n * " << value << ": ";
     ::print(os, state);
   }
 }
@@ -361,6 +361,13 @@ void AllocationAnalysis::visitOperation(mlir::Operation *op,
   } else if (mlir::isa<fir::FreeMemOp>(op)) {
     assert(op->getNumOperands() == 1 && "fir.freemem has one operand");
     mlir::Value operand = op->getOperand(0);
+
+    // Note: StackArrays is scheduled in the pass pipeline after lowering hlfir
+    // to fir. Therefore, we only need to handle `fir::DeclareOp`s.
+    if (auto declareOp =
+            llvm::dyn_cast_if_present<fir::DeclareOp>(operand.getDefiningOp()))
+      operand = declareOp.getMemref();
+
     std::optional<AllocationState> operandState = before.get(operand);
     if (operandState && *operandState == AllocationState::Allocated) {
       // don't tag things not allocated in this function as freed, so that we
@@ -452,6 +459,9 @@ StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
   };
   func->walk([&](mlir::func::ReturnOp child) { joinOperationLattice(child); });
   func->walk([&](fir::UnreachableOp child) { joinOperationLattice(child); });
+  func->walk(
+      [&](mlir::omp::TerminatorOp child) { joinOperationLattice(child); });
+
   llvm::DenseSet<mlir::Value> freedValues;
   point.appendFreedValues(freedValues);
 
@@ -518,9 +528,18 @@ AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
 
   // remove freemem operations
   llvm::SmallVector<mlir::Operation *> erases;
-  for (mlir::Operation *user : allocmem.getOperation()->getUsers())
+  for (mlir::Operation *user : allocmem.getOperation()->getUsers()) {
+    if (auto declareOp = mlir::dyn_cast_if_present<fir::DeclareOp>(user)) {
+      for (mlir::Operation *user : declareOp->getUsers()) {
+        if (mlir::isa<fir::FreeMemOp>(user))
+          erases.push_back(user);
+      }
+    }
+
     if (mlir::isa<fir::FreeMemOp>(user))
       erases.push_back(user);
+  }
+
   // now we are done iterating the users, it is safe to mutate them
   for (mlir::Operation *erase : erases)
     rewriter.eraseOp(erase);
diff --git a/flang/test/Transforms/stack-arrays-hlfir.f90 b/flang/test/Transforms/stack-arrays-hlfir.f90
new file mode 100644
index 0000000000000..7393eedc98de6
--- /dev/null
+++ b/flang/test/Transforms/stack-arrays-hlfir.f90
@@ -0,0 +1,55 @@
+! Similar to stack-arrays.f90; i.e. both test the stack-arrays pass for different
+! kinds of supported inputs. This one differs in that it takes the hlfir lowering
+! path in flag rather than the fir one. For example, temp arrays are lowered
+! differently in hlfir vs. fir and the IR that reaches the stack arrays pass looks
+! quite different.
+
+
+! RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - \
+! RUN:   | fir-opt --lower-hlfir-ordered-assignments \
+! RUN:   | fir-opt --bufferize-hlfir \
+! RUN:   | fir-opt --convert-hlfir-to-fir \
+! RUN:   | fir-opt --array-value-copy \
+! RUN:   | fir-opt --stack-arrays \
+! RUN: | FileCheck %s
+
+subroutine temp_array
+  implicit none
+  integer (8) :: lV
+  integer (8), dimension (2) :: iaVS
+
+  lV = 202
+
+  iaVS = [lV, lV]
+end subroutine temp_array
+! CHECK-LABEL: func.func @_QPtemp_array{{.*}} {
+! CHECK-NOT:     fir.allocmem
+! CHECK-NOT:     fir.freemem
+! CHECK:         fir.alloca !fir.array<2xi64>
+! CHECK-NOT:     fir.allocmem
+! CHECK-NOT:     fir.freemem
+! CHECK:         return
+! CHECK-NEXT:  }
+
+subroutine omp_temp_array
+  implicit none
+  integer (8) :: lV
+  integer (8), dimension (2) :: iaVS
+
+  lV = 202
+
+  !$omp target
+    iaVS = [lV, lV]
+  !$omp end target
+end subroutine omp_temp_array
+! CHECK-LABEL: func.func @_QPomp_temp_array{{.*}} {
+! CHECK:         omp.target {{.*}} {
+! CHECK-NOT:       fir.allocmem
+! CHECK-NOT:       fir.freemem
+! CHECK:           fir.alloca !fir.array<2xi64>
+! CHECK-NOT:       fir.allocmem
+! CHECK-NOT:       fir.freemem
+! CHECK:           omp.terminator
+! CHECK-NEXT:    }
+! CHECK:         return
+! CHECK-NEXT:  }
diff --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index a2ffe555091eb..841fea56c35d1 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -339,13 +339,10 @@ func.func @omp_placement1() {
   return
 }
 // CHECK:      func.func @omp_placement1() {
+// CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<42xi32>
+// CHECK-NEXT:   %[[MEM_CONV:.*]] = fir.convert %[[MEM]] : (!fir.ref<!fir.array<42xi32>>) -> !fir.heap<!fir.array<42xi32>>
 // CHECK-NEXT:   omp.sections {
 // CHECK-NEXT:     omp.section {
-// CHECK-NEXT:       %[[MEM:.*]] = fir.allocmem !fir.array<42xi32>
-// TODO: this allocation should be moved to the stack. Unfortunately, the data
-// flow analysis fails to propogate the lattice out of the omp region to the
-// return satement.
-// CHECK-NEXT:       fir.freemem %[[MEM]] : !fir.heap<!fir.array<42xi32>>
 // CHECK-NEXT:       omp.terminator
 // CHECK-NEXT:     }
 // CHECK-NEXT:     omp.terminator

ergawy · 2024-07-14T12:08:05Z

flang/lib/Optimizer/Transforms/StackArrays.cpp

+
+    // Note: StackArrays is scheduled in the pass pipeline after lowering hlfir
+    // to fir. Therefore, we only need to handle `fir::DeclareOp`s.
+    if (auto declareOp =


I am not sure this a proper use of the DFA framework in this case or not.

In particular, should I instead propagate the information through the fir.declare op instead of doing this indirect check of the operand? The problem is that this will complicate things a bit more:

we will either need to mark both the value returned from the allocmem and the value returned from the declare op as Allocated. I don't think this will be correct.

or we will need a new "intermediate" AllocationState, for example: AllocDeclare (?) to mark the results of fir.declare ops. If the current approach in the PR does not fit how the framework should be used, I think this option might be better.

Let me know if you have any thoughts on this.

I don't have the context on the fir dialect to comment on the memory effects model here, but what I can say is this should be fine. What you want is if the allocation state of the result of declareOp is modified, that this function gets re-invoked on the right user, which the DFA framework should ensure happens properly.

I agree this is fine and a lot simpler than tracking both references to the allocated variable

…mp regions Extends the stack-arrays pass to support `fir.declare` ops. Before that, we did not recognize malloc-free pairs for which `fir.declare` is used to declare the allocated entity. This is because the `free` op was invoked on the result of the `fir.declare` op and did not directly use the allocated memory SSA value. This also extends the pass to collect the analysis results within OpenMP regions.

Mogball

Seems fine to me but I do not have the authority to approve :^)

Mogball · 2024-07-15T21:17:43Z

flang/lib/Optimizer/Transforms/StackArrays.cpp

+
+    // Note: StackArrays is scheduled in the pass pipeline after lowering hlfir
+    // to fir. Therefore, we only need to handle `fir::DeclareOp`s.
+    if (auto declareOp =


I don't have the context on the fir dialect to comment on the memory effects model here, but what I can say is this should be fine. What you want is if the allocation state of the result of declareOp is modified, that this function gets re-invoked on the right user, which the DFA framework should ensure happens properly.

tblah

Thanks very much for this update! Getting stack arrays working after the HLFIR transition was long overdue.

+1 for the changes for the handling of declare ops.

What you have so far for handling OpenMP regions looks good to me because the pass only runs on functions and so this will not effect outlined reduction and privatisation declarations.

tblah · 2024-07-16T10:36:36Z

flang/lib/Optimizer/Transforms/StackArrays.cpp

+
+    // Note: StackArrays is scheduled in the pass pipeline after lowering hlfir
+    // to fir. Therefore, we only need to handle `fir::DeclareOp`s.
+    if (auto declareOp =


I agree this is fine and a lot simpler than tracking both references to the allocated variable

…penmp

github-actions · 2024-07-16T12:18:38Z

✅ With the latest revision this PR passed the C/C++ code formatter.

…penmp

…mp regions (#98810) Summary: Extends the stack-arrays pass to support `fir.declare` ops. Before that, we did not recognize malloc-free pairs for which `fir.declare` is used to declare the allocated entity. This is because the `free` op was invoked on the result of the `fir.declare` op and did not directly use the allocated memory SSA value. This also extends the pass to collect the analysis results within OpenMP regions. Test Plan: Reviewers: Subscribers: Tasks: Tags: Differential Revision: https://phabricator.intern.facebook.com/D60250941

llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir labels Jul 14, 2024

ergawy requested review from tblah, DominikAdamski, Mogball and dpalermo July 14, 2024 12:06

ergawy commented Jul 14, 2024

View reviewed changes

ergawy changed the title ~~[flang][stack-arrays] Extend pass to work on declare ops and within o…~~ [flang][stack-arrays] Extend pass to work on declare ops and within omp regions Jul 14, 2024

ergawy force-pushed the stack_arrays_decl_openmp branch from e1684d2 to 62457c3 Compare July 14, 2024 12:14

ergawy force-pushed the stack_arrays_decl_openmp branch from 62457c3 to 64ee7b3 Compare July 14, 2024 12:15

Mogball reviewed Jul 15, 2024

View reviewed changes

tblah approved these changes Jul 16, 2024

View reviewed changes

ergawy added 2 commits July 16, 2024 07:14

further detections of declare ops

9d642d8

Merge remote-tracking branch 'upstream/main' into stack_arrays_decl_o…

5f12f91

…penmp

ergawy added 4 commits July 16, 2024 07:22

format

be24abb

more testing

460be2e

Merge remote-tracking branch 'upstream/main' into stack_arrays_decl_o…

36ad5c8

…penmp

Merge remote-tracking branch 'upstream/main' into stack_arrays_decl_o…

f72049c

…penmp

ergawy merged commit 464d321 into llvm:main Jul 18, 2024
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[flang][stack-arrays] Extend pass to work on declare ops and within omp regions #98810

[flang][stack-arrays] Extend pass to work on declare ops and within omp regions #98810

Uh oh!

ergawy commented Jul 14, 2024 •

edited

Loading

Uh oh!

llvmbot commented Jul 14, 2024

Uh oh!

ergawy Jul 14, 2024 •

edited

Loading

Uh oh!

Mogball Jul 15, 2024

Uh oh!

tblah Jul 16, 2024

Uh oh!

Mogball left a comment

Uh oh!

Mogball Jul 15, 2024

Uh oh!

tblah left a comment

Uh oh!

tblah Jul 16, 2024

Uh oh!

github-actions bot commented Jul 16, 2024 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

[flang][stack-arrays] Extend pass to work on declare ops and within omp regions #98810

[flang][stack-arrays] Extend pass to work on declare ops and within omp regions #98810

Uh oh!

Conversation

ergawy commented Jul 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jul 14, 2024

Uh oh!

ergawy Jul 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Mogball Jul 15, 2024

Choose a reason for hiding this comment

Uh oh!

tblah Jul 16, 2024

Choose a reason for hiding this comment

Uh oh!

Mogball left a comment

Choose a reason for hiding this comment

Uh oh!

Mogball Jul 15, 2024

Choose a reason for hiding this comment

Uh oh!

tblah left a comment

Choose a reason for hiding this comment

Uh oh!

tblah Jul 16, 2024

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ergawy commented Jul 14, 2024 •

edited

Loading

ergawy Jul 14, 2024 •

edited

Loading

github-actions bot commented Jul 16, 2024 •

edited

Loading