From aa568e082b4c0aa1cfbc8d1937544af8adbde552 Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Wed, 18 May 2022 02:14:02 -0700
Subject: [PATCH] [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.
---
 mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

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());