[Mlir-commits] [mlir] [mlir] Fix scalable index handling and support scalable forall tiling (PR #235596)
MiaoMing Chen
llvmlistbot at llvm.org
Mon Oct 26 06:57:38 PDT 2026
https://github.com/chenmiaoming updated https://github.com/llvm/llvm-project/pull/235596
>From 66601e39ea7eda24e8bcd83fb7dbd1af93ee25cf Mon Sep 17 00:00:00 2001
From: MiaoMing Chen <miaomingc at stumail.nwu.edu.cn>
Date: Mon, 26 Oct 2026 10:50:28 +0800
Subject: [PATCH 1/3] [mlir] Fix scalable index handling and support scalable
forall tiling
Index-list parser overloads that do not return scalable flags silently
discard them, so [[8], 1] can be accepted and interpreted as [8, 1]. Reject
scalable indices in these overloads instead of losing their semantics.
Add scalable tile sizes to tile_using_forall. Preserve the flags in the
operation and materialize each nonzero scalable size as its static factor
times vector.vscale in each payload target's scope. Keep zero dimensions
untiled and support the new syntax in the C++ and Python builders.
Add parser, verifier, round-trip, Python builder, and tiling/fusion
regression coverage, including bounded slices and multiple targets.
Fixes https://github.com/llvm/llvm-project/issues/150203
---
.../Linalg/TransformOps/LinalgTransformOps.td | 15 +-
.../mlir/Dialect/Transform/Utils/Utils.h | 20 +-
.../mlir/Interfaces/ViewLikeInterface.h | 12 +-
.../TransformOps/LinalgTransformOps.cpp | 55 +++--
.../mlir/dialects/transform/structured.py | 29 ++-
.../Linalg/tile-to-forall-scalable.mlir | 198 ++++++++++++++++++
.../Dialect/Linalg/transform-ops-invalid.mlir | 76 +++++++
mlir/test/Dialect/Linalg/transform-ops.mlir | 33 +++
mlir/test/Dialect/Tensor/invalid.mlir | 24 +++
.../dialects/transform_structured_ext.py | 29 +++
10 files changed, 460 insertions(+), 31 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index c31cd399032986..eb6dcbba21714d 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -2417,6 +2417,11 @@ def TileUsingForallOp :
Static zero tile sizes indicate that the dimension is not tiled and can be
thought of as tiling by the full size of data.
+ A static tile size enclosed in brackets is scalable: `[8]` denotes
+ `8 * vector.vscale`. Scalable tile sizes are supported only by `tile_sizes`,
+ not by `num_threads` or packed tile size handles. A scalable zero tile size
+ also leaves the dimension untiled.
+
It is the user's responsibility to ensure that `num_threads/tile_sizes` is
a valid tiling specification (i.e. that only tiles parallel dimensions,
e.g. in the Linalg case). If the dimension is not parallelizable, a warning
@@ -2473,6 +2478,7 @@ def TileUsingForallOp :
Optional<TransformAnyParamTypeOrAnyHandle>:$packed_tile_sizes,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_num_threads,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_tile_sizes,
+ DefaultValuedOptionalAttr<DenseBoolArrayAttr, "{}">:$scalable_sizes,
OptionalAttr<DeviceMappingArrayAttr>:$mapping);
let results = (outs TransformHandleTypeInterface:$tiled_op,
TransformHandleTypeInterface:$forall_op);
@@ -2482,12 +2488,14 @@ def TileUsingForallOp :
"ArrayRef<int64_t>":$staticTileSizes,
CArg<"::mlir::transform::TileSizesSpec",
"::mlir::transform::TileSizesSpec()">,
- CArg<"ArrayAttr", "{}">:$mapping)>,
+ CArg<"ArrayAttr", "{}">:$mapping,
+ CArg<"std::optional<ArrayRef<bool>>", "std::nullopt">:$scalableSizes)>,
OpBuilder<(ins "Value":$target,
"ArrayRef<OpFoldResult>":$mixedTileSizes,
CArg<"::mlir::transform::TileSizesSpec",
"::mlir::transform::TileSizesSpec()">,
- CArg<"ArrayAttr", "{}">:$mapping)>,
+ CArg<"ArrayAttr", "{}">:$mapping,
+ CArg<"std::optional<ArrayRef<bool>>", "std::nullopt">:$scalableSizes)>,
OpBuilder<(ins "Value":$target,
"ArrayRef<int64_t>":$staticNumThreads,
CArg<"::mlir::transform::NumThreadsSpec",
@@ -2507,7 +2515,8 @@ def TileUsingForallOp :
$static_num_threads) |
`tile_sizes` custom<PackedOrDynamicIndexList>($packed_tile_sizes,
$tile_sizes,
- $static_tile_sizes))
+ $static_tile_sizes,
+ $scalable_sizes))
(`(` `mapping` `=` $mapping^ `)`)? attr-dict
`:` functional-type(operands, results)
}];
diff --git a/mlir/include/mlir/Dialect/Transform/Utils/Utils.h b/mlir/include/mlir/Dialect/Transform/Utils/Utils.h
index ea7273d7fac790..83203e07d1cfc6 100644
--- a/mlir/include/mlir/Dialect/Transform/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Transform/Utils/Utils.h
@@ -13,6 +13,7 @@
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/STLExtras.h"
namespace mlir {
class OpAsmPrinter;
@@ -72,7 +73,8 @@ inline void printPackedOrDynamicIndexList(OpAsmPrinter &printer, Operation *op,
/// custom<PackedOrDynamicIndexList>($packed, type($packed), $values,
/// type($values), $integers)
///
-/// See `printPackedOrDynamicIndexList` for details.
+/// See `printPackedOrDynamicIndexList` for details. Overloads that do not
+/// return scalable flags reject scalable indices.
ParseResult parsePackedOrDynamicIndexList(
OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> &packed,
Type &packedType, SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
@@ -83,8 +85,14 @@ inline ParseResult parsePackedOrDynamicIndexList(
Type &packedType, SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
SmallVectorImpl<Type> *valueTypes, DenseI64ArrayAttr &integers) {
DenseBoolArrayAttr scalableFlags;
- return parsePackedOrDynamicIndexList(parser, packed, packedType, values,
- scalableFlags, valueTypes, integers);
+ if (failed(parsePackedOrDynamicIndexList(parser, packed, packedType, values,
+ scalableFlags, valueTypes,
+ integers)))
+ return failure();
+ if (llvm::is_contained(scalableFlags.asArrayRef(), true))
+ return parser.emitError(parser.getNameLoc())
+ << "scalable indices are not supported";
+ return success();
}
inline ParseResult parsePackedOrDynamicIndexList(
OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> &packed,
@@ -105,9 +113,9 @@ inline ParseResult parsePackedOrDynamicIndexList(
OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> &packed,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
DenseI64ArrayAttr &integers) {
- DenseBoolArrayAttr scalableFlags;
- return parsePackedOrDynamicIndexList(parser, packed, values, scalableFlags,
- integers);
+ Type packedType;
+ return parsePackedOrDynamicIndexList(parser, packed, packedType, values,
+ /*valueTypes=*/nullptr, integers);
}
} // namespace transform
} // namespace mlir
diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.h b/mlir/include/mlir/Interfaces/ViewLikeInterface.h
index c1c2269051efd1..15152750e9db52 100644
--- a/mlir/include/mlir/Interfaces/ViewLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.h
@@ -19,6 +19,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
+#include "llvm/ADT/STLExtras.h"
namespace mlir {
@@ -207,6 +208,8 @@ inline void printDynamicIndexList(
/// 2. `values` is empty.
/// 3. `scalableFlags` is filled with `[false, true, false]`.
///
+/// Overloads that do not return scalable flags reject scalable indices.
+///
ParseResult parseDynamicIndexList(
OpAsmParser &parser,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
@@ -219,8 +222,13 @@ inline ParseResult parseDynamicIndexList(
DenseI64ArrayAttr &integers, SmallVectorImpl<Type> *valueTypes = nullptr,
AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square) {
DenseBoolArrayAttr scalableFlags;
- return parseDynamicIndexList(parser, values, integers, scalableFlags,
- valueTypes, delimiter);
+ if (failed(parseDynamicIndexList(parser, values, integers, scalableFlags,
+ valueTypes, delimiter)))
+ return failure();
+ if (llvm::is_contained(scalableFlags.asArrayRef(), true))
+ return parser.emitError(parser.getNameLoc())
+ << "scalable indices are not supported";
+ return success();
}
/// Verify that a the `values` has as many elements as the number of entries in
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 816d7499845789..dccb5f8a7c8a59 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -3901,24 +3901,22 @@ void transform::TileUsingForOp::getEffects(
// TileUsingForallOp
//===----------------------------------------------------------------------===//
-void transform::TileUsingForallOp::build(OpBuilder &builder,
- OperationState &result, Value target,
- ArrayRef<int64_t> staticTileSizes,
- transform::TileSizesSpec,
- ArrayAttr mapping) {
+void transform::TileUsingForallOp::build(
+ OpBuilder &builder, OperationState &result, Value target,
+ ArrayRef<int64_t> staticTileSizes, transform::TileSizesSpec,
+ ArrayAttr mapping, std::optional<ArrayRef<bool>> scalableSizes) {
return build(builder, result,
/*target=*/target,
/*mixedTileSizes=*/
getAsOpFoldResult(builder.getI64ArrayAttr(staticTileSizes)),
/*_=*/TileSizesSpec(),
- /*mapping=*/mapping);
+ /*mapping=*/mapping, scalableSizes);
}
-void transform::TileUsingForallOp::build(OpBuilder &builder,
- OperationState &result, Value target,
- ArrayRef<OpFoldResult> mixedTileSizes,
- transform::TileSizesSpec,
- ArrayAttr mapping) {
+void transform::TileUsingForallOp::build(
+ OpBuilder &builder, OperationState &result, Value target,
+ ArrayRef<OpFoldResult> mixedTileSizes, transform::TileSizesSpec,
+ ArrayAttr mapping, std::optional<ArrayRef<bool>> scalableSizes) {
SmallVector<int64_t> staticTileSizes;
SmallVector<Value> dynamicTileSizes;
dispatchIndexOpFoldResults(mixedTileSizes, dynamicTileSizes, staticTileSizes);
@@ -3937,6 +3935,8 @@ void transform::TileUsingForallOp::build(OpBuilder &builder,
/*packed_tile_sizes=*/Value(),
/*static_num_threads=*/builder.getDenseI64ArrayAttr({}),
/*static_tile_sizes=*/staticTileSizesAttr,
+ /*scalable_sizes=*/
+ builder.getDenseBoolArrayAttr(scalableSizes.value_or(ArrayRef<bool>{})),
/*mapping=*/mapping);
}
@@ -3974,6 +3974,7 @@ void transform::TileUsingForallOp::build(OpBuilder &builder,
/*packed_tile_sizes=*/Value(),
/*static_num_threads=*/staticNumThreadsAttr,
/*static_tile_sizes=*/builder.getDenseI64ArrayAttr({}),
+ /*scalable_sizes=*/builder.getDenseBoolArrayAttr({}),
/*mapping=*/mapping);
}
@@ -4152,9 +4153,27 @@ DiagnosedSilenceableFailure transform::TileUsingForallOp::apply(
return status;
for (Operation *target : state.getPayloadOps(getTarget())) {
+ // Materialize scalable sizes separately for each target so that the values
+ // dominate the tiled operation, even when targets are in different regions.
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(target);
+ SmallVector<OpFoldResult> tileSizes(mixedTileSizes);
+ Value vscale;
+ for (auto [i, scalable] : llvm::enumerate(getScalableSizes())) {
+ if (!scalable || isZeroInteger(tileSizes[i]))
+ continue;
+ Location loc = target->getLoc();
+ if (!vscale)
+ vscale = vector::VectorScaleOp::create(rewriter, loc,
+ rewriter.getIndexType());
+ Value factor =
+ getValueOrCreateConstantIndexOp(rewriter, loc, tileSizes[i]);
+ tileSizes[i] =
+ arith::MulIOp::create(rewriter, loc, factor, vscale).getResult();
+ }
scf::SCFTilingResult tilingResult;
DiagnosedSilenceableFailure diag = tileToForallOpImpl(
- rewriter, state, transformOp, target, mixedNumThreads, mixedTileSizes,
+ rewriter, state, transformOp, target, mixedNumThreads, tileSizes,
getMapping(), tilingResult);
if (!diag.succeeded())
return diag;
@@ -4191,6 +4210,16 @@ SmallVector<OpFoldResult> TileUsingForallOp::getMixedTileSizes() {
}
LogicalResult TileUsingForallOp::verify() {
+ ArrayRef<bool> scalableSizes = getScalableSizes();
+ if (!scalableSizes.empty() &&
+ scalableSizes.size() != getStaticTileSizes().size())
+ return emitOpError(
+ "expected as many scalable size flags as static tile sizes");
+ for (auto [size, scalable] : llvm::zip(getStaticTileSizes(), scalableSizes)) {
+ if (scalable && (ShapedType::isDynamic(size) || size < 0))
+ return emitOpError(
+ "scalable tile sizes must be non-negative static integers");
+ }
int numThreadsSpec = static_cast<int>(!getMixedNumThreads().empty()) +
static_cast<int>(getPackedNumThreads() != Value());
if (numThreadsSpec > 1)
@@ -4204,6 +4233,8 @@ LogicalResult TileUsingForallOp::verify() {
if (numThreadsSpec == 0 && tileSizesSpec == 0)
return emitOpError("either (packed_)num_threads or (packed_)tile_sizes "
"must be specified");
+ if (numThreadsSpec != 0 && llvm::is_contained(scalableSizes, true))
+ return emitOpError("scalable tile sizes cannot be used with num_threads");
return success();
}
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 9862447824bc94..65f8db017db28f 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -675,7 +675,7 @@ def __init__(
target: Union[Operation, Value, OpView],
*,
num_threads: Optional[MixedValues] = None,
- tile_sizes: MixedValues = None,
+ tile_sizes: Optional[Union[MixedValues, DynamicIndexList]] = None,
mapping=None,
loc=None,
ip=None,
@@ -688,7 +688,7 @@ def __init__(
target: Union[Operation, Value, OpView],
*,
num_threads: Optional[MixedValues] = None,
- tile_sizes: MixedValues = None,
+ tile_sizes: Optional[Union[MixedValues, DynamicIndexList]] = None,
mapping=None,
loc=None,
ip=None,
@@ -704,7 +704,7 @@ def __init__(
target_or_none: Optional[Union[Operation, Value, OpView]] = None,
*,
num_threads: MixedValues = None,
- tile_sizes: MixedValues = None,
+ tile_sizes: Optional[Union[MixedValues, DynamicIndexList]] = None,
mapping=None,
loc=None,
ip=None,
@@ -734,11 +734,23 @@ def __init__(
) = _dispatch_mixed_values(num_threads)
# Unpack mixed tile_sizes.
- (
- dynamic_tile_sizes,
- packed_tile_sizes,
- tile_sizes_attr,
- ) = _dispatch_mixed_values(tile_sizes)
+ packed_tile_sizes = None
+ if not isinstance(tile_sizes, (Operation, Value, OpView)) and any(
+ isinstance(v, Sequence) and not isinstance(v, (str, bytes, ArrayAttr))
+ for v in (tile_sizes or [])
+ ):
+ (
+ dynamic_tile_sizes,
+ tile_sizes_attr,
+ scalable_sizes,
+ ) = _dispatch_dynamic_index_list(tile_sizes)
+ else:
+ (
+ dynamic_tile_sizes,
+ packed_tile_sizes,
+ tile_sizes_attr,
+ ) = _dispatch_mixed_values(tile_sizes)
+ scalable_sizes = []
super().__init__(
loops_type,
@@ -747,6 +759,7 @@ def __init__(
tile_sizes=dynamic_tile_sizes,
packed_tile_sizes=packed_tile_sizes,
static_tile_sizes=tile_sizes_attr,
+ scalable_sizes=scalable_sizes,
num_threads=dynamic_num_threads,
packed_num_threads=packed_num_threads,
static_num_threads=num_threads_attr,
diff --git a/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir b/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir
new file mode 100644
index 00000000000000..6b89d264097df8
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir
@@ -0,0 +1,198 @@
+// RUN: mlir-opt %s --transform-interpreter --test-transform-dialect-erase-schedule --cse --split-input-file | FileCheck %s
+
+// Each target needs its own scalable tile size in the target's scope. Neither
+// extent is divisible by the scalable tile size, so the last tile is bounded.
+// CHECK-DAG: #[[$CEIL130:.+]] = affine_map<()[s0] -> (130 ceildiv s0)>
+// CHECK-DAG: #[[$OFFSET:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$MIN130:.+]] = affine_map<(d0)[s0] -> (-d0 + 130, s0)>
+// CHECK-DAG: #[[$CEIL17:.+]] = affine_map<()[s0] -> (17 ceildiv s0)>
+// CHECK-LABEL: func.func @scalable_tail
+// CHECK-SAME: %[[INPUT:.*]]: tensor<130x32xf32>, %[[INIT:.*]]: tensor<130x32xf32>
+// CHECK: %[[VS0:.*]] = vector.vscale
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[TILE0:.*]] = arith.muli %[[C8]], %[[VS0]] : index
+// CHECK: %[[UB0:.*]] = affine.apply #[[$CEIL130]]()[%[[TILE0]]]
+// CHECK: scf.forall (%[[I:.*]], %[[J:.*]]) in (%[[UB0]], 32) shared_outs(%[[OUT:.*]] = %[[INIT]])
+// CHECK: %[[OFF:.*]] = affine.apply #[[$OFFSET]](%[[I]])[%[[TILE0]]]
+// CHECK: %[[LEN:.*]] = affine.min #[[$MIN130]](%[[OFF]])[%[[TILE0]]]
+// CHECK: tensor.extract_slice %[[INPUT]][%[[OFF]], %[[J]]] [%[[LEN]], 1] [1, 1]
+// CHECK: tensor.extract_slice %[[OUT]][%[[OFF]], %[[J]]] [%[[LEN]], 1] [1, 1]
+// CHECK: linalg.copy
+// CHECK: tensor.parallel_insert_slice %{{.*}} into %[[OUT]][%[[OFF]], %[[J]]] [%[[LEN]], 1] [1, 1]
+func.func @scalable_tail(%input: tensor<130x32xf32>, %init: tensor<130x32xf32>) -> tensor<130x32xf32> {
+ %copy = linalg.copy ins(%input : tensor<130x32xf32>) outs(%init : tensor<130x32xf32>) -> tensor<130x32xf32>
+ return %copy : tensor<130x32xf32>
+}
+
+// CHECK-LABEL: func.func @second_target
+// CHECK: %[[VS1:.*]] = vector.vscale
+// CHECK: %[[C8_1:.*]] = arith.constant 8 : index
+// CHECK: %[[TILE1:.*]] = arith.muli %[[C8_1]], %[[VS1]] : index
+// CHECK: %[[UB1:.*]] = affine.apply #[[$CEIL17]]()[%[[TILE1]]]
+// CHECK: scf.forall (%[[I1:.*]], %{{.*}}) in (%[[UB1]], 32)
+// CHECK: affine.apply #[[$OFFSET]](%[[I1]])[%[[TILE1]]]
+// CHECK: linalg.copy
+func.func @second_target(%input: tensor<17x32xf32>, %init: tensor<17x32xf32>) -> tensor<17x32xf32> {
+ %copy = linalg.copy ins(%input : tensor<17x32xf32>) outs(%init : tensor<17x32xf32>) -> tensor<17x32xf32>
+ return %copy : tensor<17x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %copy = transform.structured.match ops{["linalg.copy"]} in %root : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %copy tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Both dimensions reuse one vscale value.
+// CHECK-LABEL: func.func @two_scalable_dimensions
+// CHECK: %[[VS:.*]] = vector.vscale
+// CHECK: %[[TILE0:.*]] = arith.muli %{{.*}}, %[[VS]]
+// CHECK: %[[TILE1:.*]] = arith.muli %{{.*}}, %[[VS]]
+// CHECK-NOT: vector.vscale
+// CHECK: %[[UB0:.*]] = affine.apply {{.*}}()[%[[TILE0]]]
+// CHECK: %[[UB1:.*]] = affine.apply {{.*}}()[%[[TILE1]]]
+// CHECK: scf.forall (%[[I:.*]], %[[J:.*]]) in (%[[UB0]], %[[UB1]])
+// CHECK: %[[OFF0:.*]] = affine.apply {{.*}}(%[[I]])[%[[TILE0]]]
+// CHECK: %[[OFF1:.*]] = affine.apply {{.*}}(%[[J]])[%[[TILE1]]]
+// CHECK: %[[LEN0:.*]] = affine.min {{.*}}(%[[OFF0]])[%[[TILE0]]]
+// CHECK: %[[LEN1:.*]] = affine.min {{.*}}(%[[OFF1]])[%[[TILE1]]]
+// CHECK: tensor.extract_slice %{{.*}}[%[[OFF0]], %[[OFF1]]] [%[[LEN0]], %[[LEN1]]] [1, 1]
+func.func @two_scalable_dimensions(%input: tensor<130x33xf32>, %init: tensor<130x33xf32>) -> tensor<130x33xf32> {
+ %copy = linalg.copy ins(%input : tensor<130x33xf32>) outs(%init : tensor<130x33xf32>) -> tensor<130x33xf32>
+ return %copy : tensor<130x33xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %copy = transform.structured.match ops{["linalg.copy"]} in %root : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %copy tile_sizes [[8], [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Scalable sizes may be mixed with ordinary dynamic sizes.
+// CHECK-LABEL: func.func @mixed_dynamic
+// CHECK: %[[DYNAMIC:.*]] = call @tile_size
+// CHECK: %[[VS:.*]] = vector.vscale
+// CHECK: %[[TILE:.*]] = arith.muli %{{.*}}, %[[VS]]
+// CHECK: %[[UB0:.*]] = affine.apply {{.*}}()[%[[TILE]]]
+// CHECK: %[[UB1:.*]] = affine.apply {{.*}}()[%[[DYNAMIC]]]
+// CHECK: scf.forall (%[[I:.*]], %[[J:.*]]) in (%[[UB0]], %[[UB1]])
+// CHECK: %[[OFF0:.*]] = affine.apply {{.*}}(%[[I]])[%[[TILE]]]
+// CHECK: %[[OFF1:.*]] = affine.apply {{.*}}(%[[J]])[%[[DYNAMIC]]]
+// CHECK: %[[LEN0:.*]] = affine.min {{.*}}(%[[OFF0]])[%[[TILE]]]
+// CHECK: %[[LEN1:.*]] = affine.min {{.*}}(%[[OFF1]])[%[[DYNAMIC]]]
+// CHECK: tensor.extract_slice %{{.*}}[%[[OFF0]], %[[OFF1]]] [%[[LEN0]], %[[LEN1]]] [1, 1]
+func.func @mixed_dynamic(%input: tensor<130x33xf32>, %init: tensor<130x33xf32>) -> tensor<130x33xf32> {
+ %size = func.call @tile_size() : () -> index
+ %copy = linalg.copy ins(%input : tensor<130x33xf32>) outs(%init : tensor<130x33xf32>) -> tensor<130x33xf32>
+ return %copy : tensor<130x33xf32>
+}
+func.func private @tile_size() -> index
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %size = transform.structured.match ops{["func.call"]} in %root : (!transform.any_op) -> !transform.any_op
+ %copy = transform.structured.match ops{["linalg.copy"]} in %root : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %copy tile_sizes [[8], %size] : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Zero remains untiled even when marked scalable.
+// CHECK-LABEL: func.func @scalable_zero
+// CHECK-NOT: vector.vscale
+// CHECK: scf.forall (%{{.*}}) in (8)
+// CHECK: tensor.extract_slice %{{.*}}[0, %{{.*}}] [130, 4] [1, 1]
+func.func @scalable_zero(%input: tensor<130x32xf32>, %init: tensor<130x32xf32>) -> tensor<130x32xf32> {
+ %copy = linalg.copy ins(%input : tensor<130x32xf32>) outs(%init : tensor<130x32xf32>) -> tensor<130x32xf32>
+ return %copy : tensor<130x32xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %copy = transform.structured.match ops{["linalg.copy"]} in %root : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %copy tile_sizes [[0], 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Dynamic extents, buffer semantics and device mapping remain supported.
+// CHECK-LABEL: func.func @scalable_memref
+// CHECK: vector.vscale
+// CHECK: arith.muli
+// CHECK: memref.dim
+// CHECK: scf.forall (%{{.*}})
+// CHECK: memref.subview
+// CHECK: linalg.copy
+// CHECK: mapping = [#gpu.thread<x>]
+func.func @scalable_memref(%input: memref<?x32xf32>, %init: memref<?x32xf32>) {
+ linalg.copy ins(%input : memref<?x32xf32>) outs(%init : memref<?x32xf32>)
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %copy = transform.structured.match ops{["linalg.copy"]} in %root : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %copy tile_sizes [[8], 0] (mapping = [#gpu.thread<x>]) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// The pack inner tile and the forall tile size each scale independently.
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[$PACK_OFFSET:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$PACK_MIN:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s0, s1)>
+// CHECK-DAG: #[[$SOURCE_OFFSET:.+]] = affine_map<(d0)[s0, s1] -> ((d0 * s1) * s0)>
+// CHECK-LABEL: func.func @pack_scalable_prod
+// CHECK-SAME: %[[SOURCE:.*]]: tensor<1024x32xf32>
+// CHECK: %[[EIGHT:.*]] = arith.constant 8 : index
+// CHECK: %[[SCALE:.*]] = vector.vscale
+// CHECK: %[[SCALED_TILE:.*]] = arith.muli %[[SCALE]], %[[EIGHT]] : index
+// CHECK: scf.forall (%[[PI:.*]], %[[PJ:.*]]) in (%{{.*}}, 32) shared_outs(%[[DEST:.*]] = %{{.*}})
+// CHECK: %[[POFF:.*]] = affine.apply #[[$PACK_OFFSET]](%[[PI]])[%[[SCALED_TILE]]]
+// CHECK: %[[PLEN:.*]] = affine.min #[[$PACK_MIN]](%[[POFF]])[%{{.*}}, %[[SCALED_TILE]]]
+// CHECK: %[[SOFF:.*]] = affine.apply #[[$SOURCE_OFFSET]](%[[PI]])[%[[SCALED_TILE]], %[[SCALED_TILE]]]
+// CHECK: %[[SLEN:.*]] = affine.apply #[[$PACK_OFFSET]](%[[PLEN]])[%[[SCALED_TILE]]]
+// CHECK: tensor.extract_slice %[[SOURCE]][%[[SOFF]], %[[PJ]]] [%[[SLEN]], 1] [1, 1]
+// CHECK: %[[PRODUCER:.*]] = linalg.generic
+// CHECK: %[[PACK_DEST:.*]] = tensor.extract_slice %[[DEST]][%[[POFF]], %[[PJ]], 0, 0] [%[[PLEN]], 1, %{{.*}}, 1] [1, 1, 1, 1]
+// CHECK: %[[PACK:.*]] = linalg.pack %[[PRODUCER]] inner_dims_pos = [0, 1] inner_tiles = [%[[SCALED_TILE]], 1] into %[[PACK_DEST]]
+// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[DEST]][%[[POFF]], %[[PJ]], 0, 0] [%[[PLEN]], 1, %{{.*}}, 1] [1, 1, 1, 1]
+func.func @pack_scalable_prod(%input: tensor<1024x32xf32>) -> tensor<?x32x?x1xf32> {
+ %c8 = arith.constant 8 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %vscale, %c8 : index
+ %outer_size = affine.apply affine_map<()[s0] -> (1024 ceildiv s0)>()[%c8_vscale]
+ %pack_init = tensor.empty(%outer_size, %c8_vscale) : tensor<?x32x?x1xf32>
+ %producer_init = tensor.empty() : tensor<1024x32xf32>
+ %producer = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%input : tensor<1024x32xf32>) outs(%producer_init : tensor<1024x32xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %sum = arith.addf %in, %in : f32
+ linalg.yield %sum : f32
+ } -> tensor<1024x32xf32>
+ %pack = linalg.pack %producer inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, 1] into %pack_init : tensor<1024x32xf32> -> tensor<?x32x?x1xf32>
+ return %pack : tensor<?x32x?x1xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %module : (!transform.any_op) -> !transform.any_op
+ %pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op
+ %tiled_pack, %loops = transform.structured.tile_using_forall %pack tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %fused_op, %new_containing_op = transform.structured.fuse_into_containing_op %generic into %loops : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir b/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
index 3c7f8895d2db10..4c3d4cc4322752 100644
--- a/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/Linalg/transform-ops-invalid.mlir
@@ -1,5 +1,81 @@
// RUN: mlir-opt %s --split-input-file --verify-diagnostics
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{expected as many scalable size flags as static tile sizes}}
+ %tiled:2 = "transform.structured.tile_using_forall"(%arg0) <{operandSegmentSizes = array<i32: 1, 0, 0, 0, 0>, static_tile_sizes = array<i64: 8, 1>, scalable_sizes = array<i1: true>}> : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ %size = transform.param.constant 8 : i64 -> !transform.param<i64>
+ // expected-error at below {{scalable tile sizes must be non-negative static integers}}
+ %tiled:2 = "transform.structured.tile_using_forall"(%arg0, %size) <{operandSegmentSizes = array<i32: 1, 0, 1, 0, 0>, static_tile_sizes = array<i64: -9223372036854775808>, scalable_sizes = array<i1: true>}> : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{scalable tile sizes must be non-negative static integers}}
+ %tiled, %loop = transform.structured.tile_using_forall %arg0 tile_sizes [[-8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ %size = transform.param.constant 8 : i64 -> !transform.param<i64>
+ // expected-error at below {{expected as many scalable size flags as static tile sizes}}
+ %tiled:2 = "transform.structured.tile_using_forall"(%arg0, %size) <{operandSegmentSizes = array<i32: 1, 0, 0, 0, 1>, scalable_sizes = array<i1: true>}> : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{scalable tile sizes cannot be used with num_threads}}
+ %tiled, %loop = transform.structured.tile_using_forall %arg0 num_threads [2, 4] tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+// Scalable indices must not be silently treated as fixed indices by transforms
+// that do not support scalable thread counts or interchange.
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{scalable indices are not supported}}
+ %tiled, %loop = transform.structured.tile_using_forall %arg0 num_threads [2, [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{scalable indices are not supported}}
+ %tiled, %loops:2 = transform.structured.tile_using_for %arg0 tile_sizes [8, 4] interchange = [0, [1]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{scalable indices are not supported}}
+ %tiled, %loops:2 = transform.structured.fuse %arg0 tile_sizes [4, [8]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // expected-error at below {{scalable indices are not supported}}
+ %tiled, %loops:2 = transform.structured.fuse %arg0 tile_sizes [8, 4] interchange [0, [1]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+// -----
+
transform.sequence failures(propagate) {
^bb0(%arg0: !transform.any_op):
// expected-error at below {{'transform.structured.interchange' op expects iterator_interchange to be a permutation, found 1, 1}}
diff --git a/mlir/test/Dialect/Linalg/transform-ops.mlir b/mlir/test/Dialect/Linalg/transform-ops.mlir
index fd989899dd0632..b0be5a6c2e6d75 100644
--- a/mlir/test/Dialect/Linalg/transform-ops.mlir
+++ b/mlir/test/Dialect/Linalg/transform-ops.mlir
@@ -6,6 +6,39 @@ transform.sequence failures(propagate) {
%0, %1:2 = transform.structured.tile_using_for %arg0 tile_sizes [2, 0, 3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
}
+// Scalable tile sizes round-trip for transforms that support them.
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // CHECK: transform.structured.tile_using_for %arg0 tile_sizes {{\[\[4\], 8\]}}
+ %tiled, %loops:2 = transform.structured.tile_using_for %arg0 tile_sizes [[4], 8] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // CHECK: transform.structured.tile_using_forall %arg0 tile_sizes {{\[\[8\], 1\]}}
+ %tiled, %loop = transform.structured.tile_using_forall %arg0 tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// Generic IR without scalable_sizes retains fixed tile size semantics.
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // CHECK: transform.structured.tile_using_forall %arg0 tile_sizes [8, 1]
+ %tiled:2 = "transform.structured.tile_using_forall"(%arg0) <{operandSegmentSizes = array<i32: 1, 0, 0, 0, 0>, static_tile_sizes = array<i64: 8, 1>}> : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
+// Fixed, dynamic and packed index lists remain supported without scalable flags.
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+ // CHECK: %[[SIZE:.*]] = transform.param.constant 8
+ %size = transform.param.constant 8 : i64 -> !transform.param<i64>
+ // CHECK: transform.structured.tile_using_forall %arg0 tile_sizes [%[[SIZE]], 1]
+ %tiled, %loop = transform.structured.tile_using_forall %arg0 tile_sizes [%size, 1] : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op)
+ // CHECK: transform.structured.tile_using_forall %{{.*}} tile_sizes *(%[[SIZE]])
+ %packed, %packed_loop = transform.structured.tile_using_forall %tiled tile_sizes *(%size) : (!transform.any_op, !transform.param<i64>) -> (!transform.any_op, !transform.any_op)
+ // CHECK: transform.structured.tile_using_forall %{{.*}} num_threads [2, 4]
+ %fixed, %fixed_loop = transform.structured.tile_using_forall %packed num_threads [2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}
+
// check that the Attributes of `tile_using_for` are preserved through printing
// and parsing with and without use of the optional `interchange` Attribute.
transform.sequence failures(propagate) {
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index f613d309f1db45..96e47c35e8095d 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -1,5 +1,29 @@
// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+func.func @extract_slice_scalable_offset(%source: tensor<16xf32>) {
+ // expected-error at below {{scalable indices are not supported}}
+ %0 = tensor.extract_slice %source[[4]] [8] [1] : tensor<16xf32> to tensor<8xf32>
+ return
+}
+
+// -----
+
+func.func @extract_slice_scalable_size(%source: tensor<16xf32>) {
+ // expected-error at below {{scalable indices are not supported}}
+ %0 = tensor.extract_slice %source[4] [[8]] [1] : tensor<16xf32> to tensor<8xf32>
+ return
+}
+
+// -----
+
+func.func @extract_slice_scalable_stride(%source: tensor<16xf32>) {
+ // expected-error at below {{scalable indices are not supported}}
+ %0 = tensor.extract_slice %source[4] [8] [[1]] : tensor<16xf32> to tensor<8xf32>
+ return
+}
+
+// -----
+
// Asking the dimension of a 0-D shape doesn't make sense.
func.func @dim_0_ranked(%arg : tensor<f32>, %arg1 : index) {
tensor.dim %arg, %arg1 : tensor<f32> // expected-error {{'tensor.dim' op operand #0 must be non-0-ranked or unranked tensor, but got 'tensor<f32>'}}
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index b96c8edb177c41..1a8a115bf58aac 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -685,6 +685,35 @@ def testTileToForallTileSizes(target):
# CHECK-SAME: tile_sizes [2, 3, 4]
+ at run
+ at create_sequence
+def testTileToForallScalableTileSizes(target):
+ structured.TileUsingForallOp(target, tile_sizes=[[8], 1])
+ # CHECK-LABEL: TEST: testTileToForallScalableTileSizes
+ # CHECK: = transform.structured.tile_using_forall
+ # CHECK-SAME: tile_sizes {{\[\[8\], 1\]}}
+
+
+ at run
+ at create_sequence
+def testTileToForallScalableMixedDynamic(target):
+ n = structured.MatchOp.match_op_names(target, ["test.dummy"])
+ structured.TileUsingForallOp(target, tile_sizes=[[8], n, 0])
+ # CHECK-LABEL: TEST: testTileToForallScalableMixedDynamic
+ # CHECK: = transform.structured.tile_using_forall
+ # CHECK-SAME: tile_sizes {{\[\[8\],}} %{{.*}}, 0]
+
+
+ at run
+ at create_sequence
+def testTileToForallPackedTileSizes(target):
+ n = structured.MatchOp.match_op_names(target, ["test.dummy"])
+ structured.TileUsingForallOp(target, tile_sizes=n)
+ # CHECK-LABEL: TEST: testTileToForallPackedTileSizes
+ # CHECK: = transform.structured.tile_using_forall
+ # CHECK-SAME: tile_sizes *(%0) : (!transform.any_op, !transform.any_op)
+
+
@run
@create_sequence
def testTileToForallMixedDynamic(target):
>From b35a774fd8f048ecfbbd38ba09f074661ae99b21 Mon Sep 17 00:00:00 2001
From: MiaoMing Chen <miaomingc at stumail.nwu.edu.cn>
Date: Mon, 26 Oct 2026 20:32:22 +0800
Subject: [PATCH 2/3] fixup! [mlir] Fix scalable index handling and support
scalable forall tiling
Delay scalable tile-size materialization until SCF requests it, after the
forall helper has checked the target's TilingInterface and tensor semantics.
Suppressed failures on unsupported targets then leave no scalable-size IR.
Add regressions for non-tileable operations and memref pack/unpack without
running CSE or canonicalization, so dead materialized values remain visible.
Assisted-by: OpenAI Codex (6.1-sol)
---
.../Linalg/TransformOps/LinalgTransformOps.h | 14 ++--
.../TransformOps/LinalgTransformOps.cpp | 45 ++++++------
.../tile-to-forall-scalable-failures.mlir | 71 +++++++++++++++++++
.../Linalg/tile-to-forall-scalable.mlir | 2 +-
4 files changed, 101 insertions(+), 31 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/tile-to-forall-scalable-failures.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
index b3468dd0027ad8..1da76511f2d843 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
@@ -54,14 +54,12 @@ class DialectRegistry;
namespace transform {
/// Implementation of tiling operations using `scf.forall`.
-DiagnosedSilenceableFailure
-tileToForallOpImpl(RewriterBase &rewriter, transform::TransformState &state,
- TransformOpInterface transformOp, Operation *target,
- ArrayRef<OpFoldResult> mixedNumThreads,
- ArrayRef<OpFoldResult> mixedTileSizes,
- std::optional<ArrayAttr> mapping,
- scf::SCFTilingResult &tilingResult);
-
+DiagnosedSilenceableFailure tileToForallOpImpl(
+ RewriterBase &rewriter, transform::TransformState &state,
+ TransformOpInterface transformOp, Operation *target,
+ ArrayRef<OpFoldResult> mixedNumThreads,
+ ArrayRef<OpFoldResult> mixedTileSizes, std::optional<ArrayAttr> mapping,
+ scf::SCFTilingResult &tilingResult, ArrayRef<bool> scalableSizes = {});
} // namespace transform
} // namespace mlir
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index dccb5f8a7c8a59..9b19fef5b61b77 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4069,7 +4069,7 @@ DiagnosedSilenceableFailure transform::tileToForallOpImpl(
TransformOpInterface transformOp, Operation *target,
ArrayRef<OpFoldResult> mixedNumThreads,
ArrayRef<OpFoldResult> mixedTileSizes, std::optional<ArrayAttr> mapping,
- scf::SCFTilingResult &tilingResult) {
+ scf::SCFTilingResult &tilingResult, ArrayRef<bool> scalableSizes) {
// Transform all targets one by one.
auto tileableOp = dyn_cast<TilingInterface>(target);
if (!tileableOp) {
@@ -4094,7 +4094,26 @@ DiagnosedSilenceableFailure transform::tileToForallOpImpl(
if (!mixedNumThreads.empty()) {
options.setNumThreads(mixedNumThreads);
} else {
- options.setTileSizes(mixedTileSizes);
+ // Materialize scalable sizes only after checking the target's
+ // preconditions. The tiler invokes this callback in each target's scope.
+ options.setTileSizeComputationFunction(
+ [mixedTileSizes, scalableSizes](OpBuilder &builder, Operation *op) {
+ SmallVector<OpFoldResult> tileSizes(mixedTileSizes);
+ Value vscale;
+ for (auto [i, scalable] : llvm::enumerate(scalableSizes)) {
+ if (!scalable || isZeroInteger(tileSizes[i]))
+ continue;
+ Location loc = op->getLoc();
+ if (!vscale)
+ vscale = vector::VectorScaleOp::create(builder, loc,
+ builder.getIndexType());
+ Value factor =
+ getValueOrCreateConstantIndexOp(builder, loc, tileSizes[i]);
+ tileSizes[i] =
+ arith::MulIOp::create(builder, loc, factor, vscale).getResult();
+ }
+ return tileSizes;
+ });
}
if (mapping) {
options.setMapping(mapping.value().getValue());
@@ -4153,28 +4172,10 @@ DiagnosedSilenceableFailure transform::TileUsingForallOp::apply(
return status;
for (Operation *target : state.getPayloadOps(getTarget())) {
- // Materialize scalable sizes separately for each target so that the values
- // dominate the tiled operation, even when targets are in different regions.
- OpBuilder::InsertionGuard guard(rewriter);
- rewriter.setInsertionPoint(target);
- SmallVector<OpFoldResult> tileSizes(mixedTileSizes);
- Value vscale;
- for (auto [i, scalable] : llvm::enumerate(getScalableSizes())) {
- if (!scalable || isZeroInteger(tileSizes[i]))
- continue;
- Location loc = target->getLoc();
- if (!vscale)
- vscale = vector::VectorScaleOp::create(rewriter, loc,
- rewriter.getIndexType());
- Value factor =
- getValueOrCreateConstantIndexOp(rewriter, loc, tileSizes[i]);
- tileSizes[i] =
- arith::MulIOp::create(rewriter, loc, factor, vscale).getResult();
- }
scf::SCFTilingResult tilingResult;
DiagnosedSilenceableFailure diag = tileToForallOpImpl(
- rewriter, state, transformOp, target, mixedNumThreads, tileSizes,
- getMapping(), tilingResult);
+ rewriter, state, transformOp, target, mixedNumThreads, mixedTileSizes,
+ getMapping(), tilingResult, getScalableSizes());
if (!diag.succeeded())
return diag;
if (!tilingResult.loops.empty())
diff --git a/mlir/test/Dialect/Linalg/tile-to-forall-scalable-failures.mlir b/mlir/test/Dialect/Linalg/tile-to-forall-scalable-failures.mlir
new file mode 100644
index 00000000000000..6413963ff23b3c
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/tile-to-forall-scalable-failures.mlir
@@ -0,0 +1,71 @@
+// RUN: mlir-opt %s --transform-interpreter --test-transform-dialect-erase-schedule --split-input-file | FileCheck %s
+
+// Suppressing a target precondition failure must not leave materialized tile
+// sizes behind. Do not run CSE or canonicalization: they could erase that IR.
+// CHECK-LABEL: func.func @unsupported_target
+// CHECK-NEXT: %[[SUM:.*]] = arith.addi
+// CHECK-NEXT: return %[[SUM]]
+// CHECK-NEXT: }
+func.func @unsupported_target(%lhs: i32, %rhs: i32) -> i32 {
+ %sum = arith.addi %lhs, %rhs : i32
+ return %sum : i32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ transform.sequence %root : !transform.any_op failures(suppress) {
+ ^bb0(%scope: !transform.any_op):
+ %target = transform.structured.match ops{["arith.addi"]} in %scope : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %target tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.yield
+ }
+}
+
+// -----
+
+// Memref pack implements TilingInterface, but its semantics are unsupported.
+// CHECK-LABEL: func.func @unsupported_memref_pack
+// CHECK-NEXT: linalg.pack
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @unsupported_memref_pack(%source: memref<128x256xf32>, %dest: memref<4x8x32x32xf32>) {
+ linalg.pack %source inner_dims_pos = [0, 1] inner_tiles = [32, 32] into %dest : memref<128x256xf32> -> memref<4x8x32x32xf32>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ transform.sequence %root : !transform.any_op failures(suppress) {
+ ^bb0(%scope: !transform.any_op):
+ %target = transform.structured.match ops{["linalg.pack"]} in %scope : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %target tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unsupported_memref_unpack
+// CHECK-NEXT: linalg.unpack
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+func.func @unsupported_memref_unpack(%source: memref<4x8x32x32xf32>, %dest: memref<128x256xf32>) {
+ linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 32] into %dest : memref<4x8x32x32xf32> -> memref<128x256xf32>
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ transform.sequence %root : !transform.any_op failures(suppress) {
+ ^bb0(%scope: !transform.any_op):
+ %target = transform.structured.match ops{["linalg.unpack"]} in %scope : (!transform.any_op) -> !transform.any_op
+ %tiled, %loop = transform.structured.tile_using_forall %target tile_sizes [[8], 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir b/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir
index 6b89d264097df8..707652e3a2fd81 100644
--- a/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir
+++ b/mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir
@@ -129,9 +129,9 @@ module attributes {transform.with_named_sequence} {
// Dynamic extents, buffer semantics and device mapping remain supported.
// CHECK-LABEL: func.func @scalable_memref
+// CHECK: memref.dim
// CHECK: vector.vscale
// CHECK: arith.muli
-// CHECK: memref.dim
// CHECK: scf.forall (%{{.*}})
// CHECK: memref.subview
// CHECK: linalg.copy
>From 09bae43be874dacf70d27a929a1b50d7b5b997d5 Mon Sep 17 00:00:00 2001
From: MiaoMing Chen <miaomingc at stumail.nwu.edu.cn>
Date: Mon, 26 Oct 2026 21:57:11 +0800
Subject: [PATCH 3/3] fixup! [mlir] Fix scalable index handling and support
scalable forall tiling
Cover scalable-marker rejection at the direct parser callers outside the
Transform dialect: all four scf.forall bound-list forms and linalg.pack /
linalg.unpack inner tile lists.
Keep the fixed and dynamic-SSA forms covered by their existing positive
parser and scalable pack/unpack tests.
Assisted-by: OpenAI Codex (6.1-sol)
---
mlir/test/Dialect/Linalg/invalid.mlir | 16 ++++++++++++
mlir/test/Dialect/SCF/invalid.mlir | 36 +++++++++++++++++++++++++++
2 files changed, 52 insertions(+)
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 70b375ef24941b..ce5d9ca70cd291 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -1,5 +1,21 @@
// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+func.func @pack_scalable_inner_tile(%source: tensor<16x16xf32>, %dest: tensor<2x4x8x4xf32>) {
+ // expected-error at below {{scalable indices are not supported}}
+ %0 = linalg.pack %source inner_dims_pos = [0, 1] inner_tiles = [[8], 4] into %dest : tensor<16x16xf32> -> tensor<2x4x8x4xf32>
+ return
+}
+
+// -----
+
+func.func @unpack_scalable_inner_tile(%source: tensor<2x4x8x4xf32>, %dest: tensor<16x16xf32>) {
+ // expected-error at below {{scalable indices are not supported}}
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [[8], 4] into %dest : tensor<2x4x8x4xf32> -> tensor<16x16xf32>
+ return
+}
+
+// -----
+
func.func @load_number_of_indices(%v : memref<f32>) {
// expected-error @+2 {{invalid number of indices for accessed memref, expected 0 but got 1}}
%c0 = arith.constant 0 : index
diff --git a/mlir/test/Dialect/SCF/invalid.mlir b/mlir/test/Dialect/SCF/invalid.mlir
index a81e912ab35a9c..66cc12b7f64d35 100644
--- a/mlir/test/Dialect/SCF/invalid.mlir
+++ b/mlir/test/Dialect/SCF/invalid.mlir
@@ -1,5 +1,41 @@
// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -verify-diagnostics
+func.func @forall_scalable_normalized_upper_bound() {
+ // expected-error at below {{scalable indices are not supported}}
+ scf.forall (%i, %j) in ([8], 4) {
+ }
+ return
+}
+
+// -----
+
+func.func @forall_scalable_lower_bound() {
+ // expected-error at below {{scalable indices are not supported}}
+ scf.forall (%i, %j) = ([2], 0) to (8, 4) step (1, 1) {
+ }
+ return
+}
+
+// -----
+
+func.func @forall_scalable_upper_bound() {
+ // expected-error at below {{scalable indices are not supported}}
+ scf.forall (%i, %j) = (0, 0) to ([8], 4) step (1, 1) {
+ }
+ return
+}
+
+// -----
+
+func.func @forall_scalable_step() {
+ // expected-error at below {{scalable indices are not supported}}
+ scf.forall (%i, %j) = (0, 0) to (8, 4) step ([2], 1) {
+ }
+ return
+}
+
+// -----
+
func.func @loop_for_lb(%arg0: f32, %arg1: index) {
// expected-error at +1 {{operand #0 must be signless integer or index}}
"scf.for"(%arg0, %arg1, %arg1) ({}) : (f32, index, index) -> ()
More information about the Mlir-commits
mailing list