[Mlir-commits] [mlir] [mlir][Vector] Fix out-of-bounds access in StridedSliceBroadcast (PR #238219)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: woruyu (woruyu)

<details>
<summary>Changes</summary>

This fixes https://github.com/llvm/llvm-project/issues/237496. `StridedSliceBroadcast` assumes slice attributes cover every source dimension. When a slice specifies only a prefix of the dimensions, adjusting a trailing unit dimension can access past the attribute arrays and trigger an assertion.Bound the adjustment loop by the slice attribute count, preserving unspecified trailing dimensions. 

Assisted-by: OpenAI Codex

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+2-1) 
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+34) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 7945ca3ea79634..a3abf75d7503e8 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -4960,7 +4960,8 @@ class StridedSliceBroadcast final
           getI64SubArray(op.getOffsets(), /*dropFront=*/rankDiff);
       SmallVector<int64_t> sizes =
           getI64SubArray(op.getSizes(), /*dropFront=*/rankDiff);
-      for (unsigned i = 0; i < srcRank; i++) {
+      // Slice attributes may only cover a prefix of the source dimensions.
+      for (unsigned i = 0; i < offsets.size(); i++) {
         if (srcVecType.getDimSize(i) == 1) {
           // In case this dimension was broadcasted *and* sliced, the offset
           // and size need to be updated now that there is no broadcast before
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 059d6f31260c3d..43ad8059f37c3c 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1773,6 +1773,40 @@ func.func @extract_strided_broadcast5(%arg0: vector<2x1xf32>) -> vector<2x4xf32>
 
 // -----
 
+// Slice attributes only cover the leading dimension. The trailing unit source
+// dimension is broadcasted but must not be accessed in the slice attributes.
+// CHECK-LABEL: func @extract_strided_broadcast_prefix(
+//  CHECK-SAME: %[[ARG:.+]]: vector<2x1xf32>
+//       CHECK: %[[E:.+]] = vector.extract %[[ARG]][1] : vector<1xf32> from vector<2x1xf32>
+//  CHECK-NEXT: %[[B:.+]] = vector.broadcast %[[E]] : vector<1xf32> to vector<1x4xf32>
+//  CHECK-NEXT: return %[[B]] : vector<1x4xf32>
+func.func @extract_strided_broadcast_prefix(%arg0: vector<2x1xf32>) -> vector<1x4xf32> {
+  %0 = vector.broadcast %arg0 : vector<2x1xf32> to vector<2x4xf32>
+  %1 = vector.extract_strided_slice %0
+      offsets = [1], sizes = [1], strides = [1]
+      : vector<2x4xf32> to vector<1x4xf32>
+  return %1 : vector<1x4xf32>
+}
+
+// -----
+
+// Drop the added leading dimension and adjust the sliced unit source dimension,
+// leaving the trailing source dimension outside the slice attributes.
+// CHECK-LABEL: func @extract_strided_broadcast_prefix_rank_increase(
+//  CHECK-SAME: %[[ARG:.+]]: vector<4x1x1xf32>
+//       CHECK: %[[E:.+]] = vector.extract_strided_slice %[[ARG]] offsets = [1, 0], sizes = [2, 1], strides = [1, 1] : vector<4x1x1xf32> to vector<2x1x1xf32>
+//  CHECK-NEXT: %[[B:.+]] = vector.broadcast %[[E]] : vector<2x1x1xf32> to vector<1x2x2x5xf32>
+//  CHECK-NEXT: return %[[B]] : vector<1x2x2x5xf32>
+func.func @extract_strided_broadcast_prefix_rank_increase(%arg0: vector<4x1x1xf32>) -> vector<1x2x2x5xf32> {
+  %0 = vector.broadcast %arg0 : vector<4x1x1xf32> to vector<2x4x3x5xf32>
+  %1 = vector.extract_strided_slice %0
+      offsets = [1, 1, 1], sizes = [1, 2, 2], strides = [1, 1, 1]
+      : vector<2x4x3x5xf32> to vector<1x2x2x5xf32>
+  return %1 : vector<1x2x2x5xf32>
+}
+
+// -----
+
 // CHECK-LABEL: consecutive_shape_cast
 //       CHECK:   %[[C:.*]] = vector.shape_cast %{{.*}} : vector<16xf16> to vector<4x4xf16>
 //  CHECK-NEXT:   return %[[C]] : vector<4x4xf16>

``````````

</details>


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