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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Dhairyashil R G (dhairyashilRG)

<details>
<summary>Changes</summary>

`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



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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+5) 
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+44) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 9322c11d401d9..018064077b721 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5377,6 +5377,11 @@ 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;
+  // A scalable dimension holds `vscale` times as many elements as its static
+  // size suggests and `vscale` is only known at run time, so the static size is
+  // just 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..f22dc260bb53f 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1404,6 +1404,50 @@ 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
+}
+
+// -----
+
+// Only the scalable dimension is blocked; a fixed-size dimension alongside it
+// is still folded normally.
+
+// CHECK-LABEL: func @fold_transfer_read_in_bounds_mixed_scalable
+//       CHECK:   vector.transfer_read %{{.*}} {in_bounds = [true, false]
+func.func @fold_transfer_read_in_bounds_mixed_scalable(%m: memref<8x4xf32>, %p: f32) -> vector<4x[4]xf32> {
+  %c0 = arith.constant 0 : index
+  %v = vector.transfer_read %m[%c0, %c0], %p : memref<8x4xf32>, vector<4x[4]xf32>
+  return %v : vector<4x[4]xf32>
+}
+
+// -----
+
 // CHECK-LABEL: fold_vector_transfers
 func.func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) {
   %c0 = arith.constant 0 : index

``````````

</details>


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