[Mlir-commits] [mlir] [mlir][vector] Verify non-unit strides on `masked/expand/compress` ops and fix `SparseVectorization` miscompilation (PR #210952)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 21 04:27:17 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Federico Bruzzone (FedericoBruzzone)
<details>
<summary>Changes</summary>
This closes the stride-verification gap left open by #<!-- -->204611 and #<!-- -->205869 for `vector.maskedload`, `vector.maskedstore`, `vector.expandload`, and `vector.compressstore`, and uncovers and fixes a real, pre-existing silent miscompile in the SparseVectorizer.
### First Fix
`vector.maskedload`/`maskedstore` lower to [`llvm.masked.load`](https://llvm.org/docs/LangRef.html#llvm-masked-load-intrinsics) / [`llvm.masked.store`](https://llvm.org/docs/LangRef.html#llvm-masked-store-intrinsics), and `vector.expandload`/`compressstore` lower to [`llvm.masked.expandload`](https://llvm.org/docs/LangRef.html#llvm-masked-expandload-intrinsics) / [`llvm.masked.compressstore`](https://llvm.org/docs/LangRef.html#llvm-masked-compressstore-intrinsics).
All four LLVM intrinsics read or write N *consecutive* memory locations starting from a single computed pointer. None of these four ops verified that the memref's most minor dimension actually has unit stride, so a memref like `memref<4xi32, strided<[2], offset: ?>>` passed verification and produced a masked load/store that silently read/wrote the wrong memory.
Fixed by adding a new verifier helper distinct.. It only rejects a most-minor stride that is *statically known* and `!= 1`, while a *dynamic* stride (`?`) is accepted.
**Note**: This is intentionally more permissive than `vector.load`/`vector.store`, the sparsifier legitimately produces masked loads on memrefs whose last-dim stride is contiguous by construction but not provable at the type level (see below).
### Related fix (SparseVectorization)
When SparseVectorization encounters a _direct_ access like `a[lo:hi]`, it emits a `vector.maskedload` assuming that reading N consecutive indices (`lo, lo+1, ..., lo+vl-1`) means reading N consecutive *bytes* in memory. That's only true if the underlying buffer has stride 1.
At vectorization time the memref type is still `memref<?xindex, strided<[?], offset: ?>>`, fully dynamic, even when the real stride will turn out to be 2 or 3.
The fix adds `hasKnownNonUnitStride()`, which special-cases `sparse_tensor.coordinates` results: instead of trusting the memref type, it reads the stride directly off the source tensor's encoding, which is fully known at compile time this pipeline stage.
**While all tests pass, this section needs to be carefully reviewed, as I am not very familiar with it (I studied it during this development). I am not sure if it is correct to handle it at this level. 🙏**
AI was used to investigate in this PR.
---
Full diff: https://github.com/llvm/llvm-project/pull/210952.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp (+50)
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+50)
- (modified) mlir/test/Dialect/Vector/invalid.mlir (+54)
``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index 23436a68535fc..c60ca523d81f3 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -54,6 +54,50 @@ static bool isInvariantArg(BlockArgument arg, Block *block) {
return arg.getOwner() != block;
}
+/// Returns true when `mem`'s most minor dimension has a statically known
+/// non-unit stride.
+///
+/// `genVectorLoad/genVectorStore` assume a contiguous
+/// `vector.maskedload/vector.maskedstore` is safe for consecutive loop
+/// indices, which breaks for a strided view extracting one component out
+/// of an interleaved (AoS) COO coordinate buffer.
+///
+/// Example:
+/// A `compressed(nonunique) + singleton` region stores coordinates as
+/// `[row0, col0, row1, col1, ...]`, so a 2-lane masked load of `col[0:2]`
+/// (offset=1) would read physical offsets {1, 2} = `[col0, row1]` instead
+/// of the intended {1, 3} = `[col0, col1]`: a silent miscompile.
+///
+/// NOTE: A stride that can't be proven non-unit by either means is assumed
+/// safe.
+static bool hasKnownNonUnitStride(Value mem) {
+ // sparse_tensor.coordinates isn't lowered to a concrete strided memref
+ // until sparse-tensor-codegen runs, so at this point its type
+ // still has a dynamic stride even when the true stride is already known
+ // from the tensor's encoding -- hence the special case below instead of
+ // trusting the memref type.
+ if (auto toCoords = mem.getDefiningOp<ToCoordinatesOp>()) {
+ SparseTensorType stt = getSparseTensorType(toCoords.getTensor());
+ Level cooStart = stt.getAoSCOOStart();
+ // A single trailing level (lvlRank - cooStart == 1) is not actually
+ // interleaved with anything else, so it degenerates to a contiguous
+ // buffer.
+ if (toCoords.getLevel() >= cooStart)
+ return stt.getLvlRank() - cooStart != 1;
+ return false;
+ }
+
+ auto memTp = dyn_cast<MemRefType>(mem.getType());
+ if (!memTp)
+ return false;
+ SmallVector<int64_t> strides;
+ int64_t offset;
+ if (failed(memTp.getStridesAndOffset(strides, offset)))
+ return false;
+ return !strides.empty() && !ShapedType::isDynamic(strides.back()) &&
+ strides.back() != 1;
+}
+
/// Constructs vector type for element type.
static VectorType vectorType(VL vl, Type etp) {
return VectorType::get(vl.vectorLength, etp, vl.enableVLAVectorization);
@@ -292,6 +336,8 @@ static bool vectorizeSubscripts(PatternRewriter &rewriter, scf::ForOp forOp,
if (auto load = cast.getDefiningOp<memref::LoadOp>()) {
if (!innermost)
return false;
+ if (hasKnownNonUnitStride(load.getMemRef()))
+ return false;
if (codegen) {
SmallVector<Value> idxs2(load.getIndices()); // no need to analyze
Location loc = forOp.getLoc();
@@ -408,6 +454,8 @@ static bool vectorizeExpr(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
// a[lo:hi] = ind[lo:hi], where 'lo' denotes the current index
// and 'hi = lo + vl - 1'.
if (auto load = dyn_cast<memref::LoadOp>(def)) {
+ if (hasKnownNonUnitStride(load.getMemRef()))
+ return false;
auto subs = load.getIndices();
SmallVector<Value> idxs;
if (vectorizeSubscripts(rewriter, forOp, vl, subs, codegen, vmask, idxs)) {
@@ -582,6 +630,8 @@ static bool vectorizeStmt(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
}
} else if (auto store = dyn_cast<memref::StoreOp>(last)) {
// Analyze/vectorize store operation.
+ if (hasKnownNonUnitStride(store.getMemRef()))
+ return false;
auto subs = store.getIndices();
SmallVector<Value> idxs;
Value rhs = store.getValue();
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f37083803a2a1..155ae020bf368 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6190,6 +6190,32 @@ static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,
return success();
}
+/// Unlike vector.load/store, ops such as vector.maskedload/maskedstore and
+/// vector.expandload/compressstore are also used on memrefs whose most minor
+/// dimension has a dynamic stride that is contiguous at runtime but not
+/// provable as such at the type level (e.g., buffers produced by the
+/// sparsifier).
+///
+/// Only reject strides that are statically known to be non-unit: a dynamic
+/// stride is assumed to be safe.
+static LogicalResult verifyNonStaticNonUnitStrideRejected(Operation *op,
+ VectorType vecTy,
+ MemRefType memRefTy) {
+ if (!vecTy.isScalable() &&
+ (vecTy.getRank() == 0 || vecTy.getNumElements() == 1))
+ return success();
+
+ SmallVector<int64_t> strides;
+ int64_t offset;
+ if (failed(memRefTy.getStridesAndOffset(strides, offset)))
+ return success();
+
+ if (!strides.empty() && !ShapedType::isDynamic(strides.back()) &&
+ strides.back() != 1)
+ return op->emitOpError("most minor memref dim must have unit stride");
+ return success();
+}
+
LogicalResult vector::LoadOp::verify() {
VectorType resVecTy = getVectorType();
MemRefType memRefTy = getMemRefType();
@@ -6299,6 +6325,9 @@ LogicalResult MaskedLoadOp::verify() {
VectorType resVType = getVectorType();
MemRefType memType = getMemRefType();
+ if (failed(verifyNonStaticNonUnitStrideRejected(*this, resVType, memType)))
+ return failure();
+
// Negative strides are not supported on vector.maskedload. The lowering to
// LLVM emits arithmetic operations (e.g., GEP, mul) with nuw flags that
// assume non-negative strides to avoid undefined behavior.
@@ -6365,6 +6394,9 @@ LogicalResult MaskedStoreOp::verify() {
VectorType valueVType = getVectorType();
MemRefType memType = getMemRefType();
+ if (failed(verifyNonStaticNonUnitStrideRejected(*this, valueVType, memType)))
+ return failure();
+
// Negative strides are not supported on vector.maskedstore. The lowering to
// LLVM emits arithmetic operations (e.g., GEP, mul) with nuw flags that
// assume non-negative strides to avoid undefined behavior.
@@ -6648,6 +6680,15 @@ LogicalResult ExpandLoadOp::verify() {
VectorType resVType = getVectorType();
MemRefType memType = getMemRefType();
+ if (failed(verifyNonStaticNonUnitStrideRejected(*this, resVType, memType)))
+ return failure();
+
+ // Negative strides are not supported on vector.expandload. The lowering to
+ // LLVM emits arithmetic operations (e.g., GEP, mul) with nuw flags that
+ // assume non-negative strides to avoid undefined behavior.
+ if (memref::hasNegativeStaticStride(memType))
+ return emitOpError("memref strides must be non-negative");
+
if (failed(
verifyElementTypesMatch(*this, memType, resVType, "base", "result")))
return failure();
@@ -6702,6 +6743,15 @@ LogicalResult CompressStoreOp::verify() {
VectorType valueVType = getVectorType();
MemRefType memType = getMemRefType();
+ if (failed(verifyNonStaticNonUnitStrideRejected(*this, valueVType, memType)))
+ return failure();
+
+ // Negative strides are not supported on vector.compressstore. The lowering
+ // to LLVM emits arithmetic operations (e.g., GEP, mul) with nuw flags that
+ // assume non-negative strides to avoid undefined behavior.
+ if (memref::hasNegativeStaticStride(memType))
+ return emitOpError("memref strides must be non-negative");
+
if (failed(verifyElementTypesMatch(*this, memType, valueVType, "base",
"valueToStore")))
return failure();
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index aaa55cface958..03b8f12d29cad 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1422,6 +1422,15 @@ func.func @maskedload_negative_stride(%src: memref<100x100xf32, strided<[-100, 1
// -----
+func.func @maskedload_non_unit_stride(%src: memref<?xi8, strided<[2], offset: ?>>, %mask: vector<8xi1>, %pass: vector<8xi8>) -> vector<8xi8> {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.maskedload' op most minor memref dim must have unit stride}}
+ %0 = vector.maskedload %src[%c0], %mask, %pass : memref<?xi8, strided<[2], offset: ?>>, vector<8xi1>, vector<8xi8> into vector<8xi8>
+ return %0 : vector<8xi8>
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// vector.maskedstore
//===----------------------------------------------------------------------===//
@@ -1475,6 +1484,15 @@ func.func @maskedstore_negative_stride(%src: memref<100x100xf32, strided<[-100,
// -----
+func.func @maskedstore_non_unit_stride(%src: memref<?xi8, strided<[2], offset: ?>>, %mask: vector<8xi1>, %value: vector<8xi8>) {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.maskedstore' op most minor memref dim must have unit stride}}
+ vector.maskedstore %src[%c0], %mask, %value : memref<?xi8, strided<[2], offset: ?>>, vector<8xi1>, vector<8xi8>
+ return
+}
+
+// -----
+
func.func @gather_from_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
%mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
%c0 = arith.constant 0 : index
@@ -1742,6 +1760,24 @@ func.func @expand_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector<1
// -----
+func.func @expandload_non_unit_stride(%src: memref<?xi8, strided<[2], offset: ?>>, %mask: vector<8xi1>, %pass_thru: vector<8xi8>) -> vector<8xi8> {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.expandload' op most minor memref dim must have unit stride}}
+ %0 = vector.expandload %src[%c0], %mask, %pass_thru : memref<?xi8, strided<[2], offset: ?>>, vector<8xi1>, vector<8xi8> into vector<8xi8>
+ return %0 : vector<8xi8>
+}
+
+// -----
+
+func.func @expandload_negative_stride(%src: memref<100x100xf32, strided<[-100, 1]>>, %mask: vector<8xi1>, %pass_thru: vector<8xf32>) -> vector<8xf32> {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.expandload' op memref strides must be non-negative}}
+ %0 = vector.expandload %src[%c0, %c0], %mask, %pass_thru : memref<100x100xf32, strided<[-100, 1]>>, vector<8xi1>, vector<8xf32> into vector<8xf32>
+ return %0 : vector<8xf32>
+}
+
+// -----
+
func.func @compress_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>, %value: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error at +1 {{'vector.compressstore' op base element type ('f64') does not match valueToStore element type ('f32')}}
@@ -1796,6 +1832,24 @@ func.func @compress_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector
// -----
+func.func @compressstore_non_unit_stride(%src: memref<?xi8, strided<[2], offset: ?>>, %mask: vector<8xi1>, %value: vector<8xi8>) {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.compressstore' op most minor memref dim must have unit stride}}
+ vector.compressstore %src[%c0], %mask, %value : memref<?xi8, strided<[2], offset: ?>>, vector<8xi1>, vector<8xi8>
+ return
+}
+
+// -----
+
+func.func @compressstore_negative_stride(%src: memref<100x100xf32, strided<[-100, 1]>>, %mask: vector<8xi1>, %value: vector<8xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.compressstore' op memref strides must be non-negative}}
+ vector.compressstore %src[%c0, %c0], %mask, %value : memref<100x100xf32, strided<[-100, 1]>>, vector<8xi1>, vector<8xf32>
+ return
+}
+
+// -----
+
func.func @scan_reduction_dim_constraint(%arg0: vector<2x3xi32>, %arg1: vector<3xi32>) -> vector<3xi32> {
// expected-error at +1 {{'vector.scan' op reduction dimension 5 has to be less than 2}}
%0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = true, reduction_dim = 5} :
``````````
</details>
https://github.com/llvm/llvm-project/pull/210952
More information about the Mlir-commits
mailing list