[Mlir-commits] [mlir] [MLIR][XeGPU] Enhance Layout Propagation for broadcasting both leading dimensions and inner unit dimensions (PR #185583)

[Jianhui Li]

================
@@ -143,20 +143,30 @@ xegpu::inferBroadcastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
                                   ArrayRef<int64_t> srcShape) {
 
   SmallVector<int64_t> bcastDims;
-  auto returnLayout = resLayout;
 
   // Handling broadcast from low-rank to high-rank (e.g., 1D to 2D) case.
-  int dimDiff = resShape.size() - srcShape.size();
+  size_t dimDiff = resShape.size() - srcShape.size();
+  // for (size_t i = 0; i < dimDiff; i++)
+  //   bcastDims.push_back(i);
 
-  if (dimDiff > 0) {
-    // Adding the missing leading dims
-    for (int i = 0; i < dimDiff; i++)
+  auto returnLayout = resLayout;
----------------
Jianhui-Li wrote:

done

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