[Mlir-commits] [mlir] 57227c7 - [mlir][vector] Don't fold in_bounds for scalable vector dimensions (#213506)

[llvmlistbot at llvm.org]

Author: Dhairyashil R G
Date: 2026-08-10T09:49:23+01:00
New Revision: 57227c70be689b1d81354adadcc01bb292da939a

URL: https://github.com/llvm/llvm-project/commit/57227c70be689b1d81354adadcc01bb292da939a
DIFF: https://github.com/llvm/llvm-project/commit/57227c70be689b1d81354adadcc01bb292da939a.diff

LOG: [mlir][vector] Don't fold in_bounds for scalable vector dimensions (#213506)

`isInBounds` compares `index + getVectorType().getDimSize(resultIdx)`
against the static size of the corresponding source dimension. For a
scalable dimension the vector holds `vscale * getDimSize(resultIdx)`
elements, so the static size is only a lower bound and the comparison
can succeed for a transfer that is actually out of bounds.

For example, reading `vector<[4]xf32>` at index 0 of a `memref<4xf32>`
folds to `in_bounds = [true]` because `0 + 4 <= 4`. Lowering then turns
a predicated masked load into a plain full-width load:

```mlir
// Without the fold: predicated, lanes >= 4 are inactive.
%8  = llvm.intr.stepvector : vector<[4]xi32>
%16 = llvm.icmp "slt" %8, %15 : vector<[4]xi32>
%23 = llvm.intr.masked.load %22, %16, %20 {alignment = 4 : i32} :
        (!llvm.ptr, vector<[4]xi1>, vector<[4]xf32>) -> vector<[4]xf32>

// With the fold: unpredicated, reads `4 * vscale` elements.
%9  = llvm.load %8 {alignment = 4 : i64} : !llvm.ptr -> vector<[4]xf32>
```

which reaches AArch64 as `ld1w { z1.s }, p1/z, [x1]` and `ldr z1, [x1]`
respectively. For every `vscale > 1` the folded form reads past the end
of the allocation. The same applies to `vector.transfer_write`, where
the fold turns a masked store into an out-of-bounds one.

Bail out for scalable result dimensions. Fixed-size dimensions of a
partly scalable vector are unaffected and still fold; the added
`vector<4x[4]xf32>` test covers that.

Confirmed on SVE hardware (AWS Graviton3, 256-bit vectors, vscale = 2)
with the 4-element buffer placed flush against a PROT_NONE guard page:

  predicated   ld1w -> returns 10.0, exit 0
  unpredicated ldr  -> SIGSEGV, exit 139

Assisted-by: Claude

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index e169ca4fe0c82..170c0070ff07f 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5377,6 +5377,10 @@ static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
   // op.getIndices()[indicesIdx] + vectorType < dim(op.getSource(), indicesIdx)
   if (op.getShapedType().isDynamicDim(indicesIdx))
     return false;
+  // Scalable dimensions are `vscale` times larger at runtime, so the static
+  // size is only a lower bound and cannot prove that the transfer fits.
+  if (op.getVectorType().getScalableDims()[resultIdx])
+    return false;
   Value index = op.getIndices()[indicesIdx];
   std::optional<int64_t> cstOp = getConstantIntValue(index);
   if (!cstOp.has_value())

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index ed4c908c6e5f2..bb1729cd7c3bc 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1404,6 +1404,37 @@ func.func @fold_vector_transfer_masks(%A: memref<?x?xf32>) -> (vector<4x8xf32>,
 
 // -----
 
+// A scalable vector dimension holds `vscale * N` elements, so the static size N
+// is only a lower bound and cannot prove that the transfer is in bounds. Here
+// `vector<[4]xf32>` reads `4 * vscale` elements from a 4-element memref, which
+// is out of bounds for every `vscale > 1`.
+
+// CHECK-LABEL: func @no_fold_transfer_read_in_bounds_scalable
+//       CHECK:   vector.transfer_read
+//   CHECK-NOT:   in_bounds
+//       CHECK:   : memref<4xf32>, vector<[4]xf32>
+func.func @no_fold_transfer_read_in_bounds_scalable(%m: memref<4xf32>, %p: f32) -> vector<[4]xf32> {
+  %c0 = arith.constant 0 : index
+  %v = vector.transfer_read %m[%c0], %p : memref<4xf32>, vector<[4]xf32>
+  return %v : vector<[4]xf32>
+}
+
+// -----
+
+// Same for the write path, where an unsound fold is an out-of-bounds store.
+
+// CHECK-LABEL: func @no_fold_transfer_write_in_bounds_scalable
+//       CHECK:   vector.transfer_write
+//   CHECK-NOT:   in_bounds
+//       CHECK:   : vector<[4]xf32>, memref<4xf32>
+func.func @no_fold_transfer_write_in_bounds_scalable(%m: memref<4xf32>, %v: vector<[4]xf32>) {
+  %c0 = arith.constant 0 : index
+  vector.transfer_write %v, %m[%c0] : vector<[4]xf32>, memref<4xf32>
+  return
+}
+
+// -----
+
 // CHECK-LABEL: fold_vector_transfers
 func.func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
   %c0 = arith.constant 0 : index