Remove worklist class and add negative test

Author: qedawkins

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -1319,95 +1319,6 @@ FailureOr<SmallVector<Operation *>> mlir::scf::yieldReplacementForFusedProducer(
 
 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};
-}
-
-Operation *SliceWorklist::pop() {
-  // Pop the front of the queue until we hit a valid entry.
-  while (!list.empty()) {
-    Operation *op = list.front();
-    list.pop_front();
-
-    EntryCount e = map.lookup(op);
-    // If the entry count is greater than 1 or there is no valid entry,
-    // this must be a stale entry. Decrement the map entry by one and continue.
-    if (e.count > 1 || !e.isValid) {
-      int64_t newCount = e.count - 1;
-      if (newCount <= 0)
-        map.erase(op);
-      else
-        map[op] = {e.isValid, static_cast<unsigned int>(newCount)};
-      continue;
-    }
-
-    map.erase(op);
-    return op;
-  }
-  return nullptr;
-}
-
-// Mark the operation as invalid if present. Removal from the map will
-// happen later when popping from the worklist.
-void SliceWorklist::remove(Operation *op) {
-  if (!map.contains(op))
-    return;
-
-  EntryCount e = map.lookup(op);
-  map[op] = {/*isValid=*/false, e.count};
-}
-
 //===----------------------------------------------------------------------===//
 // SliceTrackingListener
 //===----------------------------------------------------------------------===//
@@ -1430,15 +1341,18 @@ class SliceTrackingListener : public RewriterBase::Listener {
   void notifyOperationInserted(Operation *op,
                                OpBuilder::InsertPoint previous) override;
 
+  /// Shared helper for operation removal from the worklist.
+  void removeOp(Operation *op);
+
   /// Remove the operation from the worklist.
   void notifyOperationErased(Operation *op) override;
 
   /// Remove the operation from the worklist.
   void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
 
-  /// The worklist for this transformation keeps track of the operations that
-  /// need to be (re)visited.
-  SliceWorklist worklist;
+  /// The worklist for this transformation keeps track of the slices to visit
+  /// next for fusion.
+  std::deque<tensor::ExtractSliceOp> worklist;
 
 private:
   /// Optional pattern set to apply when adding new operations to the worklist.
@@ -1453,8 +1367,8 @@ SliceTrackingListener::SliceTrackingListener(
 LogicalResult
 SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
   for (Operation *op : ops) {
-    if (isa<tensor::ExtractSliceOp>(op))
-      worklist.push(op);
+    if (auto slice = dyn_cast<tensor::ExtractSliceOp>(op))
+      worklist.push_back(slice);
   }
 
   if (!patterns)
@@ -1468,18 +1382,36 @@ SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
 
 void SliceTrackingListener::notifyOperationInserted(
     Operation *op, OpBuilder::InsertPoint previous) {
+  auto slice = dyn_cast<tensor::ExtractSliceOp>(op);
+  if (!slice)
+    return;
+  worklist.push_back(slice);
+}
+
+// Scan the worklist for the given op and remove it if present. The expectation
+// is for the worklist to be small and for removal to be relatively rare.
+void SliceTrackingListener::removeOp(Operation *op) {
   if (!isa<tensor::ExtractSliceOp>(op))
     return;
-  worklist.push(op);
+  auto iter = worklist.begin();
+  while (iter != worklist.end()) {
+    if (*iter == op)
+      break;
+    iter++;
+  }
+  if (iter == worklist.end())
+    return;
+
+  worklist.erase(iter);
 }
 
 void SliceTrackingListener::notifyOperationErased(Operation *op) {
-  worklist.remove(op);
+  removeOp(op);
 }
 
 void SliceTrackingListener::notifyOperationReplaced(Operation *op,
                                                     ValueRange replacement) {
-  worklist.remove(op);
+  removeOp(op);
 }
 } // namespace
 
@@ -1547,10 +1479,9 @@ mlir::scf::tileConsumerAndFuseProducersUsingSCF(
     return rewriter.notifyMatchFailure(consumer, "cleanup patterns failed");
   }
   OpBuilder::InsertionGuard g(rewriter);
-  while (Operation *next = sliceTracker.worklist.pop()) {
-    auto candidateSlice = dyn_cast<tensor::ExtractSliceOp>(next);
-    if (!candidateSlice)
-      continue;
+  while (!sliceTracker.worklist.empty()) {
+    auto candidateSlice = sliceTracker.worklist.front();
+    sliceTracker.worklist.pop_front();
 
     auto [fusableProducer, destinationInitArg] =
         getUntiledProducerFromSliceSource(&candidateSlice.getSourceMutable(),

mlir/test/Dialect/Linalg/transform-op-fuse.mlir

@@ -243,3 +243,38 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// CHECK-LABEL: func.func @fuse_unrelated_slice
+func.func @fuse_unrelated_slices(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<10x10xf32>) {
+
+  //     CHECK: %[[SLICE1:.+]] = tensor.extract_slice
+  //     CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[SLICE1]]
+  //     CHECK: %[[RES:.*]] = scf.for
+  //     CHECK:     scf.for
+  //     CHECK:       linalg.elemwise_unary
+  //     CHECK:       linalg.elemwise_binary
+  //     CHECK: return %[[RES]], %[[SLICE2]]
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %dim0 = tensor.dim %arg1, %c0 : tensor<?x?xf32>
+  %dim1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %slice1 = tensor.extract_slice %arg0 [1, 1] [%dim0, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %slice2 = tensor.extract_slice %slice1 [1, 1] [10, 10] [1, 1] : tensor<?x?xf32> to tensor<10x10xf32>
+  %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
+                             outs(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = tensor.extract_slice %0 [1, 1] [%dim0, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+  %2 = linalg.elemwise_binary ins(%1, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+                             outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %2, %slice2 : tensor<?x?xf32>, tensor<10x10xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1], apply_cleanup = true}
+      : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)
+    transform.yield
+  }
+}