[Mlir-commits] [mlir] "[mlir][linalg] Fix masking for rank-reducing contiguous extracts" (PR #206637)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: LouisLu060211

<details>
<summary>Changes</summary>

Fix masked vectorization of contiguous tensor.extract when the source tensor rank is smaller than the surrounding loop rank.

The contiguous-load tensor.extract path builds a rank-reducing vector.transfer_read by reading only the trailing source dimensions and broadcasting leading vector dimensions. Previously this path returned the transfer_read unmasked and relied on the generic masking path, which builds a mask over the full linalg iteration space.

That full-rank mask is invalid when the source rank is smaller than the loop rank. For example, a 1-D source read inside a 2-D loop nest can produce a vector.transfer_read with inferred mask type vector<4xi1>, while the generic path wraps it with vector<1x4xi1>.

Mask the contiguous transfer_read in place using a minor-identity masking map over the dimensions actually read. This matches the transfer_read permutation map and preserves the existing full-rank behavior.

Add a regression test covering a masked contiguous tensor.extract from a 1-D source inside a 2-D linalg.generic vectorized with vectorize_nd_extract.

Fixes #<!-- -->206209.
Replaces #<!-- -->206634.

Test:
- llvm-lit -a -vv mlir/test/Dialect/Linalg/transform-op-vectorize.mlir

Assisted-by: Codex

---
Full diff: https://github.com/llvm/llvm-project/pull/206637.diff


1 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+2) 


``````````diff
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index f0a210a2ededb..703d564a07648 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -820,3 +820,5 @@ void RemoveDeadValues::runOnOperation() {
     signalPassFailure();
   }
 }
+
+

``````````

</details>


https://github.com/llvm/llvm-project/pull/206637