[Mlir-commits] [mlir] [mlir][vector] Constrain broadcast->shape_cast folding (PR #190230)

[Andrzej Warzyński]

================
@@ -128,3 +128,20 @@ func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
   return %1: vector<1xf32>
 }
 
+// -----
+
+// https://github.com/llvm/llvm-project/issues/190614 — ShapeCastBroadcastFolder
+// must not rewrite shape_cast(broadcast) into a lower-rank broadcast when that
+// changes duplication (new leading dims) vs stretching (source unit dim).
+// For %arg0 = [[1], [2]]: broadcast to 2x2x1 then shape_cast to 2x2 yields
+// [[1, 2], [1, 2]]; folding to broadcast 2x1 -> 2x2 would incorrectly yield
+// [[1, 1], [2, 2]].
+// CHECK-LABEL: @shape_cast_broadcast_folder_issue_190614
+// CHECK:         vector.broadcast %{{.*}} : vector<2x1xf32> to vector<2x2x1xf32>
+// CHECK-NEXT:    vector.shape_cast %{{.*}} : vector<2x2x1xf32> to vector<2x2xf32>
+func.func @shape_cast_broadcast_folder_issue_190614(%arg0: vector<2x1xf32>) -> vector<2x2xf32> {
+  %0 = vector.broadcast %arg0 : vector<2x1xf32> to vector<2x2x1xf32>
+  %1 = vector.shape_cast %0 : vector<2x2x1xf32> to vector<2x2xf32>
+  return %1 : vector<2x2xf32>
+}
----------------
banach-space wrote:

This is effectively a negative test. Note that in our testing guide we suggest to prefix negative tests with e.g.  `negative` or `no_fold`: https://mlir.llvm.org/getting_started/TestingGuide/#test-naming-convention

Also, since you are already referring to the GitHub issue in the comment, no need to repeat that in the test function name.

```suggestion
func.func @@no_fold_bcast_mode_switch(%arg0: vector<2x1xf32>) -> vector<2x2xf32> {
  %0 = vector.broadcast %arg0 : vector<2x1xf32> to vector<2x2x1xf32>
  %1 = vector.shape_cast %0 : vector<2x2x1xf32> to vector<2x2xf32>
  return %1 : vector<2x2xf32>
}
```

I know that this can feel nit-picky, but consistent naming really helps navigating these tests.

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