[Mlir-commits] [mlir] 33ef532 - [mlir][linalg] Fix mask rank for masked contiguous `tensor.extract` (#206207)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 8 09:13:29 PDT 2026
Author: pstarkcdpr
Date: 2026-07-08T17:13:24+01:00
New Revision: 33ef532c79b321cf22232ba0054b36d259c62079
URL: https://github.com/llvm/llvm-project/commit/33ef532c79b321cf22232ba0054b36d259c62079
DIFF: https://github.com/llvm/llvm-project/commit/33ef532c79b321cf22232ba0054b36d259c62079.diff
LOG: [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>
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorization/extract.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index b57e66a1c3580..21ca3108efcd6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1299,9 +1299,21 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
rewriter, loc, resultType, extractOp.getTensor(), transferReadIdxs,
/*padding=*/std::nullopt, permutationMap, inBounds);
+ // Mask this contiguous xfer_read here rather than relying on the generic
+ // path (the generic path assumes an identity masking map over all the loop
+ // dims, which wouldn't be valid here). A contiguous load only reads the
+ // trailing `min(dstRank, srcRank)` dims of the iteration space - the leading
+ // dims are broadcast via `permutationMap` above - so its inferred mask is
+ // rank-reduced. Build a masking map that projects the iteration space onto
+ // exactly those trailing dims so the created mask matches the xfer_read.
+ int64_t numReadDims = std::min(dstRank, srcRank);
+ auto maskingMap = AffineMap::getMinorIdentityMap(
+ linalgOp.getNumLoops(), numReadDims, rewriter.getContext());
+ Operation *maskedReadOp =
+ state.maskOperation(rewriter, transferReadOp, linalgOp, maskingMap);
+
LDBG() << "Vectorised as contiguous load: " << extractOp;
- return VectorizationHookResult{VectorizationHookStatus::NewOp,
- transferReadOp};
+ return VectorizationHookResult{VectorizationHookStatus::NewOp, maskedReadOp};
}
/// Emit reduction operations if the shapes of the value to reduce is
diff erent
diff --git a/mlir/test/Dialect/Linalg/vectorization/extract.mlir b/mlir/test/Dialect/Linalg/vectorization/extract.mlir
index 76ac4b8398069..fc0a11ef9810e 100644
--- a/mlir/test/Dialect/Linalg/vectorization/extract.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/extract.mlir
@@ -477,3 +477,52 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// A contiguous `tensor.extract` from a source whose rank is *smaller* than the
+// iteration space (here: a 1-D source inside a 2-D loop nest). The contiguous
+// `vector.transfer_read` only reads the trailing (contiguous) dim - the leading
+// dim is broadcast via the permutation map - so its mask must be rank-reduced
+// (`vector<4xi1>`) rather than the full iteration-space mask (`vector<1x4xi1>`).
+// Regression test: previously the full iteration mask was attached, producing
+// an invalid `vector.transfer_read` ("inferred mask type ... don't match").
+
+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>
+}
+
+// CHECK-LABEL: func.func @masked_contiguous_extract_rank_reducing_mask
+// CHECK-SAME: %[[SRC:.*]]: tensor<16xf32>,
+// CHECK-SAME: %[[OUTPUT:.*]]: tensor<1x3xf32>,
+// CHECK-SAME: %[[IDX_IN:.*]]: index) -> tensor<1x3xf32> {
+
+/// Full iteration-space mask (used for the output read/write).
+// CHECK: %[[MASK_2D:.*]] = vector.create_mask {{.*}} : vector<1x4xi1>
+
+/// The contiguous read of the 1-D source uses a rank-reduced mask.
+// CHECK: %[[MASK_1D:.*]] = vector.create_mask {{.*}} : vector<4xi1>
+// CHECK: vector.mask %[[MASK_1D]] {
+// CHECK-SAME: vector.transfer_read %[[SRC]]{{.*}} {in_bounds = [true, true], permutation_map = {{.*}}} : tensor<16xf32>, vector<1x4xf32>
+// CHECK-SAME: } : vector<4xi1> -> vector<1x4xf32>
+
+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
+ }
+}
More information about the Mlir-commits
mailing list