[Mlir-commits] [mlir] [mlir][vector] Verify non-unit strides on `masked/expand/compress` ops and fix `SparseVectorization` miscompilation (PR #210952)
Federico Bruzzone
llvmlistbot at llvm.org
Tue Jul 21 07:04:36 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/210952
>From ee815c22e10991ea39025101b8bc71d00558f0a4 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 21 Jul 2026 15:41:28 +0200
Subject: [PATCH 1/2] [mlir][sparse] Avoid vectorizing non-contiguous COO
coordinate loads
SparseVectorization emits vector.maskedload/maskedstore for direct
loop accesses, assuming consecutive loop indices map to consecutive
memory. This is false for sparse_tensor.coordinates of a level inside
a trailing AoS COO region, whose buffer is interleaved with other
levels -- a real, silent miscompile.
Fall back to a scalar load/store when the stride is statically known
to be non-unit, including cases only derivable from the sparse
tensor's encoding before sparse-tensor-codegen materializes the
concrete stride.
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Transforms/SparseVectorization.cpp | 50 +++++++++++++++++++
1 file changed, 50 insertions(+)
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();
>From 982fc5dcf309734db2d09dc4bf16cb5f313f5630 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 21 Jul 2026 15:41:54 +0200
Subject: [PATCH 2/2] [mlir][vector] Reject non-unit strides on
masked/expand/compress ops
vector.maskedload/maskedstore/expandload/compressstore lower to LLVM
masked intrinsics that read/write N *consecutive* elements from a
single pointer, but none of them verified the memref's minor-dim
stride, so e.g. `strided<[2]>` verified successfully and silently
miscompiled.
Reject statically-known non-unit strides; a dynamic stride is still
accepted (see linked issue for discussion on why this differs from
vector.load/store). expandload/compressstore also gain the
negative-stride check vector.load/store already have.
Depends on the SparseVectorization fix in the preceding commit: the
sparsifier's vectorizer previously relied on maskedload having no
stride check to (unsafely) vectorize non-contiguous COO coordinate
buffers.
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 51 ++++++++++++++++
mlir/test/Dialect/Vector/invalid.mlir | 77 ++++++++++++++++++++++++
2 files changed, 128 insertions(+)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f37083803a2a1..8b594c378912e 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6190,6 +6190,33 @@ static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,
return success();
}
+/// Verifies that `memRefTy`'s most minor dimension does not have a
+/// statically known non-unit stride; a dynamic (not provably unit) stride
+/// is accepted.
+///
+/// This is more permissive than vector.load/store's stride check: ops such
+/// as vector.maskedload/maskedstore and vector.expandload/compressstore are
+/// also used on memrefs whose most minor dimension is contiguous at runtime
+/// but not provable as such at the type level (e.g., buffers produced by
+/// the sparsifier).
+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 +6326,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 +6395,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 +6681,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 +6744,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..9ec17024bc1d3 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} :
@@ -2179,6 +2233,19 @@ func.func @load_non_unit_stride(%src : memref<?xi8, strided<[2], offset: ?>>) {
// -----
+// Unlike vector.maskedload/maskedstore/expandload/compressstore, a dynamic
+// (unprovable) stride is rejected here too: vector.load/store require a
+// statically known unit stride, with no exception for strides that merely
+// aren't provably non-unit.
+func.func @load_dynamic_stride(%src : memref<?xi8, strided<[?], offset: ?>>) {
+ %c0 = arith.constant 0 : index
+ // expected-error @+1 {{'vector.load' op most minor memref dim must have unit stride}}
+ %0 = vector.load %src[%c0] : memref<?xi8, strided<[?], offset: ?>>, vector<16xi8>
+ return
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// vector.store
//===----------------------------------------------------------------------===//
@@ -2215,6 +2282,16 @@ func.func @store_non_unit_stride(%src : memref<?xi8, strided<[2], offset:?>>,%va
// -----
+// See load_dynamic_stride above: vector.store also rejects a dynamic stride,
+// unlike vector.maskedstore.
+func.func @store_dynamic_stride(%src : memref<?xi8, strided<[?], offset: ?>>, %val : vector<16xi8>, %c0: index) {
+ // expected-error @below {{'vector.store' op most minor memref dim must have unit stride}}
+ vector.store %val, %src[%c0] : memref<?xi8, strided<[?], offset: ?>>, vector<16xi8>
+ return
+}
+
+// -----
+
func.func @store_negative_stride(%src: memref<100x100xf32, strided<[-100, 1]>>, %val: vector<4xf32>) {
// expected-error @+2 {{'vector.store' op memref strides must be non-negative}}
%c0 = arith.constant 0 : index
More information about the Mlir-commits
mailing list