[Mlir-commits] [mlir] [mlir][Vector] Improve support for vector.extract(broadcast) (PR #116234)

[Andrzej Warzyński]

================
@@ -652,24 +652,44 @@ func.func @fold_extract_transpose(
 
 // -----
 
-// CHECK-LABEL: fold_extract_broadcast
+// CHECK-LABEL: fold_extract_broadcast_same_type
 //  CHECK-SAME:   %[[A:.*]]: f32
 //       CHECK:   return %[[A]] : f32
-func.func @fold_extract_broadcast(%a : f32) -> f32 {
+func.func @fold_extract_broadcast_same_type(%a : f32, 
+                                            %idx0 : index, 
+                                            %idx1 : index) -> f32 {
   %b = vector.broadcast %a : f32 to vector<1x2x4xf32>
-  %r = vector.extract %b[0, 1, 2] : f32 from vector<1x2x4xf32>
+  // The indices don't batter for this folder, so we use mixed indices.
+  %r = vector.extract %b[%idx0, %idx1, 2] : f32 from vector<1x2x4xf32>
   return %r : f32
 }
 
 // -----
 
-// CHECK-LABEL: fold_extract_broadcast_0dvec
+// CHECK-LABEL: fold_extract_broadcast_same_type_vec
+//  CHECK-SAME:   %[[A:.*]]: vector<4xf32>
+//       CHECK:   return %[[A]] : vector<4xf32>
+func.func @fold_extract_broadcast_same_type_vec(%a : vector<4xf32>, 
+                                                %idx0 : index) 
+                                                -> vector<4xf32> {
----------------
banach-space wrote:

> I think that comes down to preference

Absolutely, we don't mandate a coding style for tests in LLVM. 

Still, we can opt in to reduce the cognitive load by being consistent within every test file. As pointed out in the [Coding Style](https://llvm.org/docs/CodingStandards.html#whitespace):

> As always, follow the [Golden Rule](https://llvm.org/docs/CodingStandards.html#golden-rule) above: follow the style of existing code if you are modifying and extending it.

I value **consistency** a lot and that's all I am asking for. While "Vector/canonicalize.mlir" doesn't follow a consistent style ([style 1](https://github.com/llvm/llvm-project/blob/b1a267e1b9e9b50ba5b99de014ed056bf201b762/mlir/test/Dialect/Vector/canonicalize.mlir#L71-L72), [style 2](https://github.com/llvm/llvm-project/blob/b1a267e1b9e9b50ba5b99de014ed056bf201b762/mlir/test/Dialect/Vector/canonicalize.mlir#L85), [style 3](https://github.com/llvm/llvm-project/blob/b1a267e1b9e9b50ba5b99de014ed056bf201b762/mlir/test/Dialect/Vector/canonicalize.mlir#L411-L412), [style 4](https://github.com/llvm/llvm-project/blob/b1a267e1b9e9b50ba5b99de014ed056bf201b762/mlir/test/Dialect/Vector/canonicalize.mlir#L1537-L1540), ...), lets try to find a way to improve this situation. Had the file followed a single style, this discussion likely wouldn’t be necessary.

Now, ATM this PR mixes 3 styles:

```mlir
// STYLE 1
func.func @fold_extract_broadcast_same_input_output_vec(%a : vector<4xf32>, 
                                                %idx0 : index) -> vector<4xf32> {
// STYLE 2
func.func @fold_extract_broadcast_dim1_broadcasting(%a : vector<2x1xf32>, 
                                                    %idx : index, 
                                                    %idx1 : index, 
                                                    %idx2 : index) -> f32 {
// STYLE 3
func.func @fold_extract_broadcast_to_equal_rank(%a : vector<1xf32>,
                                                         %idx0 : index) 
                                                         -> vector<8xf32> {
```

This change worsens consistency rather than improving it.

I encourage you to pick one style and apply it consistently throughout the file to enhance **readability** and **maintainability**. I am also asking that you try to choose one of the "more popular" styles within the file, e.g. https://github.com/llvm/llvm-project/blob/b1a267e1b9e9b50ba5b99de014ed056bf201b762/mlir/test/Dialect/Vector/canonicalize.mlir#L1568-L1572

>  I prefer having all the args close-by so i can read them easier.

I believe that you can achieve this within the parameters listed above ;-)

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