[Mlir-commits] [mlir] [MLIR][XeGPU] Extend propagation and sg_to_lane distribution pass support broadcast with low rank and scalar source input (PR #170409)

[Jianhui Li]

================
@@ -581,16 +581,37 @@ void LayoutInfoPropagation::visitVectorBroadCastOp(
   // Only consider vector to vector broadcasts for now.
   VectorType resultTy = broadcast.getResultVectorType();
   VectorType sourceTy = dyn_cast<VectorType>(broadcast.getSourceType());
-  if (!sourceTy) {
-    broadcast.emitWarning("Expecting source type to be a vector type.");
+  // skip layout propagation for non-vector source operand.
+  if (!sourceTy)
     return;
-  }
 
-  // Only consider nD -> nD broadcast.
+  // Hanlding broadcast from low-rank to high-rank (e.g., 1D to 2D) case.
   if (sourceTy.getRank() != resultTy.getRank()) {
-    broadcast.emitWarning("Expecting source and result to have same rank.");
+    auto sourceDims = sourceTy.getShape();
+    auto resultDims = resultTy.getShape();
+    // adding the missing leading missing dims
+    SmallVector<int64_t> bcastDims;
+    int64_t dimDiff = resultTy.getRank() - sourceTy.getRank();
+    for (int i = 0; i < dimDiff; i++) {
+      bcastDims.push_back(i);
+    }
+
+    // for the rest dims in the resultTy, if sourceTy dim is 1, then it's
+    // broadcasted dim
+    for (size_t i = 0; i < sourceDims.size(); i++) {
----------------
Jianhui-Li wrote:

yes. the source vector in the test case vector_broadcast_2d_to_2d_along_column have unit dim.  

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