[Mlir-commits] [mlir] [mlir][xegpu] Fix shape_cast source layout inference for split dims (PR #226178)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 24 07:17:31 PDT 2026
https://github.com/Naaz30 updated https://github.com/llvm/llvm-project/pull/226178
>From 4ed85e695daf317be1d65cb51d7c7f89c25dde44 Mon Sep 17 00:00:00 2001
From: Naazni Yahya <naazniyahya at Mac.lan>
Date: Thu, 24 Sep 2026 09:02:43 -0500
Subject: [PATCH] [mlir][xegpu] Fix shape_cast source layout inference for
split dims
inferShapeCastSourceLayout collapsed sg_layout and sg_data independently
when a 1-D source was split into several result dims. That is only valid
when each subgroup's elements stay a single round-robin pattern after
row-major flattening. For sg_layout=[1,2], sg_data=[2,2] on 2x4 it
inferred sg_data=[4] instead of [2], so the source and result disagreed
on which elements each subgroup owns.
Fixes #224655
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 79 ++++++++++++++++++-
.../XeGPU/propagate-layout-subgroup.mlir | 33 ++++++++
2 files changed, 110 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index e7c945f535044..3a58eda63cb80 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -835,6 +835,37 @@ xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return resLayout;
}
+/// Folds a group of adjacent result dims (listed outer to inner) into the one
+/// source dim they came from in a row-major vector.shape_cast. Returns the
+/// source {sg_layout, sg_data}, or std::nullopt when a subgroup's elements in
+/// the flattened space can't be described by a single round-robin pair.
+static std::optional<std::pair<int64_t, int64_t>>
+foldSplitGroup(ArrayRef<int64_t> layout, ArrayRef<int64_t> data,
+ ArrayRef<int64_t> shape,
+ function_ref<bool(size_t)> innerIsFaster) {
+ int64_t L = layout.back(), D = data.back(), S = shape.back();
+ for (int64_t i = static_cast<int64_t>(layout.size()) - 2; i >= 0; --i) {
+ int64_t Lo = layout[i], Do = data[i], So = shape[i];
+ if (L == 1) {
+ // (a) One subgroup owns every inner offset, so each owned outer index is
+ // a contiguous run of S elements.
+ L = Lo;
+ D = Do * S;
+ } else if (Lo == 1) {
+ // (b) Outer dim not distributed, every row repeats the inner pattern.
+ if (S % (L * D) != 0)
+ return std::nullopt;
+ } else if (Do == 1 && L * D == S && innerIsFaster(i)) {
+ // (c) One outer index per pass, inner tiled exactly once.
+ L = Lo * L;
+ } else {
+ return std::nullopt;
+ }
+ S *= So;
+ }
+ return std::make_pair(L, D);
+}
+
/// Infers the source layout attribute for a shape cast operation given the
/// result layout attribute, result shape, and source shape.
xegpu::DistributeLayoutAttr
@@ -865,13 +896,57 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return srcLayout;
}
- // Use case 2: Dim split from source to result, for multi-stage reduction
+ // Use case 2: Dim split from source to result, for multi-stage reduction.
SmallVector<SmallVector<int64_t>> splitDimGroups;
if (xegpu::matchSplitDimExpansion(srcShape, resShape, splitDimGroups)) {
+ SmallVector<int64_t> sgLayout = resLayout.getEffectiveSgLayoutAsInt();
+ SmallVector<int64_t> sgData = resLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> srcSgData;
+
+ if (!sgLayout.empty()) {
+ SmallVector<int64_t> order = resLayout.getEffectiveOrderAsInt();
+ SmallVector<int64_t> speed(order.size());
+ for (auto [pos, dim] : llvm::enumerate(order))
+ speed[dim] = pos;
+
+ for (const auto &dimGroup : splitDimGroups) {
+ SmallVector<int64_t> l, d, s;
+ for (int64_t dim : dimGroup) {
+ l.push_back(sgLayout[dim]);
+ d.push_back(sgData[dim]);
+ s.push_back(resShape[dim]);
+ }
+ auto innerIsFaster = [&](size_t i) {
+ for (size_t j = i + 1; j < dimGroup.size(); ++j)
+ if (speed[dimGroup[j]] > speed[dimGroup[i]])
+ return false;
+ return true;
+ };
+ auto folded = foldSplitGroup(l, d, s, innerIsFaster);
+ if (!folded)
+ return nullptr; // not expressible; caller emits a warning
+ srcSgData.push_back(folded->second);
+ }
+ }
+
+ // Collapse later groups first so earlier groups' dim indices stay valid.
auto srcLayout = resLayout;
- for (const auto &dimGroup : splitDimGroups)
+ for (const auto &dimGroup : llvm::reverse(splitDimGroups))
srcLayout = srcLayout.collapseDims(dimGroup);
+ // collapseDims multiplies sg_data; overwrite it with the folded values.
+ if (!srcSgData.empty() &&
+ srcSgData != srcLayout.getEffectiveSgDataAsInt()) {
+ auto plain = dyn_cast<xegpu::LayoutAttr>(srcLayout);
+ if (!plain)
+ return nullptr;
+ MLIRContext *ctx = resLayout.getContext();
+ SmallVector<int32_t> sgData32(srcSgData.begin(), srcSgData.end());
+ srcLayout = xegpu::LayoutAttr::get(
+ ctx, plain.getSgLayout(), DenseI32ArrayAttr::get(ctx, sgData32),
+ plain.getInstData(), plain.getLaneLayout(), plain.getLaneData(),
+ plain.getOrder());
+ }
return srcLayout;
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 3cf4373e1ef0d..a197f0867014d 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -577,6 +577,39 @@ gpu.module @test {
}
}
+// -----
+// shape_cast split where each subgroup's tile spans both rows, the source must
+// be round-robin [2]/[2], not contiguous [2]/[4].
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_split_outer_not_distributed(
+// CHECK: arith.constant {{.*}}layout_result_0 = #xegpu.layout<sg_layout = [2], sg_data = [2]
+ gpu.func @shape_cast_split_outer_not_distributed(%dst: memref<2x4xf32>) kernel {
+ %src = arith.constant dense<0.0> : vector<8xf32>
+ %result = vector.shape_cast %src : vector<8xf32> to vector<2x4xf32>
+ %td = xegpu.create_nd_tdesc %dst : memref<2x4xf32> -> !xegpu.tensor_desc<2x4xf32>
+ xegpu.store_nd %result, %td[0, 0]
+ <{layout = #xegpu.layout<sg_layout = [1, 2], sg_data = [2, 2], order = [0, 1]>}>
+ : vector<2x4xf32>, !xegpu.tensor_desc<2x4xf32>
+ gpu.return
+ }
+}
+
+// -----
+// Case (c): one row per subgroup, columns split once [4]/[2].
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_split_one_row_per_sg(
+// CHECK: arith.constant {{.*}}layout_result_0 = #xegpu.layout<sg_layout = [4], sg_data = [2]
+ gpu.func @shape_cast_split_one_row_per_sg(%dst: memref<2x4xf32>) kernel {
+ %src = arith.constant dense<0.0> : vector<8xf32>
+ %result = vector.shape_cast %src : vector<8xf32> to vector<2x4xf32>
+ %td = xegpu.create_nd_tdesc %dst : memref<2x4xf32> -> !xegpu.tensor_desc<2x4xf32>
+ xegpu.store_nd %result, %td[0, 0]
+ <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [1, 2]>}>
+ : vector<2x4xf32>, !xegpu.tensor_desc<2x4xf32>
+ gpu.return
+ }
+}
+
// -----
gpu.module @test {
// CHECK-LABEL: gpu.func @shape_cast_collapse_replicated(
More information about the Mlir-commits
mailing list