Adds ExtractSlicesOp to the VectorOps dialect.

ExtractSlicesOp extracts slices of its vector operand and with a specified tiling scheme.
This operation centralizes the tiling scheme around a single op, which simplifies vector op unrolling and subsequent pattern rewrite transformations.

PiperOrigin-RevId: 285761129

2 files changed, 68 insertions, 0 deletions

diff --git a/mlir/test/Dialect/VectorOps/invalid.mlir b/mlir/test/Dialect/VectorOps/invalid.mlir
index 4f56e94a092..892c10cd20d 100644
--- a/mlir/test/Dialect/VectorOps/invalid.mlir
+++ b/mlir/test/Dialect/VectorOps/invalid.mlir
@@ -704,3 +704,59 @@ func @constant_mask_with_zero_mask_dim_size() {
   %0 = vector.constant_mask [0, 2] : vector<4x3xi1>
   return
 }
+
+
+// -----
+
+func @extract_slices_non_unit_strides(%arg0 : vector<4x2xf32>) {
+  // expected-error@+1 {{requires unit strides}}
+  %0 = vector.extract_slices %arg0, [2, 2], [1, 3]
+    : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
+  return
+}
+
+// -----
+
+func @extract_slices_tuple_element_wrong_rank(%arg0 : vector<4x2xf32>) {
+  // expected-error@+1 {{requires vector tuple elements of rank 2}}
+  %0 = vector.extract_slices %arg0, [2, 2], [1, 1]
+    : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2x3xf32>>
+  return
+}
+
+// -----
+
+func @extract_slices_sizes_strides_wrong_rank(%arg0 : vector<4x2xf32>) {
+  // expected-error@+1 {{requires sizes and strides of rank}}
+  %0 = vector.extract_slices %arg0, [2, 2], [1, 1, 1]
+    : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
+  return
+}
+
+// -----
+
+func @extract_slices_invalid_tuple_element_type(%arg0 : vector<4x2xf32>) {
+  // expected-error@+1 {{invalid tuple element type}}
+  %0 = vector.extract_slices %arg0, [2, 2], [1, 1]
+    : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<4x2xf32>>
+  return
+}
+
+// -----
+
+func @tuple_of_non_vectors(%arg0 : vector<4x2xf32>) {
+  %c0 = constant 0 : index
+  // expected-error@+1 {{must be vector of any type values}}
+  %0 = vector.tuple %arg0, %c0 : vector<4x2xf32>, index
+  return
+}
+
+// -----
+
+func @tuple_get_of_non_vectors(%arg0 : tuple<vector<4x2xf32>, index>) {
+  // expected-error@+1 {{vector of any type values}}
+  %0 = vector.tuple_get %arg0, 0 : tuple<vector<4x2xf32>, index>
+  return
+}
+
+
diff --git a/mlir/test/Dialect/VectorOps/ops.mlir b/mlir/test/Dialect/VectorOps/ops.mlir
index a5bafb44366..1d28dea8282 100644
--- a/mlir/test/Dialect/VectorOps/ops.mlir
+++ b/mlir/test/Dialect/VectorOps/ops.mlir
@@ -159,3 +159,15 @@ func @constant_vector_mask() {
   %0 = vector.constant_mask [3, 2] : vector<4x3xi1>
   return
 }
+
+// CHECK-LABEL: extract_slices
+func @extract_slices(%arg0 : vector<4x2xf32>)
+  -> (tuple<vector<2x2xf32>, vector<2x2xf32>>) {
+  // CHECK: vector.extract_slices %{{.*}}, [2, 2], [1, 1] : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
+  %0 = vector.extract_slices %arg0, [2, 2], [1, 1]
+    : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
+  %1 = vector.tuple_get %0, 0 : tuple<vector<2x2xf32>, vector<2x2xf32>>
+  %2 = vector.tuple_get %0, 1 : tuple<vector<2x2xf32>, vector<2x2xf32>>
+  %3 = vector.tuple %1, %2 : vector<2x2xf32>, vector<2x2xf32>
+  return %3 : tuple<vector<2x2xf32>, vector<2x2xf32>>
+}