Masked vectorization of a contiguous `tensor.extract` produces an invalid `vector.transfer_read` mask when the source rank < loop rank
bug 🐞
### What happened?
This issue was mostly created by Claude while debugging compile errors in IREE.
`linalg::vectorize` (the `vectorizeTensorExtract` hook in `mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp`) mis-masks a `tensor.extract` that is classified as a *contiguous load* when the extracted tensor has fewer dimensions than the surrounding loop nest.
The contiguous-load path builds a `vector.transfer_read` whose permutation map **broadcasts the leading iteration dims** and only reads the trailing `min(dstRank, srcRank)` dims of the source. It then returns the read *unmasked* and lets the generic masking path wrap it. The generic path applies a mask over the **full iteration space** (rank = number of loops). When the source rank is smaller than the loop nest, that mask is over-ranked relative to the rank-reduced read, so the resulting op fails verification.
## Reproducer
`repro.mlir`:
```mlir
// RUN: mlir-opt %s -transform-interpreter -split-input-file
//
// A contiguous `tensor.extract` from a 1-D source inside a 2-D loop nest.
// Vectorizing with masking (vector_sizes [1, 4], with `vectorize_nd_extract`)
// produces an invalid masked `vector.transfer_read`: the mask has the full
// iteration-space rank (vector<1x4xi1>) while the rank-reduced contiguous read
// only needs vector<4xi1>.
func.func @masked_contiguous_extract_rank_reducing_mask(
%src: tensor<16xf32>,
%output : tensor<1x3xf32>,
%idx: index) -> tensor<1x3xf32> {
%1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} outs(%output : tensor<1x3xf32>) {
^bb0(%out: f32):
%2 = linalg.index 1 : index
%3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %idx)
%extracted = tensor.extract %src[%3] : tensor<16xf32>
linalg.yield %extracted : f32
} -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
transform.structured.vectorize %0 vector_sizes [1, 4] {vectorize_nd_extract} : !transform.any_op
transform.yield
}
}
```
```
mlir-opt repro.mlir -transform-interpreter -split-input-file
```
## Actual output
Vectorization fails verification:
```
error: 'vector.mask' op expects a 'vector<4xi1>' mask for the maskable operation
```
The offending op is the contiguous read of the 1-D source: a `vector.transfer_read` with permutation map `(d0) -> (0, d0)` producing `vector<1x4xf32>` (whose inferred mask type is `vector<4xi1>`) wrapped in a `vector.mask` carrying the full iteration-space mask `vector<1x4xi1>`.
(The same defect surfaces as the `vector.transfer_read` verifier error `inferred mask type ('vector<NxN1>') and mask operand type (...) don't match` when the masked read is later rank-reduced by canonicalization, e.g. inside a
full backend codegen pipeline.)
## Expected
Vectorization should succeed and emit a rank-reduced mask that matches the contiguous read, e.g.:
```mlir
%mask1d = vector.create_mask %c3 : vector<4xi1>
%read = vector.mask %mask1d {
vector.transfer_read %src[%i], %pad {in_bounds = [true, true],
permutation_map = affine_map<(d0) -> (0, d0)>}
: tensor<16xf32>, vector<1x4xf32>
} : vector<4xi1> -> vector<1x4xf32>
```
## Root cause
In `vectorizeTensorExtract`, the three `tensor.extract` lowering paths handle masking inconsistently:
- **Gather** masks itself (`state.maskOperation(..., gatherOp, linalgOp)`); the gather result is full-rank so the full iteration-space mask is correct.
- **Scalar broadcast** masks itself with an explicitly constructed rank-1 mask and carries a comment noting that *"the generic path assumes identity masking map, which wouldn't be valid here."*
- **Contiguous load** does **neither** — it returns the bare `transfer_read` and relies on the generic path. The generic path (`maskOperation` -> `getOrCreateMaskFor`, called with no indexing map) falls back to `AffineMap::getMultiDimIdentityMap(linalgOp.getNumLoops(), ...)`, i.e. a mask over all loop dims. But the contiguous read broadcasts its leading dims (`AffineMap::getMinorIdentityMap(srcRank, min(dstRank, srcRank))` extended with leading zeros), so its inferred mask is rank-reduced. The two disagree whenever `srcRank < numLoops`.
Existing tests only cover the case where the source is at least as high-rank as the loop nest (e.g. `tensor<80x16xf32>` in a 2-D nest), where the read is full-rank and the generic mask happens to be correct — so the rank-reducing case was never exercised.
## Proposed fix
Mask the contiguous read in place (mirroring the scalar-broadcast path), using a masking map that projects the iteration space onto exactly the trailing `min(dstRank, srcRank)` dims that are read:
```cpp
int64_t numReadDims = std::min(dstRank, srcRank);
auto maskingMap = AffineMap::getMinorIdentityMap(
linalgOp.getNumLoops(), numReadDims, rewriter.getContext());
Operation *maskedReadOp =
state.maskOperation(rewriter, transferReadOp, linalgOp, maskingMap);
```
This is **behavior-preserving for the existing full-rank case**: when `min(dstRank, srcRank) == numLoops`, `getMinorIdentityMap` collapses to the full identity map, producing the same mask (same `activeMaskCache` key) and identical IR. Only the previously-broken rank-reducing case changes.
I have a patch + regression test ready and will open a PR.
## Environment
- Reproduced on llvm-project at commit `22da7f92913937b0318e3d504b2d8e90faaf2597`; the `vectorizeTensorExtract` code path is unchanged on recent `main`.
- Surfaced originally via IREE's `llvm-cpu` backend (`GenericVectorizationPass`) when masking is triggered by a non-vector-width-multiple inner dimension (e.g. an AVX-512 target with a size-24 inner dim) reading from a 1-D source.
### Steps to reproduce your issue
`mlir-opt repro.mlir -transform-interpreter -split-input-file`
### What component(s) does this issue relate to?
MLIR
### Version information
Tried on `22da7f92913937b0318e3d504b2d8e90faaf2597` since this is currently used by IREE.
### Additional context
_No response_
关闭于 2026-06-26 2 条评论