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

[Krzysztof Parzyszek via flang-commits]

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

>From 4bf5b953eff90ca1198a5feca8c3d5dcee52b4af 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 1/2] [flang][OpenMP] Fix use-after-free 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

>From 62ec9b10a364eddc3bbe2932983df34650b42e3c Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Fri, 8 Mar 2024 07:59:44 -0600
Subject: [PATCH 2/2] Skip erased functions

---
 .../Optimizer/Transforms/OMPFunctionFiltering.cpp   | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
index f8bede88293ff5..959099d039a5e6 100644
--- a/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
+++ b/flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp
@@ -40,8 +40,6 @@ 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
@@ -81,15 +79,16 @@ class OMPFunctionFilteringPass
           // Remove the callOp
           callOp->erase();
         }
-        if (!hasTargetRegion)
-          removedFuncs.push_back(funcOp);
-        else if (declareTargetOp)
+        if (!hasTargetRegion) {
+          funcOp.erase();
+          return WalkResult::skip();
+        }
+        if (declareTargetOp)
           declareTargetOp.setDeclareTarget(declareType,
                                            omp::DeclareTargetCaptureClause::to);
       }
+      return WalkResult::advance();
     });
-    for (func::FuncOp f : removedFuncs)
-      f.erase();
   }
 };
 } // namespace