fixup! Fix invalid folding on mismatching broadcast dimensions

Author: mshockwave

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -2882,8 +2882,21 @@ struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
   }
 };
 
+// Return the broadcasted dimensions. Including broadcasts in the leading
+// dimensions and broadcasts through unit dimension (i.e. dim-1).
+static BitVector getBroadcastedDims(ArrayRef<int64_t> srcShape,
+                                    ArrayRef<int64_t> destShape) {
+  assert(destShape.size() >= srcShape.size());
+  BitVector broadcastedDims(destShape.size());
+  broadcastedDims.set(0, destShape.size() - srcShape.size());
+  auto unitDims = computeBroadcastedUnitDims(srcShape, destShape);
+  for (int64_t dim : unitDims)
+    broadcastedDims.set(dim);
+  return broadcastedDims;
+}
+
 // Fold broadcast(shape_cast(x)) into broadcast(x) if x's type is compatible
-// with broadcast's result type.
+// with broadcast's result type and the broadcasted dimensions are the same.
 struct FoldBroadcastOfShapeCast : public OpRewritePattern<BroadcastOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -2895,10 +2908,28 @@ struct FoldBroadcastOfShapeCast : public OpRewritePattern<BroadcastOp> {
 
     VectorType srcType = srcShapeCast.getSourceVectorType();
     VectorType destType = broadcastOp.getResultVectorType();
+    // Check type compatibility.
     if (vector::isBroadcastableTo(srcType, destType) !=
         BroadcastableToResult::Success)
       return failure();
 
+    // Given
+    // ```
+    // %s = shape_cast(%x)
+    // %b = broadcast(%s)
+    // ```
+    // If we want to fold %x into %b, the broadcasted dimensions from %x to
+    // %b has to be the same as that of from %s to %b.
+    ArrayRef<int64_t> shapecastShape =
+        srcShapeCast.getResultVectorType().getShape();
+    ArrayRef<int64_t> srcShape = srcType.getShape();
+    ArrayRef<int64_t> destShape = destType.getShape();
+    BitVector origBroadcastedDims =
+        getBroadcastedDims(shapecastShape, destShape);
+    BitVector newBroadcastedDims = getBroadcastedDims(srcShape, destShape);
+    if (newBroadcastedDims != origBroadcastedDims)
+      return failure();
+
     rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp, destType,
                                              srcShapeCast.getSource());
     return success();

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -1180,7 +1180,7 @@ func.func @canonicalize_shapecast_broadcast_to_broadcast(%arg0 : vector<2xf32>)
 // -----
 
 // CHECK-LABEL: func @negative_canonicalize_shapecast_broadcast_invalid_shape
-//       CHECK:   vector.shape_cast {{.+}} : vector<64xf32> to vector<4x16xf32
+//       CHECK:   vector.shape_cast {{.+}} : vector<64xf32> to vector<4x16xf32>
 //       CHECK:   vector.broadcast {{.+}} : vector<4x16xf32> to vector<2x4x16xf32>
 func.func @negative_canonicalize_shapecast_broadcast_invalid_shape(%arg0 : vector<64xf32>) -> vector<2x4x16xf32> {
   %0 = vector.shape_cast %arg0 : vector<64xf32> to vector<4x16xf32>
@@ -1190,6 +1190,17 @@ func.func @negative_canonicalize_shapecast_broadcast_invalid_shape(%arg0 : vecto
 
 // -----
 
+// CHECK-LABEL: func @negative_canonicalize_shapecast_broadcast_invalid_broadcasted_dims
+//       CHECK:   vector.shape_cast {{.+}} : vector<2x1xf32> to vector<1x2xf32>
+//       CHECK:   vector.broadcast {{.+}} : vector<1x2xf32> to vector<2x2xf32>
+func.func @negative_canonicalize_shapecast_broadcast_invalid_broadcasted_dims(%arg0 : vector<2x1xf32>) -> vector<2x2xf32> {
+  %0 = vector.shape_cast %arg0 : vector<2x1xf32> to vector<1x2xf32>
+  %1 = vector.broadcast %0 : vector<1x2xf32> to vector<2x2xf32>
+  return %1 : vector<2x2xf32>
+}
+
+// -----
+
 // CHECK-LABEL: fold_vector_transfer_masks
 func.func @fold_vector_transfer_masks(%A: memref<?x?xf32>) -> (vector<4x8xf32>, vector<4x[4]xf32>) {
   // CHECK: %[[C0:.+]] = arith.constant 0 : index