[Mlir-commits] [mlir] be82f85 - [MLIR][XeGPU] Slice the new dim in broadcast properly (#206136)

[llvmlistbot at llvm.org]

Author: Artem Kroviakov
Date: 2026-06-30T10:46:12+02:00
New Revision: be82f85c0cc80eb2e3b7ce612caf247212618708

URL: https://github.com/llvm/llvm-project/commit/be82f85c0cc80eb2e3b7ce612caf247212618708
DIFF: https://github.com/llvm/llvm-project/commit/be82f85c0cc80eb2e3b7ce612caf247212618708.diff

LOG: [MLIR][XeGPU] Slice the new dim in broadcast properly (#206136)

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
    mlir/test/Dialect/XeGPU/propagate-layout.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 61cd253508357..678f585bdcca5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -461,20 +461,32 @@ xegpu::inferBroadcastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
   SmallVector<int64_t> bcastDims;
   size_t dimDiff = resShape.size() - srcShape.size();
   auto bcastSourceLayout = resLayout;
+
+  // Right-aligned source in result, look for stretched unit dims.
   for (size_t i = dimDiff; i < resShape.size(); i++) {
     if ((srcShape[i - dimDiff] == 1) && (resShape[i] != 1))
       bcastDims.push_back(i);
   }
 
-  // the sg_layout and lane_layout for unit dimensions are preserved so it can
-  // be propagate to producer op so potentially used by the multi-reduction op.
+  // Case UnitDimStretch (e.g., 1x4 -> 4x4): the source layout data field must
+  // be 1.
   if (!bcastDims.empty())
     bcastSourceLayout = bcastSourceLayout.setUnitDimData(bcastDims);
 
-  if (dimDiff > 0) {
+  // Case RankDiff:
+  if (dimDiff) {
     SmallVector<int64_t> sliceDims;
-    for (size_t i = 0; i < dimDiff; i++)
-      sliceDims.push_back(i);
+    bool isOuterDimDiffUnitDims = llvm::all_of(
+        resShape.take_front(dimDiff), [&](int64_t dim) { return dim == 1; });
+    if (dimDiff && bcastDims.size() == dimDiff && isOuterDimDiffUnitDims) {
+      // Case RankDiffInnerDims (e.g., 1x4 -> 1x16x4):
+      //  slice the expanded inner dims
+      sliceDims.assign(bcastDims.begin(), bcastDims.end());
+    } else {
+      // Case RankDiffOuterDims (e.g., 1x4 -> 1x1x4):
+      //  slice the outer dims
+      llvm::append_range(sliceDims, llvm::seq<int64_t>(0, dimDiff));
+    }
     bcastSourceLayout = xegpu::SliceAttr::get(
         resLayout.getContext(), bcastSourceLayout,
         DenseI64ArrayAttr::get(resLayout.getContext(), sliceDims));

diff  --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 25d713ccc8a0f..1a741ba21e128 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -1100,3 +1100,30 @@ func.func @vector_deinterleave_f16(%arg0: memref<8x32xf16>, %arg1: memref<8x16xf
   return
 }
 }
+
+// -----
+// completeScatterLoadLaneLayoutFromInstData: user supplies only inst_data on a
+// scatter load; with no usable consumer lane info, the scatter default is used.
+// inst_data=[1,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @broadcast_slice_expanded_dim(
+// CHECK-DAG: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [8, 4], lane_data = [2, 1], order = [0, 1]>} dense<5.000000e-01> : vector<16x4xf4E2M1FN>
+// CHECK-DAG: %[[CST0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 8, 4], lane_data = [1, 1, 1], order = [1, 0, 2]>, dims = [1]>} dense<1.000000e+00> : vector<1x4xf8E8M0FNU>
+// CHECK-DAG: %[[BCAST:.*]] = vector.broadcast %[[CST0]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 8, 4], lane_data = [1, 2, 1], order = [1, 0, 2]>} : vector<1x4xf8E8M0FNU> to vector<1x16x4xf8E8M0FNU>
+// CHECK-DAG: %[[SCAST:.*]] = vector.shape_cast %[[BCAST]] {layout_result_0 = #xegpu.layout<lane_layout = [8, 4], lane_data = [2, 1], order = [0, 1]>} : vector<1x16x4xf8E8M0FNU> to vector<16x4xf8E8M0FNU>
+// CHECK-DAG: %[[EXT:.*]] = arith.scaling_extf %[[CST]], %[[SCAST]] {layout_result_0 = #xegpu.layout<lane_layout = [8, 4], lane_data = [2, 1], order = [0, 1]>} : vector<16x4xf4E2M1FN>, vector<16x4xf8E8M0FNU> to vector<16x4xbf16>
+// CHECK: xegpu.store_matrix %[[EXT]], %{{.*}}[0, 0] <{layout = #xegpu.layout<lane_layout = [8, 4], lane_data = [2, 1], order = [0, 1]>}>: vector<16x4xbf16>, !xegpu.mem_desc<16x4xbf16>
+func.func @broadcast_slice_expanded_dim(%dst: !xegpu.mem_desc<16x4xbf16>) {
+  %data = arith.constant dense<0.5> : vector<16x4xf4E2M1FN>
+  %loaded = arith.constant dense<1.> : vector<1x4xf8E8M0FNU>
+  %bcasted = vector.broadcast %loaded : vector<1x4xf8E8M0FNU> to vector<1x16x4xf8E8M0FNU>
+  %bcasted_shape_casted = vector.shape_cast %bcasted : vector<1x16x4xf8E8M0FNU> to vector<16x4xf8E8M0FNU>
+
+  %scaled = arith.scaling_extf %data, %bcasted_shape_casted {
+    layout_result_0 = #xegpu.layout<lane_layout = [8, 4], lane_data = [2, 1], order = [0, 1]>
+  } : vector<16x4xf4E2M1FN>, vector<16x4xf8E8M0FNU> to vector<16x4xbf16>
+
+  xegpu.store_matrix %scaled, %dst[0, 0] <{layout = #xegpu.layout<lane_layout = [8, 4], lane_data = [2, 1], order = [0, 1]>}> : vector<16x4xbf16>, !xegpu.mem_desc<16x4xbf16>
+  return
+}
+}