[Codegen][CPU] Fold reshapes into map_store per reassociation group. (#24535)
`foldReshapeIntoMapStore` composed an expand/collapse reshape into a
`map_store`'s index transformation via one global `linearize` over the
source shape plus `delinearize` over the result shape. Under dynamic
shapes that `delinearize` lowers to a dynamic integer division per
element.
Fold each reassociation group independently instead: an `expand_shape`
splits one source index by just that group's result sizes; a
`collapse_shape` merges one group of source indices. The split/merge
factors are then per-group -- typically static tile sizes -- so the
scalarized transformation is static shifts/masks with no per-element
division.
Speeds up the i8 data-tiled matmul's C-encode dispatch; i8 inner_tiled
matmul, 4096x4096 on Zen 4: 52.9ms -> 46.9ms.
Also document, at `MapStoreOpVectorizationModel::isVectorizable`, why a
CPU-tiled `map_store` reaches the vectorizer dynamically shaped and
scalarizes -- which is precisely what makes this per-group,
division-free fold matter -- while the same op vectorizes on GPU.
Comment only, no behavior change.
Progress towards #24515.
Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp b/compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp
index 6fbebb6..9db8fcb 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp
@@ -77,8 +77,16 @@
}
/// Fold a tensor::ExpandShapeOp or tensor::CollapseShapeOp into a consumer
-/// `mapStoreOp`, by linearizing and then delinearizing the source indices
-/// of the `mapStoreOp`s index transformation.
+/// `mapStoreOp` by composing the reshape's index transformation into the
+/// front of the `mapStoreOp`'s.
+///
+/// Each reassociation group is folded independently: an `expand_shape` splits
+/// one source index via a `delinearize` over just that group's result sizes,
+/// and a `collapse_shape` merges a group of source indices via a `linearize`.
+/// This keeps the split/merge factors per-group — typically static tile sizes
+/// — rather than routing every index through one global `linearize` +
+/// `delinearize` over the full (dynamic) shapes, which would force a dynamic
+/// integer division per element when the transformation is later scalarized.
template <typename ReshapeOpTy>
static FailureOr<IREE::LinalgExt::MapStoreOp>
foldReshapeIntoMapStore(RewriterBase &rewriter, ReshapeOpTy reshapeOp,
@@ -92,6 +100,7 @@
cast<RankedTensorType>(reshapeOp.getResult().getType()).getRank() == 0) {
return failure();
}
+ constexpr bool isExpand = std::is_same_v<ReshapeOpTy, tensor::ExpandShapeOp>;
Location loc = reshapeOp->getLoc();
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPointAfter(reshapeOp);
@@ -102,20 +111,53 @@
// sizes by later cleanup patterns.
SmallVector<OpFoldResult> resultDims =
tensor::getMixedSizes(rewriter, loc, reshapeOp.getResult());
+ SmallVector<ReassociationIndices> reassociation =
+ reshapeOp.getReassociationIndices();
Location mapStoreLoc = mapStoreOp->getLoc();
rewriter.modifyOpInPlace(mapStoreOp, [&]() {
mapStoreOp.insertTransformationAtStart(
rewriter,
- [&rewriter, mapStoreLoc, srcDims,
- resultDims](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
- SmallVector<Value> indexValues(indices.begin(), indices.end());
- auto linearizeIndexOp = affine::AffineLinearizeIndexOp::create(
- rewriter, mapStoreLoc, indexValues, srcDims, /*disjoint=*/true);
- auto delinearizeIndexOp = affine::AffineDelinearizeIndexOp::create(
- rewriter, mapStoreLoc, linearizeIndexOp.getResult(), resultDims,
- /*hasOuterBound=*/true);
- return delinearizeIndexOp->getResults();
+ [&rewriter, mapStoreLoc, srcDims, resultDims, reassociation,
+ isExpand](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
+ // `indices` are in the reshape's source index space; produce the
+ // reshape's result index space.
+ SmallVector<Value> result(resultDims.size());
+ for (auto [groupIdx, group] : llvm::enumerate(reassociation)) {
+ if (isExpand) {
+ // Source index `groupIdx` splits into result dims `group`.
+ if (group.size() == 1) {
+ result[group[0]] = indices[groupIdx];
+ continue;
+ }
+ SmallVector<OpFoldResult> groupSizes;
+ for (int64_t resultDim : group) {
+ groupSizes.push_back(resultDims[resultDim]);
+ }
+ auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
+ rewriter, mapStoreLoc, indices[groupIdx], groupSizes,
+ /*hasOuterBound=*/true);
+ for (auto [i, resultDim] : llvm::enumerate(group)) {
+ result[resultDim] = delinearizeOp.getResult(i);
+ }
+ } else {
+ // Source dims `group` merge into result index `groupIdx`.
+ if (group.size() == 1) {
+ result[groupIdx] = indices[group[0]];
+ continue;
+ }
+ SmallVector<Value> groupIndices;
+ SmallVector<OpFoldResult> groupSizes;
+ for (int64_t srcDim : group) {
+ groupIndices.push_back(indices[srcDim]);
+ groupSizes.push_back(srcDims[srcDim]);
+ }
+ result[groupIdx] = affine::AffineLinearizeIndexOp::create(
+ rewriter, mapStoreLoc, groupIndices, groupSizes,
+ /*disjoint=*/true);
+ }
+ }
+ return result;
},
srcDims.size());
mapStoreOp.getInputMutable().assign(reshapeOp->getOperand(0));
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/combine_result_layout_transformation.mlir b/compiler/src/iree/compiler/Codegen/Common/test/combine_result_layout_transformation.mlir
index 3284f32..872d11b 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/combine_result_layout_transformation.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/combine_result_layout_transformation.mlir
@@ -4,6 +4,12 @@
// RUN: -split-input-file %s | FileCheck %s --check-prefixes=WORKGROUP-SCOPE
// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-combine-result-layout-transformation{scope=dispatch-reshape},canonicalize,cse))" \
// RUN: -split-input-file %s | FileCheck %s --check-prefixes=RESHAPE-SCOPE
+// Scalarize the index transformations to show that the per-group fold keeps
+// the lowered body division-free, while a global `delinearize_index` over a
+// dynamic shape would scalarize to runtime-divisor `arith.divsi`/`remsi`.
+// Gated on the EXPAND prefix.
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-combine-result-layout-transformation{scope=dispatch},canonicalize,cse,affine-expand-index-ops,canonicalize,cse))" \
+// RUN: -split-input-file %s | FileCheck %s --check-prefixes=EXPAND
func.func @fold_collapse_shape_op(%source : tensor<2x4x16xf32>, %result : memref<8x16xf32>) {
%collapse = tensor.collapse_shape %source [[0, 1], [2]] : tensor<2x4x16xf32> into tensor<8x16xf32>
@@ -250,7 +256,6 @@
iree_codegen.store_to_buffer %unpack, %result : tensor<?x?xf32> into memref<?x?xf32>
return
}
-// DISPATCH-SCOPE: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 128)>
// DISPATCH-SCOPE-LABEL: @fold_unpack_op
// DISPATCH-SCOPE-SAME: %[[SOURCE:[a-zA-Z0-9_]+]]
// DISPATCH-SCOPE-SAME: %[[RESULT:[a-zA-Z0-9_]+]]
@@ -260,26 +265,65 @@
// DISPATCH-SCOPE-DAG: %[[RES_D1:.+]] = memref.dim %[[RESULT]], %[[C1]] : memref<?x?xf32>
// DISPATCH-SCOPE-DAG: %[[SRC_D0:.+]] = tensor.dim %[[SOURCE]], %[[C0]] : tensor<?x?x128x128xf32>
// DISPATCH-SCOPE-DAG: %[[SRC_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]] : tensor<?x?x128x128xf32>
-// DISPATCH-SCOPE-DAG: %[[COLLAPSE_SIZE0:.+]] = affine.apply #[[$MAP]]()[%[[SRC_D0]]]
-// DISPATCH-SCOPE-DAG: %[[COLLAPSE_SIZE1:.+]] = affine.apply #[[$MAP]]()[%[[SRC_D1]]]
// DISPATCH-SCOPE-DAG: %[[MAP_SCATTER_DEST:.+]] = tensor.empty(%[[RES_D0]], %[[RES_D1]]) : tensor<?x?xf32>
// DISPATCH-SCOPE: %[[MAP_SCATTER:.+]] = iree_linalg_ext.map_store
// DISPATCH-SCOPE-SAME: %[[SOURCE]] into %[[MAP_SCATTER_DEST]] {
// DISPATCH-SCOPE-NEXT: ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index, %[[IDX2:.+]]: index, %[[IDX3:.+]]: index):
-// DISPATCH-SCOPE: %[[LINEARIZE:.+]] = affine.linearize_index
-// DISPATCH-SCOPE-SAME: [%[[IDX0]], %[[IDX2]], %[[IDX1]], %[[IDX3]]]
-// DISPATCH-SCOPE-SAME: by (%[[SRC_D0]], 128, %[[SRC_D1]], 128)
-// DISPATCH-SCOPE: %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[LINEARIZE]]
-// DISPATCH-SCOPE-SAME: into (%[[COLLAPSE_SIZE0]], %[[COLLAPSE_SIZE1]])
-// DISPATCH-SCOPE: %[[BOUND0:.+]] = arith.cmpi ult, %[[DELINEARIZE]]#0, %[[RES_D0]]
-// DISPATCH-SCOPE: %[[BOUND1:.+]] = arith.cmpi ult, %[[DELINEARIZE]]#1, %[[RES_D1]]
+// Each collapsed dimension folds independently, so the merge factors stay
+// per-group rather than going through one global linearize + delinearize.
+// DISPATCH-SCOPE: %[[LIN0:.+]] = affine.linearize_index disjoint
+// DISPATCH-SCOPE-SAME: [%[[IDX0]], %[[IDX2]]] by (%[[SRC_D0]], 128)
+// DISPATCH-SCOPE: %[[LIN1:.+]] = affine.linearize_index disjoint
+// DISPATCH-SCOPE-SAME: [%[[IDX1]], %[[IDX3]]] by (%[[SRC_D1]], 128)
+// DISPATCH-SCOPE: %[[BOUND0:.+]] = arith.cmpi ult, %[[LIN0]], %[[RES_D0]]
+// DISPATCH-SCOPE: %[[BOUND1:.+]] = arith.cmpi ult, %[[LIN1]], %[[RES_D1]]
// DISPATCH-SCOPE: %[[MASK:.+]] = arith.andi %[[BOUND0]], %[[BOUND1]] : i1
-// DISPATCH-SCOPE: iree_linalg_ext.yield %[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[MASK]]
+// DISPATCH-SCOPE: iree_linalg_ext.yield %[[LIN0]], %[[LIN1]], %[[MASK]]
// DISPATCH-SCOPE: } : tensor<?x?x128x128xf32> into tensor<?x?xf32> -> tensor<?x?xf32>
// DISPATCH-SCOPE: iree_codegen.store_to_buffer %[[MAP_SCATTER]], %[[RESULT]] : tensor<?x?xf32> into memref<?x?xf32>
// -----
+// Dynamic-shape companion of @fold_pack_op. Mirrors @fold_unpack_op's
+// 2-reassociation-group shape (dim -> dynamic-outer + 128-tile), but going
+// the other way -- exercises the expand-case per-group `delinearize_index`
+// path on dimensions that don't canonicalize away (the case the pre-PR
+// global linearize + delinearize would have scalarized to a per-element
+// dynamic integer division).
+func.func @fold_pack_op_dynamic(%source : tensor<?x?xf32>, %result : memref<?x?x128x128xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.0 : f32
+ %d0 = memref.dim %result, %c0 : memref<?x?x128x128xf32>
+ %d1 = memref.dim %result, %c1 : memref<?x?x128x128xf32>
+ %dest = tensor.empty(%d0, %d1) : tensor<?x?x128x128xf32>
+ %pack = linalg.pack %source padding_value(%cst : f32)
+ outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [128, 128]
+ into %dest : tensor<?x?xf32> -> tensor<?x?x128x128xf32>
+ iree_codegen.store_to_buffer %pack, %result : tensor<?x?x128x128xf32> into memref<?x?x128x128xf32>
+ return
+}
+// DISPATCH-SCOPE-LABEL: @fold_pack_op_dynamic
+// DISPATCH-SCOPE-SAME: %[[SOURCE:[a-zA-Z0-9_]+]]
+// DISPATCH-SCOPE-SAME: %[[RESULT:[a-zA-Z0-9_]+]]
+// DISPATCH-SCOPE-DAG: %[[C0:.+]] = arith.constant 0 : index
+// DISPATCH-SCOPE-DAG: %[[C1:.+]] = arith.constant 1 : index
+// DISPATCH-SCOPE-DAG: %[[RES_D0:.+]] = memref.dim %[[RESULT]], %[[C0]] : memref<?x?x128x128xf32>
+// DISPATCH-SCOPE-DAG: %[[RES_D1:.+]] = memref.dim %[[RESULT]], %[[C1]] : memref<?x?x128x128xf32>
+// DISPATCH-SCOPE-DAG: %[[MAP_SCATTER_DEST:.+]] = tensor.empty(%[[RES_D0]], %[[RES_D1]]) : tensor<?x?x128x128xf32>
+// DISPATCH-SCOPE: %[[MAP_SCATTER:.+]] = iree_linalg_ext.map_store
+// DISPATCH-SCOPE-SAME: %[[SOURCE]] into %[[MAP_SCATTER_DEST]] {
+// DISPATCH-SCOPE-NEXT: ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index):
+// Each expanded dimension folds independently, so the split factors stay
+// per-group rather than going through one global linearize + delinearize.
+// DISPATCH-SCOPE: %[[DELIN0:.+]]:2 = affine.delinearize_index %[[IDX0]] into (%[[RES_D0]], 128)
+// DISPATCH-SCOPE: %[[DELIN1:.+]]:2 = affine.delinearize_index %[[IDX1]] into (%[[RES_D1]], 128)
+// DISPATCH-SCOPE: iree_linalg_ext.yield %[[DELIN0]]#0, %[[DELIN1]]#0, %[[DELIN0]]#1, %[[DELIN1]]#1
+// DISPATCH-SCOPE: } : tensor<?x?xf32> into tensor<?x?x128x128xf32> -> tensor<?x?x128x128xf32>
+// DISPATCH-SCOPE: iree_codegen.store_to_buffer %[[MAP_SCATTER]], %[[RESULT]] : tensor<?x?x128x128xf32> into memref<?x?x128x128xf32>
+
+// -----
+
func.func @fold_pack_op(%source : tensor<250x250xf32>, %result : memref<2x2x128x128xf32>) {
%cst = arith.constant 0.0 : f32
%dest = tensor.empty() : tensor<2x2x128x128xf32>
@@ -600,3 +644,55 @@
// RESHAPE-SCOPE-LABEL: @pure_transpose_skipped_under_dispatch_reshape
// RESHAPE-SCOPE: linalg.transpose
// RESHAPE-SCOPE-NOT: iree_linalg_ext.map_store
+
+// -----
+
+// Lowering contrast: scalarizing the per-group `(de)linearize_index` the
+// combine pass emits leaves the body division-free -- multiplies and adds by
+// the static tile size -- while a global `delinearize_index` over the full
+// dynamic result shape scalarizes to runtime-divisor `arith.divsi`/`remsi`.
+
+func.func @lowered_per_group_no_division(%source : tensor<?x?x128x128xf32>,
+ %result : memref<?x?xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %d0 = memref.dim %result, %c0 : memref<?x?xf32>
+ %d1 = memref.dim %result, %c1 : memref<?x?xf32>
+ %dest = tensor.empty(%d0, %d1) : tensor<?x?xf32>
+ %unpack = linalg.unpack %source
+ outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [128, 128]
+ into %dest : tensor<?x?x128x128xf32> -> tensor<?x?xf32>
+ iree_codegen.store_to_buffer %unpack, %result : tensor<?x?xf32> into memref<?x?xf32>
+ return
+}
+// The per-group `linearize_index` lowers to `muli`/`addi` by the static tile
+// size 128 -- no division of any kind.
+// EXPAND-LABEL: @lowered_per_group_no_division
+// EXPAND-DAG: %[[C128:.+]] = arith.constant 128 : index
+// EXPAND: iree_linalg_ext.map_store
+// EXPAND: arith.muli %{{.+}}, %[[C128]]
+// EXPAND: arith.addi
+// EXPAND: arith.muli %{{.+}}, %[[C128]]
+// EXPAND: arith.addi
+// EXPAND-NOT: arith.divsi
+// EXPAND-NOT: arith.floordivsi
+// EXPAND-NOT: arith.remsi
+
+// -----
+
+// Negative case for contrast: a global `affine.delinearize_index` over a
+// fully dynamic result shape -- the per-group fold avoids producing this IR.
+func.func @lowered_global_dynamic_divisor(
+ %idx: index, %d0: index, %d1: index, %d2: index, %d3: index
+) -> (index, index, index, index) {
+ %r:4 = affine.delinearize_index %idx into (%d0, %d1, %d2, %d3)
+ : index, index, index, index
+ return %r#0, %r#1, %r#2, %r#3 : index, index, index, index
+}
+// Scalarizes to runtime-divisor `arith.divsi`/`remsi` (divisors are
+// `arith.muli` products of the dynamic factors) -- the cost the per-group
+// fold avoids.
+// EXPAND-LABEL: @lowered_global_dynamic_divisor
+// EXPAND: arith.muli
+// EXPAND: arith.divsi
+// EXPAND: arith.remsi
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_map_store.mlir b/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_map_store.mlir
index afe51c6..8df64a3 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_map_store.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/generic_vectorization_map_store.mlir
@@ -139,3 +139,78 @@
}
// CHECK-LABEL: @map_store_f4_mask_depends_on_inner_index
// CHECK-NOT: vector
+
+// -----
+
+// A `map_store` whose body uses the per-group `affine.linearize_index disjoint`
+// form -- the IR shape `foldReshapeIntoMapStore` emits when folding a
+// `collapse_shape` per reassociation group rather than via one global linearize
+// + delinearize. Vectorization accepts it when the input is statically shaped
+// (the relevant criterion is `inputType.hasStaticShape()`, the index
+// transformation in the body is opaque to that check).
+func.func @vectorize_map_store_per_group_linearize(
+ %input: tensor<2x4x16xf32>, %output: tensor<8x16xf32>
+) -> tensor<8x16xf32> {
+ %0 = iree_linalg_ext.map_store %input into %output {
+ ^bb0(%idx0: index, %idx1: index, %idx2: index):
+ %lin = affine.linearize_index disjoint [%idx0, %idx1] by (2, 4) : index
+ %mask = arith.constant true
+ iree_linalg_ext.yield %lin, %idx2, %mask : index, index, i1
+ } : tensor<2x4x16xf32> into tensor<8x16xf32> -> tensor<8x16xf32>
+ return %0 : tensor<8x16xf32>
+}
+// CHECK-LABEL: @vectorize_map_store_per_group_linearize
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[OUTPUT:[a-zA-Z0-9_]+]]
+// CHECK: %[[READ:.+]] = vector.transfer_read %[[INPUT]]
+// CHECK: %[[MAP_SCATTER:.+]] = iree_linalg_ext.map_store
+// CHECK-SAME: %[[READ]] into %[[OUTPUT]]
+// CHECK: : vector<2x4x16xf32> into tensor<8x16xf32> -> tensor<8x16xf32>
+
+// -----
+
+// Expand-direction companion of the above: per-group `affine.delinearize_index`
+// in the body (the shape emitted when folding an `expand_shape` per
+// reassociation group). Same story -- vectorization is gated on the input's
+// static shape, not on what the index transformation looks like.
+func.func @vectorize_map_store_per_group_delinearize(
+ %input: tensor<8x16xf32>, %output: tensor<2x4x16xf32>
+) -> tensor<2x4x16xf32> {
+ %0 = iree_linalg_ext.map_store %input into %output {
+ ^bb0(%idx0: index, %idx1: index):
+ %delin:2 = affine.delinearize_index %idx0 into (2, 4) : index, index
+ %mask = arith.constant true
+ iree_linalg_ext.yield %delin#0, %delin#1, %idx1, %mask : index, index, index, i1
+ } : tensor<8x16xf32> into tensor<2x4x16xf32> -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
+// CHECK-LABEL: @vectorize_map_store_per_group_delinearize
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]
+// CHECK-SAME: %[[OUTPUT:[a-zA-Z0-9_]+]]
+// CHECK: %[[READ:.+]] = vector.transfer_read %[[INPUT]]
+// CHECK: %[[MAP_SCATTER:.+]] = iree_linalg_ext.map_store
+// CHECK-SAME: %[[READ]] into %[[OUTPUT]]
+// CHECK: : vector<8x16xf32> into tensor<2x4x16xf32> -> tensor<2x4x16xf32>
+
+// -----
+
+// Dynamic-input counterpart of @vectorize_map_store_per_group_delinearize:
+// the per-group form is what reaches this point in practice for CPU-tiled
+// encoding relayouts, but the vectorizer still skips it because the input is
+// dynamically shaped -- consistent with @no_vectorize_map_store_dynamic above.
+// The per-group fold's payoff is at the later scalarization stage (static
+// shifts/masks instead of dynamic integer division per element), not at
+// vectorization.
+func.func @no_vectorize_map_store_per_group_delinearize_dynamic(
+ %input: tensor<?xf32>, %output: tensor<?x4xf32>, %d0 : index
+) -> tensor<?x4xf32> {
+ %0 = iree_linalg_ext.map_store %input into %output {
+ ^bb0(%idx0: index):
+ %delin:2 = affine.delinearize_index %idx0 into (%d0, 4) : index, index
+ %mask = arith.constant true
+ iree_linalg_ext.yield %delin#0, %delin#1, %mask : index, index, i1
+ } : tensor<?xf32> into tensor<?x4xf32> -> tensor<?x4xf32>
+ return %0 : tensor<?x4xf32>
+}
+// CHECK-LABEL: @no_vectorize_map_store_per_group_delinearize_dynamic
+// CHECK-NOT: vector
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp b/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
index 9ae3075..ac79fbb 100644
--- a/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
+++ b/compiler/src/iree/compiler/Codegen/Interfaces/VectorizableOpInterface.cpp
@@ -462,6 +462,18 @@
return false;
}
ShapedType inputType = mapStoreOp.getInputType();
+ // `vectorize()` below emits an unmasked, full-tile `vector.transfer_read`,
+ // so the input must be statically shaped. A `map_store` iterates its
+ // *input*, so it only vectorizes when that input is already statically
+ // tiled. For a data-tiling encoding relayout the input is the un-padded
+ // source tensor, which has no static tile structure -- so a CPU-tiled
+ // `map_store` reaches here dynamically shaped and is left to scalarize.
+ // The padded, statically-tiled side of an encoding relayout is its
+ // *output*; a relayout rooted there -- a `map_load`, which iterates its
+ // result -- tiles into complete static tiles by construction. So a CPU
+ // encoding relayout should be combined into a `map_load`, not a
+ // `map_store`. (GPU also forms `map_store` for these but tiles them
+ // statically via its thread-grid distribution, so it does not hit this.)
if (!inputType.hasStaticShape()) {
return false;
}
