[mlir][vector] reject negative strides for vector.load/vector.store (#204611)

This PR follows up #204309 and #204309.

It simply rejects negative strides for vector.load/vector.store :D


AI Disclaimer: I used AI for the tests.

---------

Signed-off-by: Federico Bruzzone <federico.bruzzone.i@gmail.com>
Co-authored-by: Andrzej Warzyński <andrzej.warzynski@gmail.com>

Author: FedericoBruzzone

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

@@ -194,6 +194,10 @@ LogicalResult resolveSourceIndicesRankReducingSubview(
     Location loc, OpBuilder &b, memref::SubViewOp subViewOp, ValueRange indices,
     SmallVectorImpl<Value> &sourceIndices);
 
+/// Returns true if any stride of `memRefTy` is statically known to be
+/// negative.
+bool hasNegativeStaticStride(MemRefType memRefTy);
+
 } // namespace memref
 } // namespace mlir
 

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

@@ -1717,6 +1717,9 @@ def Vector_LoadOp : Vector_Op<"load", [
     %result = vector.load %memref[%i, %j] : memref<200x100xvector<4x8xf32>>, vector<4x8xf32>
     ```
 
+    The memref must have non-negative strides. Negative strides are not supported
+    and will trigger a verification error.
+
     Representation-wise, the 'vector.load' operation permits out-of-bounds
     reads. Support and implementation of out-of-bounds vector loads is
     target-specific. No assumptions should be made on the value of elements
@@ -1835,6 +1838,9 @@ def Vector_StoreOp : Vector_Op<"store", [
     vector.store %valueToStore, %memref[%i, %j] : memref<200x100xvector<4x8xf32>>, vector<4x8xf32>
     ```
 
+    The memref must have non-negative strides. Negative strides are not supported
+    and will trigger a verification error.
+
     Representation-wise, the 'vector.store' operation permits out-of-bounds
     writes. Support and implementation of out-of-bounds vector stores are
     target-specific. No assumptions should be made on the memory written out of

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

@@ -18,6 +18,7 @@
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
@@ -223,16 +224,6 @@ static void replaceLoadOrStoreOp(vector::MaskedStoreOp storeOp,
       storeOp, adaptor.getValueToStore(), ptr, adaptor.getMask(), align);
 }
 
-/// Returns true if all strides of `memRefTy` are static and non-negative. A
-/// negative (or dynamic, hence unknown-sign) stride would make `mul nuw` on the
-/// index arithmetic wrap, so `nuw` must not be emitted in that case.
-static bool hasNonNegativeStrides(MemRefType memRefTy) {
-  auto [strides, offset] = memRefTy.getStridesAndOffset();
-  return llvm::all_of(strides, [](int64_t stride) {
-    return !ShapedType::isDynamic(stride) && stride >= 0;
-  });
-}
-
 /// Conversion pattern for a vector.load, vector.store, vector.maskedload, and
 /// vector.maskedstore.
 template <class LoadOrStoreOp>
@@ -283,10 +274,12 @@ class VectorLoadStoreConversion : public ConvertOpToLLVMPattern<LoadOrStoreOp> {
              "vector.load/store requires unit trailing memref stride");
       if (enableGEPInboundsNuw) {
         noWrapFlags = noWrapFlags | LLVM::GEPNoWrapFlags::inbounds;
-        // `nuw` additionally requires non-negative strides; skip it when the
-        // memref has dynamic or negative strides to avoid emitting poison.
-        if (hasNonNegativeStrides(memRefTy))
-          noWrapFlags = noWrapFlags | LLVM::GEPNoWrapFlags::nuw;
+
+        // `nuw` additionally requires non-negative strides.
+        assert(
+            !(memref::hasNegativeStaticStride(memRefTy)) &&
+            "Invalid MemRef type - should have been rejected by Op verifier.");
+        noWrapFlags = noWrapFlags | LLVM::GEPNoWrapFlags::nuw;
       }
     }
     auto vtype = cast<VectorType>(

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

@@ -344,5 +344,12 @@ LogicalResult resolveSourceIndicesRankReducingSubview(
   return success();
 }
 
+bool hasNegativeStaticStride(MemRefType memRefTy) {
+  auto [strides, offset] = memRefTy.getStridesAndOffset();
+  return llvm::any_of(strides, [](int64_t stride) {
+    return ShapedType::isStatic(stride) && stride < 0;
+  });
+}
+
 } // namespace memref
 } // namespace mlir

mlir/lib/Dialect/Vector/IR/CMakeLists.txt

@@ -24,6 +24,7 @@ add_mlir_dialect_library(MLIRVectorDialect
   MLIRMaskableOpInterface
   MLIRMaskingOpInterface
   MLIRMemRefDialect
+  MLIRMemRefUtils
   MLIRSideEffectInterfaces
   MLIRTensorDialect
   MLIRUBDialect

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -20,6 +20,7 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/UB/IR/UBMatchers.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
@@ -6194,6 +6195,10 @@ LogicalResult vector::LoadOp::verify() {
   if (failed(verifyLoadStoreMemRefLayout(*this, resVecTy, memRefTy)))
     return failure();
 
+  // Negative strides are not supported on vector.load.
+  if (memref::hasNegativeStaticStride(memRefTy))
+    return emitOpError("memref strides must be non-negative");
+
   if (memRefTy.getRank() < resVecTy.getRank())
     return emitOpError(
         "destination memref has lower rank than the result vector");
@@ -6240,6 +6245,10 @@ LogicalResult vector::StoreOp::verify() {
   if (failed(verifyLoadStoreMemRefLayout(*this, valueVecTy, memRefTy)))
     return failure();
 
+  // Negative strides are not supported on vector.store.
+  if (memref::hasNegativeStaticStride(memRefTy))
+    return emitOpError("memref strides must be non-negative");
+
   if (memRefTy.getRank() < valueVecTy.getRank())
     return emitOpError("source memref has lower rank than the vector to store");
 

mlir/test/Dialect/Vector/invalid.mlir

@@ -2145,6 +2145,15 @@ func.func @store_non_unit_stride(%src : memref<?xi8, strided<[2], offset:?>>,%va
 
 // -----
 
+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
+  vector.store %val, %src[%c0, %c0] : memref<100x100xf32, strided<[-100, 1]>>, vector<4xf32>
+  return
+}
+
+// -----
+
 // Verify that vector.bitcast rejects vectors with i0 (zero-bitwidth) element type.
 func.func @bitcast_i0(%a: vector<4xi0>) -> vector<4xi0> {
   // expected-error @+1 {{'vector.bitcast' op operand #0 must be vector of non-zero-bitwidth type values, but got 'vector<4xi0>'}}
@@ -2198,3 +2207,12 @@ func.func @scan_i0(%a: vector<4xi0>, %init: vector<1xi0>) -> (vector<4xi0>, vect
     vector<4xi0>, vector<1xi0>
   return %0#0, %0#1 : vector<4xi0>, vector<1xi0>
 }
+
+// -----
+
+func.func @load_negative_stride(%src: memref<100x100xf32, strided<[-100, 1]>>) -> vector<8xf32> {
+  // expected-error @+2 {{'vector.load' op memref strides must be non-negative}}
+  %c0 = arith.constant 0 : index
+  %v = vector.load %src[%c0, %c0] : memref<100x100xf32, strided<[-100, 1]>>, vector<8xf32>
+  return %v : vector<8xf32>
+}