[Mlir-commits] [mlir] [MLIR] [Vector] Added canonicalizer for folding from_elements + transpose (PR #161841)

[Andrzej Warzyński]

================
@@ -308,6 +308,41 @@ func.func @constant_mask_transpose_to_transposed_constant_mask() -> (vector<2x3x
 
 // -----
 
+// CHECK-LABEL: transpose_from_elements_1d
+func.func @transpose_from_elements_1d(%arg0: i32, %arg1: i32) -> vector<2xi32> {
+  %v = vector.from_elements %arg0, %arg1 : vector<2xi32>
+  %t = vector.transpose %v, [0] : vector<2xi32> to vector<2xi32>
+  return %t : vector<2xi32>
+  // CHECK: %[[R:.*]] = vector.from_elements %arg0, %arg1 : vector<2xi32>
+  // CHECK-NOT: vector.transpose
+}
+
+// CHECK-LABEL: transpose_from_elements_2d
+func.func @transpose_from_elements_2d(
+  %arg0: i32, %arg1: i32, %arg2: i32,
+  %arg3: i32, %arg4: i32, %arg5: i32
+) -> vector<3x2xi32> {
+  %arg6 = vector.from_elements %arg0, %arg1, %arg2, %arg3, %arg4, %arg5 : vector<2x3xi32>
+  %arg7 = vector.transpose %arg6, [1, 0] : vector<2x3xi32> to vector<3x2xi32>
+  return %arg7 : vector<3x2xi32>
+  // CHECK: %[[R:.*]] = vector.from_elements %arg0, %arg3, %arg1, %arg4, %arg2, %arg5 : vector<3x2xi32>
+  // CHECK-NOT: vector.transpose
+}
+
+// CHECK-LABEL: transpose_from_elements_3d
+func.func @transpose_from_elements_3d(
+  %arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32,
+  %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: i32, %arg11: i32
----------------
banach-space wrote:

Kind request - use variable names that make it easy to track the actual transposition. For example:
```suggestion
  %el_0_0_0: i32, %el_0_0_1: i32, %el_0_0_2: i32, %el_0_1_0: i32, ...
```

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