[Mlir-commits] [mlir] [mlir][Vector] Add `vector.shuffle` tree transformation (PR #145740)

[Andrzej Warzyński]

================
@@ -0,0 +1,423 @@
+// RUN: mlir-opt -lower-vector-to-from-elements-to-shuffle-tree -split-input-file %s | FileCheck %s
+
+// Captured variable names for `vector.shuffle` operations follow the L#SH# convention,
+// where L# refers to the level of the tree the shuffle belongs to, and SH# refers to
+// the shuffle index within that level.
+
+func.func @trivial_forwarding(%a: vector<8xf32>) -> vector<8xf32> {
+  %0:8 = vector.to_elements %a : vector<8xf32>
+  %1 = vector.from_elements %0#0, %0#1, %0#2, %0#3, %0#4, %0#5, %0#6, %0#7 : vector<8xf32>
+  return %1 : vector<8xf32>
+}
+
+// No shuffle tree needed for trivial forwarding case.
+
+// CHECK-LABEL: func @trivial_forwarding(
+//  CHECK-SAME:     %[[A:.*]]: vector<8xf32>
+//       CHECK:   return %[[A]] : vector<8xf32>
+
+// -----
+
+func.func @unsupported_multi_dim_vector_inputs(%a: vector<2x4xf32>, %b: vector<2x4xf32>) -> vector<4xf32> {
+  %0:8 = vector.to_elements %a : vector<2x4xf32>
+  %1:8 = vector.to_elements %b : vector<2x4xf32>
+  %2 = vector.from_elements %0#0, %0#7,
+                            %1#0, %1#7 : vector<4xf32>
+  return %2 : vector<4xf32>
+}
+
+// -----
+
+func.func @unsupported_multi_dim_vector_output(%a: vector<8xf32>, %b: vector<8xf32>) -> vector<2x2xf32> {
+  %0:8 = vector.to_elements %a : vector<8xf32>
+  %1:8 = vector.to_elements %b : vector<8xf32>
+  %2 = vector.from_elements %0#0, %0#7,
+                            %1#0, %1#7 : vector<2x2xf32>
+  return %2 : vector<2x2xf32>
+}
+
+// -----
+
+func.func @single_input_shuffle(%a: vector<8xf32>) -> vector<8xf32> {
+  %0:8 = vector.to_elements %a : vector<8xf32>
+  %1 = vector.from_elements %0#7, %0#0, %0#6, %0#1, %0#5, %0#2, %0#4, %0#3 : vector<8xf32>
+  return %1 : vector<8xf32>
+}
+
+// CHECK-LABEL: func @single_input_shuffle(
+//  CHECK-SAME:     %[[A:.*]]: vector<8xf32>
+      // CHECK:   %[[L0SH0:.*]] = vector.shuffle %[[A]], %[[A]] [7, 0, 6, 1, 5, 2, 4, 3] : vector<8xf32>, vector<8xf32>
+      // CHECK:   return %[[L0SH0]]
+
+// -----
+
+func.func @from_elements_to_elements_single_shuffle(%a: vector<8xf32>,
----------------
banach-space wrote:

Looking at the tests that follow, I would update these names for consistency  - at least in my head it feels _more_ consistent :)

```suggestion
func.func @shuffle_no_tree_single_input(%a: vector<8xf32>) -> vector<8xf32> {
  %0:8 = vector.to_elements %a : vector<8xf32>
  %1 = vector.from_elements %0#7, %0#0, %0#6, %0#1, %0#5, %0#2, %0#4, %0#3 : vector<8xf32>
  return %1 : vector<8xf32>
}

// CHECK-LABEL: func @single_input_shuffle(
//  CHECK-SAME:     %[[A:.*]]: vector<8xf32>
      // CHECK:   %[[L0SH0:.*]] = vector.shuffle %[[A]], %[[A]] [7, 0, 6, 1, 5, 2, 4, 3] : vector<8xf32>, vector<8xf32>
      // CHECK:   return %[[L0SH0]]

// -----

func.func @shuffle_no_tree_multiple_inputs(%a: vector<8xf32>,
```

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