Add a vector.InsertStridedSliceOp

This new op is the counterpart of vector.StridedSliceOp and will be used for in the pattern rewrites for vector unrolling.

PiperOrigin-RevId: 282447414

2 files changed, 56 insertions, 7 deletions

diff --git a/mlir/test/Dialect/VectorOps/invalid.mlir b/mlir/test/Dialect/VectorOps/invalid.mlir
index c0f4d1648aa..60d57740e69 100644
--- a/mlir/test/Dialect/VectorOps/invalid.mlir
+++ b/mlir/test/Dialect/VectorOps/invalid.mlir
@@ -278,6 +278,48 @@ func @test_vector.transfer_write(%arg0: memref<?x?x?xf32>) {
 
 // -----
 
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{expected offsets of same size as destination vector rank}}
+  %1 = vector.insert_strided_slice %a, %b {offsets = [100], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>
+}
+
+// -----
+
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{expected strides of same size as source vector rank}}
+  %1 = vector.insert_strided_slice %a, %b {offsets = [2, 2, 2], strides = [1]} : vector<4x4xf32> into vector<4x8x16xf32>
+}
+
+// -----
+
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{expected source rank to be smaller than destination rank}}
+  %1 = vector.insert_strided_slice %b, %a {offsets = [2, 2], strides = [1, 1, 1]} : vector<4x8x16xf32> into vector<4x4xf32>
+}
+
+// -----
+
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{op expected offsets dimension 0 to be confined to [0, 4)}}
+  %1 = vector.insert_strided_slice %a, %b {offsets = [100,100,100], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>
+}
+
+// -----
+
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{op expected strides to be confined to [1, 2)}}
+  %1 = vector.insert_strided_slice %a, %b {offsets = [2, 2, 2], strides = [100, 100]} : vector<4x4xf32> into vector<4x8x16xf32>
+}
+
+// -----
+
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{op expected sum(offsets, source vector shape) dimension 1 to be confined to [1, 9)}}
+  %1 = vector.insert_strided_slice %a, %b {offsets = [2, 7, 2], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>
+}
+
+// -----
+
 func @strided_slice(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected offsets, sizes and strides attributes of same size}}
   %1 = vector.strided_slice %arg0 {offsets = [100], sizes = [2, 2], strides = [1, 1]} : vector<4x8x16xf32> to vector<2x2x16xf32>
@@ -300,14 +342,14 @@ func @strided_slice(%arg0: vector<4x8x16xf32>) {
 // -----
 
 func @strided_slice(%arg0: vector<4x8x16xf32>) {
-  // expected-error@+1 {{op expected offsets to be confined to [0, 4)}}
+  // expected-error@+1 {{op expected offsets dimension 0 to be confined to [0, 4)}}
   %1 = vector.strided_slice %arg0 {offsets = [100], sizes = [100], strides = [100]} : vector<4x8x16xf32> to vector<100x8x16xf32>
 }
 
 // -----
 
 func @strided_slice(%arg0: vector<4x8x16xf32>) {
-  // expected-error@+1 {{op expected sizes to be confined to [1, 5)}}
+  // expected-error@+1 {{op expected sizes dimension 0 to be confined to [1, 5)}}
   %1 = vector.strided_slice %arg0 {offsets = [2], sizes = [100], strides = [100]} : vector<4x8x16xf32> to vector<100x8x16xf32>
 }
 
@@ -328,7 +370,7 @@ func @strided_slice(%arg0: vector<4x8x16xf32>) {
 // -----
 
 func @strided_slice(%arg0: vector<4x8x16xf32>) {
-  // expected-error@+1 {{op expected sum(offsets, sizes) to be confined to [1, 5)}}
+  // expected-error@+1 {{op expected sum(offsets, sizes) dimension 0 to be confined to [1, 5)}}
   %1 = vector.strided_slice %arg0 {offsets = [2], sizes = [3], strides = [1]} : vector<4x8x16xf32> to vector<3x8x16xf32>
 }
 
diff --git a/mlir/test/Dialect/VectorOps/ops.mlir b/mlir/test/Dialect/VectorOps/ops.mlir
index 09b8815a05c..a2b1ac34142 100644
--- a/mlir/test/Dialect/VectorOps/ops.mlir
+++ b/mlir/test/Dialect/VectorOps/ops.mlir
@@ -22,7 +22,7 @@ func @vector_transfer_ops(%arg0: memref<?x?xf32>) {
   return
 }
 
-// CHECK-LABEL: extractelement
+// CHECK-LABEL: @extractelement
 func @extractelement(%arg0: vector<4x8x16xf32>) -> (vector<8x16xf32>, vector<16xf32>, f32) {
   //      CHECK: vector.extractelement {{.*}}[3 : i32] : vector<4x8x16xf32>
   %1 = vector.extractelement %arg0[3 : i32] : vector<4x8x16xf32>
@@ -33,7 +33,7 @@ func @extractelement(%arg0: vector<4x8x16xf32>) -> (vector<8x16xf32>, vector<16x
   return %1, %2, %3 : vector<8x16xf32>, vector<16xf32>, f32
 }
 
-// CHECK-LABEL: insertelement
+// CHECK-LABEL: @insertelement
 func @insertelement(%a: f32, %b: vector<16xf32>, %c: vector<8x16xf32>, %res: vector<4x8x16xf32>) {
   //      CHECK: vector.insertelement %{{.*}}, %{{.*}}[3 : i32] : vector<8x16xf32> into vector<4x8x16xf32>
   %1 = vector.insertelement %c, %res[3 : i32] : vector<8x16xf32> into vector<4x8x16xf32>
@@ -44,7 +44,7 @@ func @insertelement(%a: f32, %b: vector<16xf32>, %c: vector<8x16xf32>, %res: vec
   return
 }
 
-// CHECK-LABEL: outerproduct
+// CHECK-LABEL: @outerproduct
 func @outerproduct(%arg0: vector<4xf32>, %arg1: vector<8xf32>, %arg2: vector<4x8xf32>) -> vector<4x8xf32> {
   //     CHECK: vector.outerproduct {{.*}} : vector<4xf32>, vector<8xf32>
   %0 = vector.outerproduct %arg0, %arg1 : vector<4xf32>, vector<8xf32>
@@ -53,7 +53,14 @@ func @outerproduct(%arg0: vector<4xf32>, %arg1: vector<8xf32>, %arg2: vector<4x8
   return %1 : vector<4x8xf32>
 }
 
-// CHECK-LABEL: strided_slice
+// CHECK-LABEL: @insert_strided_slice
+func @insert_strided_slice(%a: vector<4x4xf32>, %b: vector<4x8x16xf32>) {
+  //      CHECK: vector.insert_strided_slice %{{.*}}, %{{.*}} {offsets = [2, 2, 2], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>
+  %1 = vector.insert_strided_slice %a, %b {offsets = [2, 2, 2], strides = [1, 1]} : vector<4x4xf32> into vector<4x8x16xf32>
+  return
+}
+
+// CHECK-LABEL: @strided_slice
 func @strided_slice(%arg0: vector<4x8x16xf32>) -> vector<2x2x16xf32> {
   //      CHECK: vector.strided_slice %{{.*}} {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x8x16xf32>
   %1 = vector.strided_slice %arg0 {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x8x16xf32> to vector<2x2x16xf32>