[Mlir-commits] [mlir] Add worklist comparator to tile and fuse options (PR #211523)
Hagai Lev Hacohen
llvmlistbot at llvm.org
Thu Jul 23 04:10:47 PDT 2026
https://github.com/HagaiLevHacohen updated https://github.com/llvm/llvm-project/pull/211523
>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 1/2] 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)
}];
}
>From 0d4d897399253a3b31be6c8c4c4dbcd1bd98e910 Mon Sep 17 00:00:00 2001
From: Hagai Lev Hacohen <hagai4000 at gmail.com>
Date: Thu, 23 Jul 2026 14:10:21 +0300
Subject: [PATCH 2/2] lint
---
.../SCF/Transforms/TileUsingInterface.h | 27 +++++++--------
.../SCF/Transforms/TileUsingInterface.cpp | 34 +++++++++----------
.../TestTilingInterfaceTransformOps.cpp | 13 ++++---
3 files changed, 36 insertions(+), 38 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index 125ff0edb7a89..034bfbec5b77b 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -312,20 +312,19 @@ 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;
- }
+ /// 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
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 00a27b4dc66c5..b856174b787cf 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1645,9 +1645,9 @@ class SliceTrackingListener : public RewriterBase::Listener {
/// `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);
+ tensor::ExtractSliceOp
+ popNext(const scf::SCFTileAndFuseOptions::TileOrderControlFnTy
+ &tileOrderControlFn);
/// The worklist for this transformation keeps track of the slices to visit
/// next for fusion.
@@ -1716,22 +1716,22 @@ void SliceTrackingListener::notifyOperationReplaced(Operation *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();
+ 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;
}
-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
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index d891e576b7dd4..42139c110707d 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -60,13 +60,12 @@ static llvm::SmallDenseSet<Operation *> collectTiledAndFusedOps(Operation *op) {
/// Apply a tile and fuse transformation to all payload ops and store both the
/// tiled operation as well as the created tile loops.
template <typename Range>
-static LogicalResult
-applyTileAndFuseToAll(RewriterBase &rewriter, Operation *transformOp,
- Range &&payloadOps, unsigned numLoops,
- scf::SCFTilingOptions tilingOptions,
- TransformResults &transformResults,
- scf::SCFTileAndFuseOptions::TileOrderControlFnTy
- tileOrderControlFn = nullptr) {
+static LogicalResult applyTileAndFuseToAll(
+ RewriterBase &rewriter, Operation *transformOp, Range &&payloadOps,
+ unsigned numLoops, scf::SCFTilingOptions tilingOptions,
+ TransformResults &transformResults,
+ scf::SCFTileAndFuseOptions::TileOrderControlFnTy tileOrderControlFn =
+ nullptr) {
SmallVector<Operation *> tiledOps;
SmallVector<SmallVector<Operation *>> loopOps(numLoops);
More information about the Mlir-commits
mailing list