fixup! Rewrite as a folding pattern

Author: mshockwave

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

@@ -2841,9 +2841,59 @@ LogicalResult BroadcastOp::verify() {
   llvm_unreachable("unexpected vector.broadcast op error");
 }
 
+// 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 and the broadcasted dimensions are the same.
+static LogicalResult foldBroadcastOfShapeCast(BroadcastOp broadcastOp) {
+  auto srcShapeCast = broadcastOp.getSource().getDefiningOp<ShapeCastOp>();
+  if (!srcShapeCast)
+    return failure();
+
+  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();
+
+  broadcastOp.getSourceMutable().assign(srcShapeCast.getSource());
+  return success();
+}
+
 OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
   if (getSourceType() == getResultVectorType())
     return getSource();
+  if (succeeded(foldBroadcastOfShapeCast(*this)))
+    return getResult();
+
   if (!adaptor.getSource())
     return {};
   auto vectorType = getResultVectorType();
@@ -2881,67 +2931,13 @@ struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
     return success();
   }
 };
-
-// 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 and the broadcasted dimensions are the same.
-struct FoldBroadcastOfShapeCast : public OpRewritePattern<BroadcastOp> {
-  using OpRewritePattern::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(BroadcastOp broadcastOp,
-                                PatternRewriter &rewriter) const override {
-    auto srcShapeCast = broadcastOp.getSource().getDefiningOp<ShapeCastOp>();
-    if (!srcShapeCast)
-      return failure();
-
-    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();
-  }
-};
 } // namespace
 
 void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
   // BroadcastToShapeCast is not a default canonicalization, it is opt-in by
   // calling `populateCastAwayVectorLeadingOneDimPatterns`
-  results.add<BroadcastFolder, FoldBroadcastOfShapeCast>(context);
+  results.add<BroadcastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//