[flang][OpenMP] Fix use-after-free in OMPFunctionFiltering

[kparzysz]

When walking over functions (in pre-order), if the function being visited needs to be erased, skip visiting its regions.

This was detected by address sanitizer.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

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

Author: Krzysztof Parzyszek (kparzysz)

Changes

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.

---

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

1 Files Affected:

• (modified) flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp (+5-1)

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

[skatrak]

Thank you for noticing this issue. I think we can simplify the fix a bit, though. The docs for Operation::walk() state that the callback is allowed to erase that operation when walking in pre-order if it is skipped after the erasure. It seems simpler and cheaper to just return WalkResult::skip() at that point rather than constructing a list to be deleted later.

[kparzysz]

Done,

[skatrak]

Thanks, LGTM