[Mlir-commits] [mlir] [MLIR][XeGPU] Slice the new dim in broadcast properly (PR #206136)
Artem Kroviakov
llvmlistbot at llvm.org
Mon Jun 29 04:26:21 PDT 2026
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/206136
>From f4e421962313921c0ad9ff0785b36f0018fc64a7 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Fri, 26 Jun 2026 17:47:34 +0000
Subject: [PATCH 1/2] [MLIR][XeGPU] Slice the new dim in broadcast properly
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 13 ++++++---
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 27 +++++++++++++++++++
2 files changed, 37 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 61cd253508357..9eb7c710d2df7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -471,10 +471,17 @@ xegpu::inferBroadcastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
if (!bcastDims.empty())
bcastSourceLayout = bcastSourceLayout.setUnitDimData(bcastDims);
- if (dimDiff > 0) {
+ if (dimDiff) {
SmallVector<int64_t> sliceDims;
- for (size_t i = 0; i < dimDiff; i++)
- sliceDims.push_back(i);
+ bool isOuterDimDiffUnitDims =
+ llvm::all_of(llvm::seq<int64_t>(0, dimDiff),
+ [&](int64_t i) { return resShape[i] == 1; });
+ if (dimDiff && bcastDims.size() == dimDiff && isOuterDimDiffUnitDims) {
+ sliceDims.assign(bcastDims.begin(), bcastDims.end());
+ } else {
+ for (int64_t i = 0; i < dimDiff; ++i) // original behavior
+ sliceDims.push_back(i);
+ }
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
+}
+}
>From 0b028054067822b6d487cf18936513deee031b9d Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Mon, 29 Jun 2026 11:26:03 +0000
Subject: [PATCH 2/2] Add comments
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 19 ++++++++++++-------
1 file changed, 12 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 9eb7c710d2df7..678f585bdcca5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -461,26 +461,31 @@ 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);
+ // Case RankDiff:
if (dimDiff) {
SmallVector<int64_t> sliceDims;
- bool isOuterDimDiffUnitDims =
- llvm::all_of(llvm::seq<int64_t>(0, dimDiff),
- [&](int64_t i) { return resShape[i] == 1; });
+ 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 {
- for (int64_t i = 0; i < dimDiff; ++i) // original behavior
- sliceDims.push_back(i);
+ // 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,
More information about the Mlir-commits
mailing list