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

[llvmlistbot at llvm.org]

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

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

>From 98ccee9c44232e46b54b643535e0111bdfa929ab Mon Sep 17 00:00:00 2001
From: woruyu <1214539920 at qq.com>
Date: Tue, 3 Nov 2026 14:19:27 +0800
Subject: [PATCH] [mlir][Vector] Fix out-of-bounds access in
 StridedSliceBroadcast

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   |  3 +-
 mlir/test/Dialect/Vector/canonicalize.mlir | 34 ++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

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>