[Mlir-commits] [mlir] aa568e0 - [mlir:GreedyDriver] Return WalkResult::skip after deleting a known constant

[River Riddle]

Author: River Riddle
Date: 2022-05-18T02:14:02-07:00
New Revision: aa568e082b4c0aa1cfbc8d1937544af8adbde552

URL: https://github.com/llvm/llvm-project/commit/aa568e082b4c0aa1cfbc8d1937544af8adbde552
DIFF: https://github.com/llvm/llvm-project/commit/aa568e082b4c0aa1cfbc8d1937544af8adbde552.diff

LOG: [mlir:GreedyDriver] Return WalkResult::skip after deleting a known constant

This avoids use-after-free when trying to access the regions after visiting
the operation.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b2e965945101..a08657cb7f6e 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -159,11 +159,15 @@ bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
       }
     } else {
       // Add all nested operations to the worklist in preorder.
-      for (auto &region : regions)
+      for (auto &region : regions) {
         region.walk<WalkOrder::PreOrder>([&](Operation *op) {
-          if (!insertKnownConstant(op))
+          if (!insertKnownConstant(op)) {
             worklist.push_back(op);
+            return WalkResult::advance();
+          }
+          return WalkResult::skip();
         });
+      }
 
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());