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

[Andrzej Warzyński]

================
@@ -0,0 +1,329 @@
+// 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 @to_from_elements_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 @to_from_elements_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>,
+                                                    %b: vector<8xf32>) -> vector<8xf32> {
+  %0:8 = vector.to_elements %a : vector<8xf32>
+  %1:8 = vector.to_elements %b : vector<8xf32>
+  %2 = vector.from_elements %0#7, %1#0, %0#6, %1#1, %0#5, %1#2, %0#4, %1#3 : vector<8xf32>
+  return %2 : vector<8xf32>
+}
+
+// CHECK-LABEL: func @from_elements_to_elements_single_shuffle(
+//  CHECK-SAME:     %[[A:.*]]: vector<8xf32>, %[[B:.*]]: vector<8xf32>
+//       CHECK:   %[[L0SH0:.*]] = vector.shuffle %[[A]], %[[B]] [7, 8, 6, 9, 5, 10, 4, 11] : vector<8xf32>
+//       CHECK:   return %[[L0SH0]]
+
+// -----
+
+func.func @to_from_elements_shuffle_tree_concat_4x8_to_32(%a: vector<8xf32>,
+                                                          %b: vector<8xf32>,
+                                                          %c: vector<8xf32>,
+                                                          %d: vector<8xf32>) -> vector<32xf32> {
+  %0:8 = vector.to_elements %a : vector<8xf32>
+  %1:8 = vector.to_elements %b : vector<8xf32>
+  %2:8 = vector.to_elements %c : vector<8xf32>
+  %3:8 = vector.to_elements %d : vector<8xf32>
+  %4 = vector.from_elements %0#0, %0#1, %0#2, %0#3, %0#4, %0#5, %0#6, %0#7,
+                            %1#0, %1#1, %1#2, %1#3, %1#4, %1#5, %1#6, %1#7,
+                            %2#0, %2#1, %2#2, %2#3, %2#4, %2#5, %2#6, %2#7,
+                            %3#0, %3#1, %3#2, %3#3, %3#4, %3#5, %3#6, %3#7 : vector<32xf32>
+  return %4 : vector<32xf32>
+}
+
+// CHECK-LABEL: func @to_from_elements_shuffle_tree_concat_4x8_to_32(
+//  CHECK-SAME:     %[[A:.*]]: vector<8xf32>, %[[B:.*]]: vector<8xf32>, %[[C:.*]]: vector<8xf32>, %[[D:.*]]: vector<8xf32>
+//       CHECK:   %[[L0SH0:.*]] = vector.shuffle %[[A]], %[[B]] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xf32>, vector<8xf32>
+//       CHECK:   %[[L0SH1:.*]] = vector.shuffle %[[C]], %[[D]] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15] : vector<8xf32>, vector<8xf32>
+//       CHECK:   %[[L1SH0:.*]] = vector.shuffle %[[L0SH0]], %[[L0SH1]] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
+//       CHECK:   return %[[L1SH0]] : vector<32xf32>
+
+// -----
+
+func.func @to_from_elements_shuffle_tree_concat_3x4_to_12(%a: vector<4xf32>,
+                                                          %b: vector<4xf32>,
+                                                          %c: vector<4xf32>) -> vector<12xf32> {
+  %0:4 = vector.to_elements %a : vector<4xf32>
+  %1:4 = vector.to_elements %b : vector<4xf32>
+  %2:4 = vector.to_elements %c : vector<4xf32>
+  %3 = vector.from_elements %0#0, %0#1, %0#2, %0#3, %1#0, %1#1, %1#2, %1#3, %2#0, %2#1, %2#2, %2#3 : vector<12xf32>
+  return %3 : vector<12xf32>
+}
+
+// CHECK-LABEL: func @to_from_elements_shuffle_tree_concat_3x4_to_12(
+//  CHECK-SAME:     %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[C:.*]]: vector<4xf32>
+//       CHECK:   %[[L0SH0:.*]] = vector.shuffle %[[A]], %[[B]] [0, 1, 2, 3, 4, 5, 6, 7] : vector<4xf32>, vector<4xf32>
+//       CHECK:   %[[L0SH1:.*]] = vector.shuffle %[[C]], %[[C]] [0, 1, 2, 3, -1, -1, -1, -1] : vector<4xf32>, vector<4xf32>
+//       CHECK:   %[[L1SH0:.*]] = vector.shuffle %[[L0SH0]], %[[L0SH1]] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] : vector<8xf32>, vector<8xf32>
+//       CHECK:   return %[[L1SH0]] : vector<12xf32>
+
+// -----
+
+func.func @to_from_elements_shuffle_tree_concat_64x4_256(
----------------
banach-space wrote:

IMHO, this example is a bit too long and I'm not sure whether it adds much unique coverage. Do we believe that jumping from e.g. 4 to 64 input vectors changes much?

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