[Mlir-commits] [mlir] [mlir][xegpu] Fix shape_cast source layout inference for split dims (PR #226178)

[llvmlistbot at llvm.org]

https://github.com/Naaz30 created https://github.com/llvm/llvm-project/pull/226178

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

>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(