address review comments, fix rebase conflicts WRT

Author: newling

mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp

@@ -29,7 +29,7 @@ using namespace mlir::vector;
 
 namespace {
 
-/// Convert a vector.broadcast without a scalar operand to a lower rank
+/// Convert a vector.broadcast with a vector operand to a lower rank
 /// vector.broadcast. vector.broadcast with a scalar operand is expected to be
 /// convertible to the lower level target dialect (LLVM, SPIR-V, etc.) directly.
 class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
@@ -45,12 +45,15 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
 
     // A broadcast from a scalar is considered to be in the lowered form.
     if (!srcType)
-      return failure();
+      return rewriter.notifyMatchFailure(
+          op, "broadcast from scalar already in lowered form");
 
     // Determine rank of source and destination.
     int64_t srcRank = srcType.getRank();
     int64_t dstRank = dstType.getRank();
 
+    // Here we are broadcasting to a rank-1 vector. Ensure that the source is a
+    // scalar.
     if (srcRank <= 1 && dstRank == 1) {
       SmallVector<int64_t> fullRankPosition(srcRank, 0);
       Value ext = vector::ExtractOp::create(rewriter, loc, op.getSource(),

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

@@ -1045,7 +1045,6 @@ struct ReorderElementwiseOpsOnBroadcast final
     Type resultElemType = resultType.getElementType();
 
     // Get the type of the first non-constant operand
-    // Operation *firstBroadcastOrSplat = nullptr;
     Value splatSource;
     for (Value operand : op->getOperands()) {
       Operation *definingOp = operand.getDefiningOp();
@@ -1057,26 +1056,24 @@ struct ReorderElementwiseOpsOnBroadcast final
       break;
     }
     if (!splatSource)
-      // TODO: why?
       return failure();
-
     Type unbroadcastResultType =
         cloneOrReplace(splatSource.getType(), resultElemType);
 
-    Type lhsBcastOrSplatType = splatSource.getType();
-
     // Make sure that all operands are broadcast from identically-shaped types:
     //  * scalar (`vector.broadcast` + `vector.splat`), or
     //  * vector (`vector.broadcast`).
     // Otherwise the re-ordering wouldn't be safe.
-    if (!llvm::all_of(op->getOperands(), [lhsBcastOrSplatType](Value val) {
+    if (!llvm::all_of(op->getOperands(), [splatSource](Value val) {
           if (auto source = getBroadcastLikeSource(val))
-            return source.getType() == lhsBcastOrSplatType;
+            return haveSameShapeAndScaling(source.getType(),
+                                           splatSource.getType());
           SplatElementsAttr splatConst;
           return matchPattern(val, m_Constant(&splatConst));
         })) {
       return rewriter.notifyMatchFailure(
-          op, "not all operands are broadcasts from the sametype");
+          op,
+          "not all operands are constants or broadcasts from the same type");
     }
 
     // Collect the source values before broadcasting