[Mlir-commits] [mlir] [mlir] Fix scalable index handling and support scalable forall tiling (PR #235596)
MiaoMing Chen
llvmlistbot at llvm.org
Sun Oct 25 22:32:39 PDT 2026
https://github.com/chenmiaoming created https://github.com/llvm/llvm-project/pull/235596
Fixes https://github.com/llvm/llvm-project/issues/150203.
`tile_using_forall` accepts `tile_sizes [[8], 1]`, but the parser drops the scalable flag and prints it back as `[8, 1]`. At `vscale = 2`, this requests a tile size of 8 instead of 16. The bug changes the requested tiling even when numerical results remain correct.
Preserve the flags in `scalable_sizes` and materialize each nonzero scalable factor times `vector.vscale` before calling the SCF tiler. Materialization happens separately for each payload target to preserve SSA dominance across functions. Zero sizes remain untiled, and omitted flags retain fixed-size behavior.
Parser overloads that cannot return scalable flags now diagnose scalable indices instead of discarding them. Scalable thread counts and packed tile size handles remain unsupported. The C++ and Python builders accept scalable tile sizes.
Regression tests cover parsing/printing, invalid specifications, legacy generic IR, Python builders, multiple targets, tail slices, mixed dynamic sizes, and pack tiling with producer fusion. For pack, the forall tile size is in the packed outer iteration space; it scales independently of the pack's inner tile.
Validation:
- Assertions-enabled build; 384 dialect tests and 20 Transform Python tests pass.
- The original reproducer succeeds. Fixed/scalable IR specialization checks pass at `vscale = 1, 2, 4`.
- QEMU AArch64: 15 numerical cases pass across `vscale = 1, 2, 4`.
- Native Redmi K70E (Termux): 5 numerical cases pass at `vscale = 1`, with zero mismatches.
Execution checks serialize forall loops and check padding and partial tiles. Wider vector lengths were tested under QEMU only.
Assisted-by: OpenAI Codex
>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] [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):
More information about the Mlir-commits
mailing list