[flang-commits] [flang] [flang][OpenMP] Fix use-after-tree in OMPFunctionFiltering (PR #84373)

[Krzysztof Parzyszek via flang-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/84373

Erasing the element of the range that it being iterated on can invalidate the range. Instead of erasing function as we see them, store them in a separate list, then erase them after the range has been traversed.

This was detected by address sanitizer.

>From 986f335202954a061c4aac04ba5d5f3d2e2d7791 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Thu, 7 Mar 2024 14:11:33 -0600
Subject: [PATCH] [flang][OpenMP] Fix use-after-tree in OMPFunctionFiltering

Erasing the element of the range that it being iterated on can
invalidate the range. Instead of erasing function as we see them,
store them in a separate list, then erase them after the range
has been traversed.

This was detected by address sanitizer.
---
 flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
index 466bf53e8dbd60..f8bede88293ff5 100644
--- a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
+++ b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
@@ -40,6 +40,8 @@ class OMPFunctionFilteringPass
     if (!op || !op.getIsTargetDevice())
       return;
 
+    llvm::SmallVector<func::FuncOp> removedFuncs;
+
     op->walk<WalkOrder::PreOrder>([&](func::FuncOp funcOp) {
       // Do not filter functions with target regions inside, because they have
       // to be available for both host and device so that regular and reverse
@@ -80,12 +82,14 @@ class OMPFunctionFilteringPass
           callOp->erase();
         }
         if (!hasTargetRegion)
-          funcOp.erase();
+          removedFuncs.push_back(funcOp);
         else if (declareTargetOp)
           declareTargetOp.setDeclareTarget(declareType,
                                            omp::DeclareTargetCaptureClause::to);
       }
     });
+    for (func::FuncOp f : removedFuncs)
+      f.erase();
   }
 };
 } // namespace