ITADN

Masked vectorization of a contiguous `tensor.extract` produces an invalid `vector.transfer_read` mask when the source rank < loop rank

#24637Closedpstarkcdpr 创建于 2026-06-26
bug 🐞
P
pstarkcdprcommented
### 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 条评论