Bugfix: Keep worklistMap in sync with worklist in GreedyPatternRewriter

When we removed a pattern, we removed it from worklist but not from
worklistMap. Then, when we tried to add a new pattern on the same Operation
again, the pattern wasn't added since it already existed in the
worklistMap (but not in the worklist).

Closes tensorflow/mlir#211

PiperOrigin-RevId: 277319669

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 1619671..56701e3 100644 (file)
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -79,6 +79,7 @@ public:
     if (it != worklistMap.end()) {
       assert(worklist[it->second] == op && "malformed worklist data structure");
       worklist[it->second] = nullptr;
+      worklistMap.erase(it);
     }
   }