Add VectorContractionOp to the VectorOps dialect.

PiperOrigin-RevId: 281605471

2 files changed, 98 insertions, 0 deletions

diff --git a/mlir/test/Dialect/VectorOps/invalid.mlir b/mlir/test/Dialect/VectorOps/invalid.mlir
index d8ebb9808b9..3ec8b4eca55 100644
--- a/mlir/test/Dialect/VectorOps/invalid.mlir
+++ b/mlir/test/Dialect/VectorOps/invalid.mlir
@@ -303,3 +303,71 @@ func @strided_slice(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{op expected result type to be 'vector<2x8x16xf32>'}}
   %1 = vector.strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4x8x16xf32> to vector<3x1xf32>
 }
+
+// -----
+
+func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>,
+                  %arg2: vector<8x15x5xf32>, %arg3 :  vector<8x15x8x5xf32>,
+                  %arg4 : index) {
+  // expected-error@+1 {{op expected at least one contracting dimension pair}}
+  %0 = vector.contract %arg0, %arg1, %arg2
+    { batch_dim_map = [[1, 0]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
+  return
+}
+
+// -----
+
+func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>,
+                  %arg2: vector<8x15x5xf32>, %arg3 :  vector<8x15x8x5xf32>,
+                  %arg4 : index) {
+  // expected-error@+1 {{invalid contracting dimension map}}
+  %0 = vector.contract %arg0, %arg1, %arg2
+    { batch_dim_map = [[1, 0]], contracting_dim_map = [[1, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
+  return
+}
+
+// -----
+
+func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>,
+                  %arg2: vector<8x15x5xf32>, %arg3 :  vector<8x15x8x5xf32>,
+                  %arg4 : index) {
+  // expected-error@+1 {{invalid batch dimension map}}
+  %0 = vector.contract %arg0, %arg1, %arg2
+    { batch_dim_map = [[1, 2]], contracting_dim_map = [[0, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
+
+  return
+}
+
+// -----
+
+func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>,
+                  %arg2: vector<8x15x88xf32>, %arg3 :  vector<8x15x8x5xf32>,
+                  %arg4 : index) {
+  // expected-error@+1 {{invalid accumulator/result vector shape}}
+  %0 = vector.contract %arg0, %arg1, %arg2
+    { batch_dim_map = [[1, 0]], contracting_dim_map = [[0, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x88xf32>
+
+  return
+}
+
+// -----
+
+func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>,
+                  %arg2: vector<8x15x5xf32>, %arg3 :  vector<8x15x8x5xf32>,
+                  %arg4 : index) {
+  %lhs_mask = vector.make_index_tuple %arg4, %arg4, %arg4, %arg4
+    : tuple<index, index, index, index>
+  %rhs_mask = vector.make_index_tuple %arg4, %arg4, %arg4, %arg4
+    : tuple<index, index, index, index>
+  // expected-error@+1 {{expected zero or exactly 2 vector mask operands}}
+  %0 = vector.contract %arg0, %arg1, %arg2, %lhs_mask
+    { batch_dim_map = [[1, 0]], contracting_dim_map = [[0, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
+  return
+}
+
+
diff --git a/mlir/test/Dialect/VectorOps/ops.mlir b/mlir/test/Dialect/VectorOps/ops.mlir
index 7ad46c28339..40b00291a05 100644
--- a/mlir/test/Dialect/VectorOps/ops.mlir
+++ b/mlir/test/Dialect/VectorOps/ops.mlir
@@ -48,3 +48,33 @@ func @strided_slice(%arg0: vector<4x8x16xf32>) -> vector<2x2x16xf32> {
   %1 = vector.strided_slice %arg0 {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x8x16xf32> to vector<2x2x16xf32>
   return %1: vector<2x2x16xf32>
 }
+
+// CHECK-LABEL: contraction
+func @contraction(%arg0 : vector<7x8x16x15xf32>, %arg1 : vector<8x16x7x5xf32>,
+                  %arg2 : vector<8x15x5xf32>, %arg3 : vector<8x15x8x5xf32>,
+                  %arg4 : index) {
+  // Test contraction with batch and contracting dims.
+  // CHECK: vector.contract {{.*}}, {{.*}}, {{.*}} {batch_dim_map = {{.*}}1, 0{{.*}}, contracting_dim_map = {{.*}}0, 2{{.*}}, {{.*}}2, 1{{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
+  %0 = vector.contract %arg0, %arg1, %arg2
+    { batch_dim_map = [[1, 0]], contracting_dim_map = [[0, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
+
+  // Test contraction with only contracting dims. In this case the lhs/rhs
+  // dimension of size 8 will be considered a free dim for lhs/rhs and will
+  // appear twice in the output.
+  // CHECK: vector.contract {{.*}}, {{.*}}, {{.*}} {contracting_dim_map = {{.*}}0, 2{{.*}}, {{.*}}2, 1{{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
+  %1 = vector.contract %arg0, %arg1, %arg3
+    { contracting_dim_map = [[0, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
+
+  // Test contraction with optional vector mask arguments.
+  %lhs_mask = vector.make_index_tuple %arg4, %arg4, %arg4, %arg4
+    : tuple<index, index, index, index>
+  %rhs_mask = vector.make_index_tuple %arg4, %arg4, %arg4, %arg4
+    : tuple<index, index, index, index>
+  // CHECK: vector.contract {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}} {contracting_dim_map = {{.*}}0, 2{{.*}}, {{.*}}2, 1{{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
+  %2 = vector.contract %arg0, %arg1, %arg3, %lhs_mask, %rhs_mask
+    { contracting_dim_map = [[0, 2], [2, 1]] }
+      : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
+  return
+}