[Mlir-commits] [mlir] [mlir] Add option for a cleanup pattern set to SCF tiling helper (PR #109554)

[Quinn Dawkins]

================
@@ -1315,6 +1317,172 @@ FailureOr<SmallVector<Operation *>> mlir::scf::yieldReplacementForFusedProducer(
   return generatedSlices;
 }
 
+namespace {
+
+//===----------------------------------------------------------------------===//
+// SliceWorklist
+//===----------------------------------------------------------------------===//
+
+/// Struct for tracking the number of stale entries on the worklist and whether
+/// there is a remaining valid entry.
+struct EntryCount {
+  bool isValid = true;
+  unsigned count = 0;
+};
+
+/// A FIFO worklist of operations with efficient removal and set semantics.
+///
+/// This class maintains a queue of operations and a mapping of operations to
+/// positions in the vector, so that operations can be removed efficiently at
+/// random. When an operation is removed, it is replaced with nullptr. Such
+/// nullptr are skipped when pop'ing elements.
+///
+/// This is similar to the worklist used by the GreedyPatternRewriteDriver,
+/// except instead FIFO so that slices for fusion can be processed breadth
+/// first.
+class SliceWorklist {
+public:
+  SliceWorklist() = default;
+
+  /// Push an operation to the end of the worklist. This assumes that
+  /// the given operation is not already on the worklist.
+  void push(Operation *op);
+
+  /// Pop the an operation from the end of the worklist. Returns nullptr if
+  /// there are no remaining valid operations.
+  Operation *pop();
+
+  /// Remove an operation from the worklist.
+  void remove(Operation *op);
+
+protected:
+  /// The queue of operations.
+  std::deque<Operation *> list;
+
+  /// A mapping of operations to the number of stale copies in the queue.
+  DenseMap<Operation *, EntryCount> map;
+};
+
+void SliceWorklist::push(Operation *op) {
+  assert(op && "cannot push nullptr to worklist");
+  list.push_back(op);
+  EntryCount newCount = map.lookup(op);
+  // Because operations are only pushed on creation, valid duplicates are
+  // never added.
+  assert((!map.contains(op) || !newCount.isValid) &&
+         "cannot push a duplicate operation");
+  map[op] = {/*isValid=*/true, newCount.count + 1};
----------------
qedawkins wrote:

Dropped the class and changed the worklist to remove on the fly as discussed offline.

https://github.com/llvm/llvm-project/pull/109554