[Mlir-commits] [mlir] [mlir][Vector] Add `vector.shuffle` fold for poison inputs (PR #125608)

[Andrzej Warzyński]

================
@@ -2023,6 +2023,45 @@ func.func @shuffle_1d_poison_idx() -> vector<4xi32> {
 
 // -----
 
+// CHECK-LABEL: func @shuffle_1d_rhs_lhs_poison
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   %[[V:.+]] = ub.poison : vector<4xi32>
+//       CHECK:   return %[[V]]
+func.func @shuffle_1d_rhs_lhs_poison() -> vector<4xi32> {
+  %v0 = ub.poison : vector<3xi32>
+  %v1 = ub.poison : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @shuffle_1d_lhs_poison
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   %[[V:.+]] = arith.constant dense<[5, 4, 5, 5]> : vector<4xi32>
+//       CHECK:   return %[[V]]
+func.func @shuffle_1d_lhs_poison() -> vector<4xi32> {
+  %v0 = arith.constant dense<[5, 4, 3]> : vector<3xi32>
+  %v1 = ub.poison : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>
----------------
banach-space wrote:

[nit] There's a value _and_ index == 5, so it's not obvious that the first element of `%v0` is in any way significant. Perhaps use some more distinct number? (e.g. `123`).
```suggestion
  %v0 = arith.constant dense<[123, 4, 3]> : vector<3xi32>
  %v1 = ub.poison : vector<3xi32>
  %shuffle = vector.shuffle %v0, %v1 [3, 1, 123, 4] : vector<3xi32>, vector<3xi32>
```

I appreciate that this is obvious _right now_, but lets also cater for our future selves :) 

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