[Mlir-commits] [mlir] [mlir] Fix scalable index handling and support scalable forall tiling (PR #235596)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 25 22:33:18 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: MiaoMing Chen (chenmiaoming)
<details>
<summary>Changes</summary>
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
---
Patch is 36.16 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/235596.diff
10 Files Affected:
- (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+12-3)
- (modified) mlir/include/mlir/Dialect/Transform/Utils/Utils.h (+14-6)
- (modified) mlir/include/mlir/Interfaces/ViewLikeInterface.h (+10-2)
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+43-12)
- (modified) mlir/python/mlir/dialects/transform/structured.py (+21-8)
- (added) mlir/test/Dialect/Linalg/tile-to-forall-scalable.mlir (+198)
- (modified) mlir/test/Dialect/Linalg/transform-ops-invalid.mlir (+76)
- (modified) mlir/test/Dialect/Linalg/transform-ops.mlir (+33)
- (modified) mlir/test/Dialect/Tensor/invalid.mlir (+24)
- (modified) mlir/test/python/dialects/transform_structured_ext.py (+29)
``````````diff
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.mu...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/235596
More information about the Mlir-commits
mailing list