address some of reviewer's comments

Author: newling

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

@@ -143,7 +143,10 @@ static FailureOr<SmallVector<int64_t>> intsFromArrayAttr(ArrayAttr attrs) {
 /// Example: insert a 1x2 vector into a 4x5 vector at position (1,3). The 2
 /// positions written to are (1,3) and (1,4), which have linearized indices 8
 /// and 9. So [8,9] is returned.
-SmallVector<int64_t> static getFlattenedStridedSliceIndices(
+///
+/// The length of the returned vector is equal to the number of elements in
+/// the shape `small` (i.e. the product of dimensions of `small`).
+SmallVector<int64_t> static getStridedSliceInsertionIndices(
     ArrayRef<int64_t> small, ArrayRef<int64_t> large,
     ArrayRef<int64_t> offsets) {
 
@@ -153,8 +156,10 @@ SmallVector<int64_t> static getFlattenedStridedSliceIndices(
   //  offsets  =  2, 3, 0
   //
   // `offsets` has implicit trailing 0s, `small` has implicit leading 1s.
-  assert(large.size() >= small.size());
-  assert(large.size() >= offsets.size());
+  assert((large.size() >= small.size()) &&
+         "rank of 'large' cannot be lower than rank of 'small'");
+  assert((large.size() >= offsets.size()) &&
+         "rank of 'large' cannot be lower than the number of offsets");
   unsigned delta = large.size() - small.size();
   unsigned nOffsets = offsets.size();
   auto getSmall = [&](int64_t i) { return i >= delta ? small[i - delta] : 1; };
@@ -223,10 +228,8 @@ struct LinearizeVectorExtractStridedSlice final
         extractStridedSliceOp.getType());
     assert(flatOutputType && "vector type expected");
 
-    if (!stridesAllOne(extractStridedSliceOp)) {
-      return rewriter.notifyMatchFailure(extractStridedSliceOp,
-                                         "strides other than 1 not supported");
-    }
+    assert(stridesAllOne(extractStridedSliceOp) &&
+           "has extract_strided_slice's verifier not checked strides are 1?");
 
     FailureOr<SmallVector<int64_t>> offsets =
         intsFromArrayAttr(extractStridedSliceOp.getOffsets());
@@ -240,7 +243,7 @@ struct LinearizeVectorExtractStridedSlice final
 
     ArrayRef<int64_t> outputShape = extractStridedSliceOp.getType().getShape();
 
-    SmallVector<int64_t> indices = getFlattenedStridedSliceIndices(
+    SmallVector<int64_t> indices = getStridedSliceInsertionIndices(
         outputShape, inputShape, offsets.value());
 
     Value srcVector = adaptor.getVector();
@@ -287,10 +290,9 @@ struct LinearizeVectorInsertStridedSlice final
                   OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
 
-    if (!stridesAllOne(insertStridedSliceOp)) {
-      return rewriter.notifyMatchFailure(insertStridedSliceOp,
-                                         "strides other than 1 not supported");
-    }
+    // See InsertStridedSliceOp's verify method.
+    assert(stridesAllOne(insertStridedSliceOp) &&
+           "has insert_strided_slice's verifier not checked strides are 1?");
 
     VectorType inputType = insertStridedSliceOp.getValueToStore().getType();
     ArrayRef<int64_t> inputShape = inputType.getShape();
@@ -305,10 +307,10 @@ struct LinearizeVectorInsertStridedSlice final
       return rewriter.notifyMatchFailure(insertStridedSliceOp,
                                          "failed to get integer offsets");
     }
-    SmallVector<int64_t> sliceIndices = getFlattenedStridedSliceIndices(
+    SmallVector<int64_t> sliceIndices = getStridedSliceInsertionIndices(
         inputShape, outputShape, offsets.value());
 
-    SmallVector<int64_t> indices(nOutputElements, 0);
+    SmallVector<int64_t> indices(nOutputElements);
     std::iota(indices.begin(), indices.end(), 0);
     for (auto [index, sliceIndex] : llvm::enumerate(sliceIndices)) {
       indices[sliceIndex] = index + nOutputElements;

mlir/test/Dialect/Vector/linearize.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -split-input-file -test-vector-linearize -verify-diagnostics | FileCheck %s 
+// RUN: mlir-opt %s -split-input-file -test-vector-linearize -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: test_linearize
 // CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>)
@@ -153,7 +153,7 @@ func.func @test_extract_strided_slice_2D_scalable(%arg0: vector<4x[8]xf32>) -> v
 
   // CHECK-NOT: vector.shuffle
   // CHECK-NOT: vector.shape_cast
-  // CHECK: %[[RES:.*]] = vector.extract_strided_slice %[[VAL_0]] 
+  // CHECK: %[[RES:.*]] = vector.extract_strided_slice %[[VAL_0]]
   %0 = vector.extract_strided_slice %arg0 { sizes = [2, 8], strides = [1, 1], offsets = [1, 0] } : vector<4x[8]xf32> to vector<2x[8]xf32>
 
   // CHECK: return %[[RES]] : vector<2x[8]xf32>
@@ -179,7 +179,7 @@ func.func @test_extract_strided_slice_3D(%arg0 : vector<2x8x2xf32>) -> vector<1x
 // -----
 
 // Test of insert_strided_slice -> shuffle.
-// This is a contiguous insertion of 4 elements at offset 6 into a vector of 12 elements. 
+// This is a contiguous insertion of 4 elements at offset 6 into a vector of 12 elements.
 // CHECK-LABEL: insert_strided_slice_2D_into_4D
 func.func @insert_strided_slice_2D_into_4D(%arg0 : vector<2x2xi8>, %arg1 : vector<2x1x3x2xi8>) -> vector<2x1x3x2xi8> {
 
@@ -196,19 +196,19 @@ func.func @insert_strided_slice_2D_into_4D(%arg0 : vector<2x2xi8>, %arg1 : vecto
 
 // -----
 
-// Test of insert_strided_slice -> shuffle. 
+// Test of insert_strided_slice -> shuffle.
 // [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]], [[12, 13], [14, 15]], [[16, 17]]]
 //                                         ^         ^
 //                                         |         |
 //                          where the 2 elements are inserted into the 3x3x2 vector
 // CHECK-LABEL: insert_strided_slice_3D
-func.func @insert_strided_slice_3D(%arg0 : vector<1x2x1xi8>, %arg2 : vector<3x3x2xi8>) -> vector<3x3x2xi8> {
+func.func @insert_strided_slice_3D(%arg0 : vector<1x2x1xi8>, %arg1 : vector<3x3x2xi8>) -> vector<3x3x2xi8> {
 
 //   CHECK-DAG:     %[[ARG0:.*]] = vector.shape_cast {{.*}}  to vector<2xi8>
 //   CHECK-DAG:     %[[ARG1:.*]] = vector.shape_cast {{.*}}  to vector<18xi8>
 //       CHECK:     vector.shuffle %[[ARG1]], %[[ARG0]]
 //  CHECK-SAME:       [0, 1, 2, 3, 4, 5, 6, 7, 8, 18, 10, 19, 12, 13, 14, 15, 16, 17] : vector<18xi8>, vector<2xi8>
-  %0 = vector.insert_strided_slice %arg0, %arg2 {offsets = [1, 1, 1], sizes = [1, 2, 1], strides = [1, 1, 1]} : vector<1x2x1xi8> into vector<3x3x2xi8>
+  %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [1, 1, 1], sizes = [1, 2, 1], strides = [1, 1, 1]} : vector<1x2x1xi8> into vector<3x3x2xi8>
 
 //       CHECK:     %[[RES:.*]] = vector.shape_cast {{.*}} to vector<3x3x2xi8>
 //       CHECK:     return %[[RES]] : vector<3x3x2xi8>
@@ -217,6 +217,37 @@ func.func @insert_strided_slice_3D(%arg0 : vector<1x2x1xi8>, %arg2 : vector<3x3x
 
 // -----
 
+// CHECK-LABEL: insert_strided_slice_2D_higher_offsets
+func.func @insert_strided_slice_2D_higher_offsets(%arg0 : vector<2x1xi8>, %arg1 : vector<2x2xi8>, %arg2 : vector<5x2xi8>) -> vector<5x2xi8> {
+
+  // CHECK: [0, 1, 2, 3, 10, 11, 12, 13, 8, 9]
+  //                     ^^^ ^^^ ^^^ ^^^
+  //                    insertion indices
+  %0 = vector.insert_strided_slice %arg1, %arg2 {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<2x2xi8> into vector<5x2xi8>
+
+  // CHECK: [0, 1, 2, 3, 10, 5, 11, 7, 8, 9]
+  //                     ^^^    ^^^
+  %1 = vector.insert_strided_slice %arg0, %0 {offsets = [2, 0], sizes = [2, 1], strides = [1, 1]} : vector<2x1xi8> into vector<5x2xi8>
+
+  // CHECK: [0, 1, 2, 3, 4, 5, 6, 10, 8, 11]
+  //                              ^^^    ^^^
+  %2 = vector.insert_strided_slice %arg0, %1 {offsets = [3, 1], sizes = [2, 1], strides = [1, 1]} : vector<2x1xi8> into vector<5x2xi8>
+
+  return %2 : vector<5x2xi8>
+}
+
+// -----
+
+// CHECK-LABEL: negative_insert_strided_slice_scalable
+// CHECK-NOT:   vector.shuffle
+// CHECK:       return
+func.func @negative_insert_strided_slice_scalable(%arg0 : vector<1x[2]xi8>, %arg1 : vector<2x[2]xi8>) -> vector<2x[2]xi8> {
+  %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0], strides = [1,1]} : vector<1x[2]xi8> into vector<2x[2]xi8>
+  return %0 : vector<2x[2]xi8>
+}
+
+// -----
+
 // CHECK-LABEL: test_vector_shuffle
 // CHECK-SAME: (%[[ORIG_ARG0:.*]]: vector<4x2xf32>, %[[ORIG_ARG1:.*]]: vector<4x2xf32>) -> vector<8x2xf32> {
 func.func @test_vector_shuffle(%arg0: vector<4x2xf32>, %arg1: vector<4x2xf32>) -> vector<8x2xf32> {