[VectorOps] Add [insert/extract]element definition together with lowering to LLVM

Similar to insert/extract vector instructions but
(1) work on 1-D vectors only
(2) allow for a dynamic index

  %c3 = constant 3 : index
  %0 = vector.insertelement %arg0, %arg1[%c : index] : vector<4xf32>
  %1 = vector.extractelement %arg0[%c3 : index] : vector<4xf32>

PiperOrigin-RevId: 285792205

2 files changed, 43 insertions, 1 deletions

diff --git a/mlir/test/Dialect/VectorOps/invalid.mlir b/mlir/test/Dialect/VectorOps/invalid.mlir
index 892c10cd20d..c04c8ea486a 100644
--- a/mlir/test/Dialect/VectorOps/invalid.mlir
+++ b/mlir/test/Dialect/VectorOps/invalid.mlir
@@ -60,12 +60,20 @@ func @shuffle_index_out_of_range(%arg0: vector<2xf32>, %arg1: vector<2xf32>) {
 // -----
 
 func @shuffle_empty_mask(%arg0: vector<2xf32>, %arg1: vector<2xf32>) {
-  // expected-error@+1 {{custom op 'vector.shuffle' invalid mask length}}
+  // expected-error@+1 {{'vector.shuffle' invalid mask length}}
   %1 = vector.shuffle %arg0, %arg1 [] : vector<2xf32>, vector<2xf32>
 }
 
 // -----
 
+func @extract_element(%arg0: vector<4x4xf32>) {
+  %c = constant 3 : index
+  // expected-error@+1 {{'vector.extractelement' op expected 1-D vector}}
+  %1 = vector.extractelement %arg0[%c : index] : vector<4x4xf32>
+}
+
+// -----
+
 func @extract_vector_type(%arg0: index) {
   // expected-error@+1 {{expected vector type}}
   %1 = vector.extract %arg0[] : index
@@ -115,6 +123,22 @@ func @extract_position_overflow(%arg0: vector<4x8x16xf32>) {
 
 // -----
 
+func @insert_element(%arg0: f32, %arg1: vector<4x4xf32>) {
+  %c = constant 3 : index
+  // expected-error@+1 {{'vector.insertelement' op expected 1-D vector}}
+  %0 = vector.insertelement %arg0, %arg1[%c : index] : vector<4x4xf32>
+}
+
+// -----
+
+func @insert_element_wrong_type(%arg0: i32, %arg1: vector<4xf32>) {
+  %c = constant 3 : index
+  // expected-error@+1 {{'vector.insertelement' op failed to verify that source operand and result have same element type}}
+  %0 = "vector.insertelement" (%arg0, %arg1, %c) : (i32, vector<4xf32>, index) -> (vector<4xf32>)
+}
+
+// -----
+
 func @insert_vector_type(%a: f32, %b: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected non-empty position attribute}}
   %1 = vector.insert %a, %b[] : f32 into vector<4x8x16xf32>
diff --git a/mlir/test/Dialect/VectorOps/ops.mlir b/mlir/test/Dialect/VectorOps/ops.mlir
index 1d28dea8282..69af80f46bc 100644
--- a/mlir/test/Dialect/VectorOps/ops.mlir
+++ b/mlir/test/Dialect/VectorOps/ops.mlir
@@ -53,6 +53,15 @@ func @shuffle2D(%a: vector<1x4xf32>, %b: vector<2x4xf32>) -> vector<3x4xf32> {
   return %1 : vector<3x4xf32>
 }
 
+// CHECK-LABEL: @extract_element
+func @extract_element(%a: vector<16xf32>) -> f32 {
+  // CHECK:      %[[C15:.*]] = constant 15 : index
+  %c = constant 15 : index
+  // CHECK-NEXT: vector.extractelement %{{.*}}[%[[C15]] : index] : vector<16xf32>
+  %1 = vector.extractelement %a[%c : index] : vector<16xf32>
+  return %1 : f32
+}
+
 // CHECK-LABEL: @extract
 func @extract(%arg0: vector<4x8x16xf32>) -> (vector<8x16xf32>, vector<16xf32>, f32) {
   // CHECK: vector.extract {{.*}}[3 : i32] : vector<4x8x16xf32>
@@ -64,6 +73,15 @@ func @extract(%arg0: vector<4x8x16xf32>) -> (vector<8x16xf32>, vector<16xf32>, f
   return %1, %2, %3 : vector<8x16xf32>, vector<16xf32>, f32
 }
 
+// CHECK-LABEL: @insert_element
+func @insert_element(%a: f32, %b: vector<16xf32>) -> vector<16xf32> {
+  // CHECK:      %[[C15:.*]] = constant 15 : index
+  %c = constant 15 : index
+  // CHECK-NEXT: vector.insertelement %{{.*}}, %{{.*}}[%[[C15]] : index] : vector<16xf32>
+  %1 = vector.insertelement %a, %b[%c : index] : vector<16xf32>
+  return %1 : vector<16xf32>
+}
+
 // CHECK-LABEL: @insert
 func @insert(%a: f32, %b: vector<16xf32>, %c: vector<8x16xf32>, %res: vector<4x8x16xf32>) -> vector<4x8x16xf32> {
   // CHECK: vector.insert %{{.*}}, %{{.*}}[3 : i32] : vector<8x16xf32> into vector<4x8x16xf32>