[Mlir-commits] [mlir] [mlir][vector] support unrolling vector.contract with scalar result (PR #232203)

[Federico Bruzzone]

================
@@ -935,3 +935,66 @@ func.func @masked_vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf3
 // CHECK-LABEL: func @masked_vector_multi_reduction
 //  CHECK-NEXT:   %[[R:.*]] = vector.mask %{{.*}} { vector.multi_reduction <add>, {{.*}} : vector<4x6xf32> to vector<4xf32> }
 //  CHECK-NEXT:   return %[[R]]
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0) -> ()>
+module {
+  func.func @vector_contract_to_scalar_f32_from_1d(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
+    %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
+    return %r : f32
+  }
+}
+
+// CHECK-LABEL: func @vector_contract_to_scalar_f32_from_1d(
+// CHECK-SAME:    %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[ACC:.*]]: f32) -> f32 {
+// CHECK:         %[[A0:.*]] = vector.extract_strided_slice %[[A]] offsets = [0], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK:         %[[B0:.*]] = vector.extract_strided_slice %[[B]] offsets = [0], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK:         %[[R0:.*]] = vector.contract {{.*}} %[[A0]], %[[B0]], %[[ACC]] : vector<2xf32>, vector<2xf32> into f32
+// CHECK:         %[[A1:.*]] = vector.extract_strided_slice %[[A]] offsets = [2], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK:         %[[B1:.*]] = vector.extract_strided_slice %[[B]] offsets = [2], sizes = [2], strides = [1] : vector<4xf32> to vector<2xf32>
+// CHECK:         %[[R1:.*]] = vector.contract {{.*}} %[[A1]], %[[B1]], %[[R0]] : vector<2xf32>, vector<2xf32> into f32
+// CHECK:         return %[[R1]] : f32
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0) -> ()>
+module {
----------------
FedericoBruzzone wrote:

Can we drop the `module` wrapper?

https://github.com/llvm/llvm-project/pull/232203