[Mlir-commits] [mlir] Add worklist comparator to tile and fuse options (PR #211523)
Hagai Lev Hacohen
llvmlistbot at llvm.org
Thu Jul 23 04:08:46 PDT 2026
https://github.com/HagaiLevHacohen created https://github.com/llvm/llvm-project/pull/211523
## Summary
`tileConsumerAndFuseProducersUsingSCF` currently processes its `tensor.extract_slice` worklist in FIFO order, resulting in a breadth-first traversal.
This MR adds an optional comparator to `SCFTileAndFuseOptions`, allowing callers to prioritize worklist entries. The default remains FIFO. Our intended use is a topological consumer-to-producer traversal.
## Motivation
Consider this dependency graph, where arrows point from a consumer to its producer \[For a pseudo MLIR example - see below\]:
```text
1 → 2 → 3 → 4 → 5
└────────→ 4 → 5
```
Node 1 consumes nodes 2 and 4, forming a skip connection. Suppose we tile node 1 and fuse the entire producer graph, and nodes 1 and 3 request the same slice from node 4.
With FIFO processing, the fusion order is:
```text
2, 4, 3, 5
```
Node 4 is processed before node 3. Consequently:
1. Processing node 1 creates an extract slice from node 4.
2. That slice is processed and removed before node 3 is fused.
3. Fusing node 3 creates an identical extract slice from node 4.
Because the identical slices never coexist in the worklist, cleanup patterns cannot deduplicate them.
When this skip-connection pattern is repeated, redundant slices and fusion work grow exponentially. In our workloads, this causes compilation time to explode and may prevent compilation from completing.
With topological prioritization, the order becomes:
```text
2, 3, 4, 5
```
Node 3 is processed before node 4. Therefore, both identical slices from node 4 coexist before either is processed, allowing cleanup patterns to eliminate the duplicate.
This change lets callers select such an ordering while preserving FIFO behavior by default.
## Pseudo MLIR Example
```mlir
// Pseudo-MLIR; operation details and types are omitted.
// Node 5
%node5 = linalg.fill ...
// Node 4
%node4 = linalg.generic ins(%node5) ...
// Node 3
%node3 = linalg.generic ins(%node4) ...
// Node 2
%node2 = linalg.generic ins(%node3) ...
// Node 1: contains the skip connection to node 4.
%node1 = linalg.add ins(%node2, %node4) ...
```
When node 1 is tiled, it creates a slice from node 4:
```mlir
%slice_from_node1 = tensor.extract_slice %node4[%offset] [%size] [1]
```
Later, fusing node 3 creates the same slice:
```mlir
%slice_from_node3 = tensor.extract_slice %node4[%offset] [%size] [1]
```
With FIFO ordering, `%slice_from_node1` is processed before `%slice_from_node3` is created. With topological ordering, both slices coexist, allowing cleanup patterns to deduplicate them.
>From 389b3b7cd4326e8f3f076ac160c8d99693b5b59b Mon Sep 17 00:00:00 2001
From: Hagai Lev Hacohen <hagai4000 at gmail.com>
Date: Wed, 22 Jul 2026 15:23:09 +0300
Subject: [PATCH] Adding to tile and fuse options a worklist comparator
---
.../SCF/Transforms/TileUsingInterface.h | 15 +++
.../SCF/Transforms/TileUsingInterface.cpp | 30 +++++-
.../tile-fuse-and-yield-using-interface.mlir | 98 +++++++++++++++++++
.../TestTilingInterfaceTransformOps.cpp | 15 ++-
.../TestTilingInterfaceTransformOps.td | 10 +-
5 files changed, 162 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index 0bd5d19b136d0..125ff0edb7a89 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -312,6 +312,21 @@ struct SCFTileAndFuseOptions {
return *this;
}
+ /// Comparator used to select the next `tensor.extract_slice` to process from
+ /// the fusion worklist. Returns true if `lhs` should be processed before
+ /// `rhs`. This allows callers to enforce a desired tiling order. For example,
+ /// to process producers in topological order, a producer that is an ancestor
+ /// of another producer in the defining-op chain can be ordered before it.
+ /// By default, the worklist is processed in FIFO order.
+ using TileOrderControlFnTy =
+ std::function<bool(tensor::ExtractSliceOp lhs,
+ tensor::ExtractSliceOp rhs)>;
+ TileOrderControlFnTy tileOrderControlFn = nullptr;
+ SCFTileAndFuseOptions &setTileOrderControlFn(TileOrderControlFnTy controlFn) {
+ tileOrderControlFn = std::move(controlFn);
+ return *this;
+ }
+
/// An optional set of rewrite patterns to apply to the results of tiling
/// before fusion. This will track deleted and newly inserted
/// `tensor.extract_slice` ops and update the worklist.
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index ad4aff893a03a..00a27b4dc66c5 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1641,6 +1641,14 @@ class SliceTrackingListener : public RewriterBase::Listener {
/// Remove the operation from the worklist.
void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
+ /// Pop the next slice to process from the worklist. When
+ /// `tileOrderControlFn` is set, returns the preferred slice according to the
+ /// callback. Otherwise pops from the front (FIFO). The worklist must be
+ /// non-empty.
+ tensor::ExtractSliceOp popNext(
+ const scf::SCFTileAndFuseOptions::TileOrderControlFnTy
+ &tileOrderControlFn);
+
/// The worklist for this transformation keeps track of the slices to visit
/// next for fusion.
std::deque<tensor::ExtractSliceOp> worklist;
@@ -1707,6 +1715,24 @@ void SliceTrackingListener::notifyOperationReplaced(Operation *op,
removeOp(op);
}
+tensor::ExtractSliceOp SliceTrackingListener::popNext(
+ const scf::SCFTileAndFuseOptions::TileOrderControlFnTy
+ &tileOrderControlFn) {
+assert(!worklist.empty() && "expected non-empty worklist");
+if (!tileOrderControlFn) {
+ auto slice = worklist.front();
+ worklist.pop_front();
+ return slice;
+}
+auto it = llvm::min_element(worklist, [&](tensor::ExtractSliceOp lhs,
+ tensor::ExtractSliceOp rhs) {
+ return tileOrderControlFn(lhs, rhs);
+});
+auto slice = *it;
+worklist.erase(it);
+return slice;
+}
+
//===----------------------------------------------------------------------===//
// ReplacementListener
//===----------------------------------------------------------------------===//
@@ -1814,8 +1840,8 @@ mlir::scf::tileConsumerAndFuseProducersUsingSCF(
}
OpBuilder::InsertionGuard g(rewriter);
while (!sliceTracker.worklist.empty()) {
- auto candidateSlice = sliceTracker.worklist.front();
- sliceTracker.worklist.pop_front();
+ tensor::ExtractSliceOp candidateSlice =
+ sliceTracker.popNext(options.tileOrderControlFn);
auto [fusableProducer, destinationInitArg] =
getUntiledProducerFromSliceSource(&candidateSlice.getSourceMutable(),
diff --git a/mlir/test/Interfaces/TilingInterface/tile-fuse-and-yield-using-interface.mlir b/mlir/test/Interfaces/TilingInterface/tile-fuse-and-yield-using-interface.mlir
index 3c0ada9d2cabc..c11c7c4a5fc50 100644
--- a/mlir/test/Interfaces/TilingInterface/tile-fuse-and-yield-using-interface.mlir
+++ b/mlir/test/Interfaces/TilingInterface/tile-fuse-and-yield-using-interface.mlir
@@ -120,3 +120,101 @@ module attributes {transform.with_named_sequence} {
// CHECK: %[[INSERT2:.+]] = tensor.insert_slice %[[GENERIC_TILE]]#1 into %[[ITERARG2]][0, %[[IV]]]
// CHECK: scf.yield %[[INSERT0]], %[[INSERT1]], %[[INSERT2]]
// CHECK: return %[[RESULT]]#1, %[[RESULT]]#2, %[[RESULT]]#0
+
+// -----
+
+// Verify that the default FIFO worklist reconstructs the earlier `linalg.fill`
+// before the later `linalg.copy`. Their loop results are therefore passed to
+// the downstream `linalg.mul` as result #1 and result #2, respectively.
+
+// CHECK-LABEL: func.func @worklist_fifo_order(
+// CHECK: %[[RESULT:.+]]:3 = scf.for
+// CHECK: %[[FILL_TILE:.+]] = linalg.fill
+// CHECK: %[[COPY_TILE:.+]] = linalg.copy
+// CHECK: %[[CONSUMER_TILE:.+]] = linalg.add
+// CHECK: tensor.insert_slice %[[CONSUMER_TILE]]
+// CHECK: tensor.insert_slice %[[FILL_TILE]]
+// CHECK: tensor.insert_slice %[[COPY_TILE]]
+// CHECK: scf.yield
+// CHECK: %[[FINAL:.+]] = linalg.mul
+// CHECK-SAME: ins(%[[RESULT]]#1, %[[RESULT]]#2
+// CHECK-SAME: outs(%[[RESULT]]#0
+// CHECK: return %[[FINAL]]
+
+func.func @worklist_fifo_order(
+ %input: tensor<32x32xf32>, %out: tensor<32x32xf32>)
+ -> tensor<32x32xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%c0 : f32)
+ outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
+ %copy = linalg.copy ins(%input : tensor<32x32xf32>)
+ outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
+ %add = linalg.add
+ ins(%fill, %copy : tensor<32x32xf32>, tensor<32x32xf32>)
+ outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
+ %result = linalg.mul
+ ins(%fill, %copy : tensor<32x32xf32>, tensor<32x32xf32>)
+ outs(%add : tensor<32x32xf32>) -> tensor<32x32xf32>
+ return %result : tensor<32x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %arg0 : !transform.any_op {transform.readonly}) {
+ %add = transform.structured.match ops{["linalg.add"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.test.fuse_and_yield %add [16]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Verify that preferring later producers reconstructs the `linalg.copy` before
+// the `linalg.fill`. Their loop result numbers are reversed, but the downstream
+// `linalg.mul` still receives the fill and copy values in its original order.
+
+// CHECK-LABEL: func.func @worklist_prefer_later_producers(
+// CHECK: %[[RESULT:.+]]:3 = scf.for
+// CHECK: %[[FILL_TILE:.+]] = linalg.fill
+// CHECK: %[[COPY_TILE:.+]] = linalg.copy
+// CHECK: %[[CONSUMER_TILE:.+]] = linalg.add
+// CHECK: tensor.insert_slice %[[CONSUMER_TILE]]
+// CHECK: tensor.insert_slice %[[COPY_TILE]]
+// CHECK: tensor.insert_slice %[[FILL_TILE]]
+// CHECK: scf.yield
+// CHECK: %[[FINAL:.+]] = linalg.mul
+// CHECK-SAME: ins(%[[RESULT]]#2, %[[RESULT]]#1
+// CHECK-SAME: outs(%[[RESULT]]#0
+// CHECK: return %[[FINAL]]
+
+
+func.func @worklist_prefer_later_producers(
+ %input: tensor<32x32xf32>, %out: tensor<32x32xf32>)
+ -> tensor<32x32xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%c0 : f32)
+ outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
+ %copy = linalg.copy ins(%input : tensor<32x32xf32>)
+ outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
+ %add = linalg.add
+ ins(%fill, %copy : tensor<32x32xf32>, tensor<32x32xf32>)
+ outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
+ %result = linalg.mul
+ ins(%fill, %copy : tensor<32x32xf32>, tensor<32x32xf32>)
+ outs(%add : tensor<32x32xf32>) -> tensor<32x32xf32>
+ return %result : tensor<32x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %arg0 : !transform.any_op {transform.readonly}) {
+ %add = transform.structured.match ops{["linalg.add"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.test.fuse_and_yield %add [16]
+ worklist_prefer_later_producers true
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index 9467c925e543c..d891e576b7dd4 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -64,7 +64,9 @@ static LogicalResult
applyTileAndFuseToAll(RewriterBase &rewriter, Operation *transformOp,
Range &&payloadOps, unsigned numLoops,
scf::SCFTilingOptions tilingOptions,
- TransformResults &transformResults) {
+ TransformResults &transformResults,
+ scf::SCFTileAndFuseOptions::TileOrderControlFnTy
+ tileOrderControlFn = nullptr) {
SmallVector<Operation *> tiledOps;
SmallVector<SmallVector<Operation *>> loopOps(numLoops);
@@ -98,6 +100,7 @@ applyTileAndFuseToAll(RewriterBase &rewriter, Operation *transformOp,
yieldProducerReplacement};
};
tileAndFuseOptions.setFusionControlFn(controlFn);
+ tileAndFuseOptions.setTileOrderControlFn(tileOrderControlFn);
rewriter.setInsertionPoint(target);
FailureOr<scf::SCFTileAndFuseResult> tiledResults =
@@ -159,10 +162,18 @@ transform::TestFuseAndYieldOp::apply(TransformRewriter &rewriter,
tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);
}
+ scf::SCFTileAndFuseOptions::TileOrderControlFnTy tileOrderControlFn;
+ if (getWorklistPreferLaterProducers()) {
+ tileOrderControlFn = [](tensor::ExtractSliceOp lhs,
+ tensor::ExtractSliceOp rhs) {
+ return rhs.getOperation()->isBeforeInBlock(lhs.getOperation());
+ };
+ }
+
LogicalResult result = applyTileAndFuseToAll(
rewriter, getOperation(), state.getPayloadOps(getTarget()),
tileSizes.size() - llvm::count(tileSizes, 0), tilingOptions,
- transformResults);
+ transformResults, tileOrderControlFn);
return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
: DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
index efa16212f0fef..03032e7efb7ca 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td
@@ -30,6 +30,9 @@ def TestFuseAndYieldOp : Op<Transform_Dialect, "test.fuse_and_yield",
producers greedily using the options provided as attributes.
It also yields some of the fused producers for testing.
+ The `worklist_prefer_later_producers` attribute processes slices that occur
+ later in the tiled IR before earlier slices.
+
On success returns the tiled operations as well as generated loops. Emits
a definite failure if tiling fails.
}];
@@ -38,13 +41,16 @@ def TestFuseAndYieldOp : Op<Transform_Dialect, "test.fuse_and_yield",
(ins TransformHandleTypeInterface:$target,
DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,
DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_interchange,
- DefaultValuedAttr<BoolAttr, "false">:$use_forall);
+ DefaultValuedAttr<BoolAttr, "false">:$use_forall,
+ DefaultValuedAttr<BoolAttr, "false">:$worklist_prefer_later_producers);
let results = (outs TransformHandleTypeInterface:$transfomed,
Variadic<TransformHandleTypeInterface>:$loops);
let assemblyFormat = [{
$target ($tile_sizes^)? (`interchange` $tile_interchange^)?
- (`use_forall` $use_forall^)? attr-dict
+ (`use_forall` $use_forall^)?
+ (`worklist_prefer_later_producers` $worklist_prefer_later_producers^)?
+ attr-dict
`:` functional-type(operands, results)
}];
}
More information about the Mlir-commits
mailing list