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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Federico Bruzzone (FedericoBruzzone)

<details>
<summary>Changes</summary>

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.


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+38) 
- (modified) mlir/test/Dialect/Vector/invalid.mlir (+38) 


``````````diff
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')}}

``````````

</details>


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