[Mlir-commits] [mlir] 3e223e3 - [mlir][vector] Fix out-of-bounds access (#126734)

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2025-02-13T09:17:43+08:00
New Revision: 3e223e3a202c046b6553aac91d79b6abd089ee8d

URL: https://github.com/llvm/llvm-project/commit/3e223e3a202c046b6553aac91d79b6abd089ee8d
DIFF: https://github.com/llvm/llvm-project/commit/3e223e3a202c046b6553aac91d79b6abd089ee8d.diff

LOG: [mlir][vector] Fix out-of-bounds access (#126734)

This PR fixes an out-of-bounds bug that occurs when there are no overlap
dimensions between the `sizes` and source of
`vector.extract_strided_slice`, causing access to `sizes` to go out of
bounds. Fixes #126196.

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 d36bfddc13279..d5f3634377e4c 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -4023,8 +4023,8 @@ class ContiguousExtractStridedSliceToExtract final
     // Avoid generating slices that have leading unit dimensions. The shape_cast
     // op that we create below would take bad generic fallback patterns
     // (ShapeCastOpRewritePattern).
-    while (sizes[numOffsets] == 1 &&
-           numOffsets < static_cast<int>(sizes.size()) - 1) {
+    while (numOffsets < static_cast<int>(sizes.size()) - 1 &&
+           sizes[numOffsets] == 1) {
       ++numOffsets;
     }
 

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 93581cbfbe5e4..f17d917ca521e 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2984,10 +2984,10 @@ func.func @insert_multiple_poison_idx(%a: vector<4x5x8xf32>, %b: vector<8xf32>)
 
 // -----
 
-// CHECK-LABEL: @contiguous_extract_strided_slices_to_extract
+// CHECK-LABEL: @contiguous_extract_strided_slices_to_extract_sizes_and_outer_source_dims_overlap
 // CHECK:        %[[EXTRACT:.+]] = vector.extract {{.*}}[0, 0, 0, 0, 0] : vector<4xi32> from vector<8x1x2x1x1x4xi32>
 // CHECK-NEXT:   return %[[EXTRACT]] :  vector<4xi32>
-func.func @contiguous_extract_strided_slices_to_extract(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<4xi32> {
+func.func @contiguous_extract_strided_slices_to_extract_sizes_and_outer_source_dims_overlap(%arg0 : vector<8x1x2x1x1x4xi32>) -> vector<4xi32> {
   %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0, 0, 0, 0, 0], sizes = [1, 1, 1, 1, 1, 4], strides = [1, 1, 1, 1, 1, 1]} : vector<8x1x2x1x1x4xi32> to vector<1x1x1x1x1x4xi32>
   %2 = vector.shape_cast %1 : vector<1x1x1x1x1x4xi32> to vector<4xi32>
   return %2 : vector<4xi32>
@@ -2995,6 +2995,17 @@ func.func @contiguous_extract_strided_slices_to_extract(%arg0 : vector<8x1x2x1x1
 
 // -----
 
+// CHECK-LABEL: @contiguous_extract_strided_slices_to_extract_sizes_and_outer_source_dims_no_overlap
+// CHECK:        %[[EXTRACT:.+]] = vector.extract {{.*}}[0, 0] : vector<4xi32> from vector<8x2x4xi32>
+// CHECK-NEXT:   return %[[EXTRACT]] :  vector<4xi32>
+func.func @contiguous_extract_strided_slices_to_extract_sizes_and_outer_source_dims_no_overlap(%arg0 : vector<8x2x4xi32>) -> vector<4xi32> {
+  %1 = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]} : vector<8x2x4xi32> to vector<1x1x4xi32>
+  %2 = vector.shape_cast %1 : vector<1x1x4xi32> to vector<4xi32>
+  return %2 : vector<4xi32>
+}
+
+// -----
+
 // CHECK-LABEL: @contiguous_extract_strided_slices_to_extract_shorter_size_list
 // CHECK:        %[[EXTRACT:.+]] = vector.extract {{.*}}[0, 0, 0, 0] : vector<1x4xi32> from vector<8x1x2x1x1x4xi32>
 // CHECK-NEXT:   return %[[EXTRACT]] :  vector<1x4xi32>