[Mlir-commits] [mlir] [MLIR][Vector] Fix ShapeCastBroadcastFolder when shape_cast changes non-unit dim strides (PR #227574)

[Federico Bruzzone]

================
@@ -144,3 +145,135 @@ func.func @no_fold_bcast_mode_switch(%arg0: vector<2x1xf32>) -> vector<2x2xf32>
   %1 = vector.shape_cast %0 : vector<2x2x1xf32> to vector<2x2xf32>
   return %1 : vector<2x2xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_axis_shift
+// CHECK:         vector.broadcast %{{.*}} : vector<1x4x1xf32> to vector<1x4x4xf32>
+// CHECK-NEXT:    vector.shape_cast %{{.*}} : vector<1x4x4xf32> to vector<4x4x1xf32>
+func.func @no_fold_bcast_axis_shift(%arg0: vector<1x4x1xf32>) -> vector<4x4x1xf32> {
+  %0 = vector.broadcast %arg0 : vector<1x4x1xf32> to vector<1x4x4xf32>
+  %1 = vector.shape_cast %0 : vector<1x4x4xf32> to vector<4x4x1xf32>
+  return %1 : vector<4x4x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_leading_dims
+// CHECK:         %[[RES:.*]] = vector.broadcast %{{.*}} : vector<3xf32> to vector<8x3xf32>
+// CHECK-NEXT:    return %[[RES]] : vector<8x3xf32>
+func.func @fold_bcast_leading_dims(%arg0: vector<3xf32>) -> vector<8x3xf32> {
+  %0 = vector.broadcast %arg0 : vector<3xf32> to vector<2x4x3xf32>
+  %1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
+  return %1 : vector<8x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_src_leading_unit_dims
+// CHECK:         %[[RES:.*]] = vector.broadcast %{{.*}} : vector<1x3xf32> to vector<4x3xf32>
+// CHECK-NEXT:    return %[[RES]] : vector<4x3xf32>
+func.func @fold_bcast_src_leading_unit_dims(%arg0: vector<1x3xf32>) -> vector<4x3xf32> {
+  %0 = vector.broadcast %arg0 : vector<1x3xf32> to vector<4x1x3xf32>
+  %1 = vector.shape_cast %0 : vector<4x1x3xf32> to vector<4x3xf32>
+  return %1 : vector<4x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_consecutive_unit_dims
+// CHECK:         %[[RES:.*]] = vector.broadcast %{{.*}} : vector<2x1x1x3xf32> to vector<2x2x2x3xf32>
+// CHECK-NEXT:    return %[[RES]] : vector<2x2x2x3xf32>
+func.func @fold_bcast_consecutive_unit_dims(%arg0: vector<2x1x1x3xf32>) -> vector<2x2x2x3xf32> {
+  %0 = vector.broadcast %arg0 : vector<2x1x1x3xf32> to vector<2x1x4x3xf32>
+  %1 = vector.shape_cast %0 : vector<2x1x4x3xf32> to vector<2x2x2x3xf32>
+  return %1 : vector<2x2x2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_scalable_vs_fixed
+// CHECK:         vector.broadcast %{{.*}} : vector<2x1x1xf32> to vector<2x[4]x2xf32>
+// CHECK-NEXT:    vector.shape_cast %{{.*}} : vector<2x[4]x2xf32> to vector<[1]x2x4x2xf32>
+func.func @no_fold_bcast_scalable_vs_fixed(%arg0: vector<2x1x1xf32>) -> vector<[1]x2x4x2xf32> {
+  %0 = vector.broadcast %arg0 : vector<2x1x1xf32> to vector<2x[4]x2xf32>
+  %1 = vector.shape_cast %0 : vector<2x[4]x2xf32> to vector<[1]x2x4x2xf32>
+  return %1 : vector<[1]x2x4x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @no_fold_bcast_scalable_axis_shift
+// CHECK:         vector.broadcast %{{.*}} : vector<1x[4]x1xf32> to vector<1x[4]x4xf32>
+// CHECK-NEXT:    vector.shape_cast %{{.*}} : vector<1x[4]x4xf32> to vector<4x[4]x1xf32>
+func.func @no_fold_bcast_scalable_axis_shift(%arg0: vector<1x[4]x1xf32>) -> vector<4x[4]x1xf32> {
+  %0 = vector.broadcast %arg0 : vector<1x[4]x1xf32> to vector<1x[4]x4xf32>
+  %1 = vector.shape_cast %0 : vector<1x[4]x4xf32> to vector<4x[4]x1xf32>
+  return %1 : vector<4x[4]x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_leading_scalable_dims
+// CHECK:         %[[RES:.*]] = vector.broadcast %{{.*}} : vector<[4]xf32> to vector<8x[4]xf32>
+// CHECK-NEXT:    return %[[RES]] : vector<8x[4]xf32>
+func.func @fold_bcast_leading_scalable_dims(%arg0: vector<[4]xf32>) -> vector<8x[4]xf32> {
+  %0 = vector.broadcast %arg0 : vector<[4]xf32> to vector<2x4x[4]xf32>
+  %1 = vector.shape_cast %0 : vector<2x4x[4]xf32> to vector<8x[4]xf32>
+  return %1 : vector<8x[4]xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @fold_bcast_leading_scalable_dim_fixed_src
+// CHECK:         %[[RES:.*]] = vector.broadcast %{{.*}} : vector<4xf32> to vector<[8]x4xf32>
+// CHECK-NEXT:    return %[[RES]] : vector<[8]x4xf32>
+func.func @fold_bcast_leading_scalable_dim_fixed_src(%arg0: vector<4xf32>) -> vector<[8]x4xf32> {
----------------
FedericoBruzzone wrote:

Ok, I literally don't know what the underlying GH implementation is for posting comments once one has put it on another line! 😤

Sorry for the noise, the line was: https://github.com/llvm/llvm-project/pull/227574/changes#diff-633e27bc0e3c45e45a44713cfb76f89163967fc4b2ae7debb4fa4e8cf6915c40R195-R202

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