[Mlir-commits] [mlir] [mlir][vector] add consistent stride verification to `masked load/store` and `gather/scatter` ops (PR #204842)

[Federico Bruzzone]

https://github.com/FedericoBruzzone created https://github.com/llvm/llvm-project/pull/204842

Extend negative stride checks to MaskedLoadOp, MaskedStoreOp, GatherOp, and ScatterOp to match LoadOp and StoreOp behavior (see #204611).

Depends indirectly on: #204611.

AI Disclaimer: I used AI for the tests.


>From b7d7f3e4e5d0ab52aeffb057a61dc5ec3b9a5516 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Fri, 19 Jun 2026 16:17:54 +0200
Subject: [PATCH] Add consistent stride verification to masked load/store and
 gather/scatter ops

Extend negative stride checks to MaskedLoadOp, MaskedStoreOp, GatherOp, and
ScatterOp to match LoadOp and StoreOp behavior. This ensures all vector
memory operations consistently reject negative strides that would cause
undefined behavior during lowering (e.g., poison from nuw arithmetic).

Also add unit stride verification to masked operations for consistency.

Co-Authored-By: Claude Haiku 4.5 <noreply at anthropic.com>
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 38 ++++++++++++++++++++++++
 mlir/test/Dialect/Vector/invalid.mlir    | 38 ++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 67c31730f4b65..edacca4255cd6 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6284,6 +6284,16 @@ LogicalResult MaskedLoadOp::verify() {
   VectorType resVType = getVectorType();
   MemRefType memType = getMemRefType();
 
+  if (failed(verifyLoadStoreMemRefLayout(*this, resVType, memType)))
+    return failure();
+
+  // Negative strides are not supported on vector.maskedload.
+  auto [strides, offset] = memType.getStridesAndOffset();
+  for (int64_t stride : strides) {
+    if (ShapedType::isStatic(stride) && stride < 0)
+      return emitOpError("memref strides must be non-negative");
+  }
+
   if (failed(
           verifyElementTypesMatch(*this, memType, resVType, "base", "result")))
     return failure();
@@ -6344,6 +6354,16 @@ LogicalResult MaskedStoreOp::verify() {
   VectorType valueVType = getVectorType();
   MemRefType memType = getMemRefType();
 
+  if (failed(verifyLoadStoreMemRefLayout(*this, valueVType, memType)))
+    return failure();
+
+  // Negative strides are not supported on vector.maskedstore.
+  auto [strides, offset] = memType.getStridesAndOffset();
+  for (int64_t stride : strides) {
+    if (ShapedType::isStatic(stride) && stride < 0)
+      return emitOpError("memref strides must be non-negative");
+  }
+
   if (failed(verifyElementTypesMatch(*this, memType, valueVType, "base",
                                      "valueToStore")))
     return failure();
@@ -6405,6 +6425,15 @@ LogicalResult GatherOp::verify() {
   if (!llvm::isa<MemRefType, RankedTensorType>(baseType))
     return emitOpError("requires base to be a memref or ranked tensor type");
 
+  // Negative strides are not supported on vector.gather.
+  if (auto memRefType = dyn_cast<MemRefType>(baseType)) {
+    auto [strides, offset] = memRefType.getStridesAndOffset();
+    for (int64_t stride : strides) {
+      if (ShapedType::isStatic(stride) && stride < 0)
+        return emitOpError("memref strides must be non-negative");
+    }
+  }
+
   if (failed(
           verifyElementTypesMatch(*this, baseType, resVType, "base", "result")))
     return failure();
@@ -6519,6 +6548,15 @@ LogicalResult ScatterOp::verify() {
   if (!llvm::isa<MemRefType, RankedTensorType>(baseType))
     return emitOpError("requires base to be a memref or ranked tensor type");
 
+  // Negative strides are not supported on vector.scatter.
+  if (auto memRefType = dyn_cast<MemRefType>(baseType)) {
+    auto [strides, offset] = memRefType.getStridesAndOffset();
+    for (int64_t stride : strides) {
+      if (ShapedType::isStatic(stride) && stride < 0)
+        return emitOpError("memref strides must be non-negative");
+    }
+  }
+
   if (failed(verifyElementTypesMatch(*this, baseType, valueVType, "base",
                                      "valueToStore")))
     return failure();
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 2fed3002596a3..50bc17a102211 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1413,6 +1413,15 @@ func.func @maskedload_memref_mismatch(%base: memref<?xf32>, %mask: vector<16xi1>
 
 // -----
 
+func.func @maskedload_negative_stride(%flip: memref<100x100xf32, strided<[-100, 1]>>, %mask: vector<8xi1>, %pass: vector<8xf32>) -> vector<8xf32> {
+  // expected-error @+2 {{'vector.maskedload' op memref strides must be non-negative}}
+  %c0 = arith.constant 0 : index
+  %0 = vector.maskedload %flip[%c0, %c0], %mask, %pass : memref<100x100xf32, strided<[-100, 1]>>, vector<8xi1>, vector<8xf32> into vector<8xf32>
+  return %0 : vector<8xf32>
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // vector.maskedstore
 //===----------------------------------------------------------------------===//
@@ -1457,6 +1466,15 @@ func.func @maskedstore_memref_mismatch(%base: memref<?xf32>, %mask: vector<16xi1
 
 // -----
 
+func.func @maskedstore_negative_stride(%flip: memref<100x100xf32, strided<[-100, 1]>>, %mask: vector<8xi1>, %value: vector<8xf32>) {
+  // expected-error @+2 {{'vector.maskedstore' op memref strides must be non-negative}}
+  %c0 = arith.constant 0 : index
+  vector.maskedstore %flip[%c0, %c0], %mask, %value : memref<100x100xf32, strided<[-100, 1]>>, vector<8xi1>, vector<8xf32>
+  return
+}
+
+// -----
+
 func.func @gather_from_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
                                 %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
   %c0 = arith.constant 0 : index
@@ -1554,6 +1572,16 @@ func.func @gather_tensor_alignment(%base: tensor<16xf32>, %indices: vector<16xi3
 
 // -----
 
+func.func @gather_negative_stride(%flip: memref<100x100xf32, strided<[-100, 1]>>, %indices: vector<16xi32>,
+                                  %mask: vector<16xi1>, %pass_thru: vector<16xf32>, %c0 : index) -> vector<16xf32> {
+  // expected-error @+1 {{'vector.gather' op memref strides must be non-negative}}
+  %0 = vector.gather %flip[%c0, %c0][%indices], %mask, %pass_thru
+    : memref<100x100xf32, strided<[-100, 1]>>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
+  return %0 : vector<16xf32>
+}
+
+// -----
+
 func.func @scatter_to_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
                              %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
   %c0 = arith.constant 0 : index
@@ -1642,6 +1670,16 @@ func.func @scatter_tensor_alignment(%base: tensor<?xf32>, %indices: vector<16xi3
 
 // -----
 
+func.func @scatter_negative_stride(%flip: memref<100x100xf32, strided<[-100, 1]>>, %indices: vector<16xi32>,
+                                   %mask: vector<16xi1>, %value: vector<16xf32>, %c0: index) {
+  // expected-error @+1 {{'vector.scatter' op memref strides must be non-negative}}
+  vector.scatter %flip[%c0, %c0][%indices], %mask, %value
+    : memref<100x100xf32, strided<[-100, 1]>>, vector<16xi32>, vector<16xi1>, vector<16xf32>
+  return
+}
+
+// -----
+
 func.func @expand_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
   %c0 = arith.constant 0 : index
   // expected-error at +1 {{'vector.expandload' op base element type ('f64') does not match result element type ('f32')}}