[all-commits] [llvm/llvm-project] 33ef53: [mlir][linalg] Fix mask rank for masked contiguous...

[pstarkcdpr via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 33ef532c79b321cf22232ba0054b36d259c62079
      https://github.com/llvm/llvm-project/commit/33ef532c79b321cf22232ba0054b36d259c62079
  Author: pstarkcdpr <paul.stark at cdprojektred.com>
  Date:   2026-07-08 (Wed, 08 Jul 2026)

  Changed paths:
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    M mlir/test/Dialect/Linalg/vectorization/extract.mlir

  Log Message:
  -----------
  [mlir][linalg] Fix mask rank for masked contiguous `tensor.extract` (#206207)

### Summary

Note: This fix was made mostly by Claude based on a failure case in
IREE. It addresses issue
https://github.com/llvm/llvm-project/issues/206209

When `vectorizeTensorExtract` lowers a `tensor.extract` recognized as a
*contiguous load*, it 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. Until now this read was
returned unmasked and masked later by the generic path, which applies a
**full iteration-space identity mask**. When the source rank is smaller
than the loop nest, that mask is over-ranked relative to the
(rank-reduced) read — e.g. a `vector<1x4xi1>` mask on a read whose
inferred mask type is `vector<4xi1>` — and the op fails verification:

```
'vector.mask' op expects a 'vector<4xi1>' mask for the maskable operation
```

(or, after the masked read is rank-reduced by canonicalization in a full
pipeline, `'vector.transfer_read' op inferred mask type ... don't
match`).

### Root cause

The three `tensor.extract` lowering paths in `vectorizeTensorExtract`
handle masking inconsistently:

- **gather** masks itself — correct, its result is full-rank;
- **scalar broadcast** masks itself with a rank-1 mask, explicitly
noting that the generic identity-map masking "wouldn't be valid here";
- **contiguous load** does neither and relies on the generic path, which
masks over all loop dims. That is only correct when the read is
full-rank (`srcRank >= numLoops`); when `srcRank < numLoops` the read
broadcasts its leading dims and needs a rank-reduced mask.

Existing tests only covered the full-rank case (e.g. a
`tensor<80x16xf32>` source in a 2-D nest), so the rank-reducing case was
never exercised.

### 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.

### Testing

- New regression test `@masked_contiguous_extract_rank_reducing_mask` in
`mlir/test/Dialect/Linalg/vectorization/extract.mlir` (1-D source inside
a 2-D loop nest). It fails verification before this change and checks
the rank-reduced mask after it.
- `mlir/test/Dialect/Linalg` (164 tests) and `mlir/test/Dialect/Vector`
+ `mlir/test/Dialect/Linalg/vectorization` (113 tests) all pass.

### Notes for reviewers

- The change is localized to the contiguous-load branch of
`vectorizeTensorExtract`; the gather and scalar-broadcast branches are
untouched.
- No new mask is created in the full-rank case — the existing cached
iteration-space mask is reused, so there is no codegen/IR churn for
current callers.

---------

Signed-off-by: Paul Stark <paul.stark at cdprojektred.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications