[Mlir-commits] [mlir] [mlir][xegpu] Raise the reduction vector size to match maxReduceVectorSize (PR #225122)

[Jianhui Li]

https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/225122

>From df8ac1580e107fca33d595a880f1a0f9eb0c2313 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 19 Sep 2026 22:24:09 +0000
Subject: [PATCH 1/6] [mlir][xegpu] Raise maxReduceVectorSize and derive
 reduction lane layouts

Bump maxReduceVectorSize from 1 to 16 so a multi_reduction's source layout can
pack several elements per lane instead of one, and derive that layout from a
single rule instead of two competing policies.

The previous code chose between a consumer-driven path that spent spare lanes on
the reduction dim and a fallback that kept the reduction inside a lane. Which
one ran depended on the op's shape at the time, so XeGPUPeepHoleOptimizer
splitting a 2D reduction changed the answer between propagation runs and left an
unlowerable convert_layout behind. computeReductionLaneLayoutAndData now applies
one rule:

  * lane_layout: lanes live only in the innermost two dims. Surviving dims take
    the consumer's lane_layout; leftover lanes go to the innermost dim when it
    is reduced. Lanes are not claimed further out, because that is where a
    producer spreads its own and a redistribution has no lowering. A consumer
    already spending the whole subgroup therefore leaves the reduction dim
    lane-less and the reduction lane-local, so [16, 1] falls out of the rule
    rather than from a separate policy.

  * lane_data: maxReduceVectorSize is a per-lane budget shared across dims. The
    consumer's lane_data counts against it and only the remainder is placed, on
    a lane-less unpacked dim, preferring the one just outside the innermost two
    -- the scale computation shape, where those two come from a shape_cast that
    split a single dim.

  * A consumer sliced over exactly the reduction dims already describes every
    source dim, so its grid and packing are reused directly.

inst_data is now lane_layout * lane_data on every dim, with no exception; the
previous fold of the reduce-vector size into inst_data on an arbitrary
lane-less dim is gone.

Also clamp lane_data in setupGenericLoadAnchorLayout, shared by load_gather and
load_matrix. Those read one contiguous span per lane, so a lane can only pack
along the innermost dim and only as far as the op's chunk size allows. A
consumer requesting lane_data > 1 on an outer, strided dim is not lowerable:
load_matrix becomes a single contiguous llvm.load, which would fetch adjacent
inner elements instead of the intended strided ones and silently drop rows in a
cross-subgroup reduction SLM round-trip.

Also initialize `intra` in XeGPUPeepHoleOptimizer's getReductionDimOrder. Only
`cross` was initialized, so when every reduction dim carries lanes the loop
never assigns `intra` and the caller's `== -1` guard cannot fire, indexing out
of bounds in accShape.erase.

Fixes seven imex integration failures (flash_attention_fwd and its two
array_len variants, four reduction cases), all of which failed to legalize
convert_layout and multi_reduction. All XeGPU and VectorToXeGPU tests pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 213 +++++++++++++-----
 .../Transforms/XeGPUPeepHoleOptimizer.cpp     |   2 +-
 .../XeGPU/propagate-layout-inst-data.mlir     |  53 +++--
 mlir/test/Dialect/XeGPU/propagate-layout.mlir |  26 +--
 4 files changed, 202 insertions(+), 92 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 05b3f24d660b4e..85a4fcdd69168d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1269,42 +1269,136 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
   return {laneLayout, laneData, order};
 }
 
-/// Computes the (lane_layout, lane_data) for a multi-reduction's source layout.
-/// Only the innermost two dims are distributed; leading dims are assumed unit.
-/// `subgroupSize` lanes go on one dim; up to `maxReduceVectorSize` elements are
-/// packed into lane_data on the other. To minimize cross-lane reduction, lanes
-/// are spread across a non-reduction dim when possible so the reduction happens
-/// within a lane. inst_data is the element-wise product lane_layout *
-/// lane_data.
+/// Computes the (lane_layout, lane_data, inst_data) for a multi-reduction's
+/// source layout, with inst_data = lane_layout * lane_data on every dim.
 ///
-/// e.g. with srcShape=[32, 128], subgroupSize=16, maxReduceVectorSize=2:
-///   - Switch: reductionDims=[1] and consumerReductionDims=[] -> lanes move
-///     to the non-reduction dim 0: lane_layout=[16, 1], lane_data=[1, 2].
-///   - Default: reductionDims=[0, 1] (both reduced) -> lanes stay on the
-///     innermost dim: lane_layout=[1, 16], lane_data=[2, 1].
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
-                                  ArrayRef<int64_t> reductionDims,
-                                  int subgroupSize, int64_t maxReduceVectorSize,
-                                  bool verticalLaneLayout = false) {
+/// Lanes live only in the innermost two dims. Surviving dims take the
+/// consumer's lane_layout; leftover lanes go to the innermost dim when it is
+/// reduced, so the reduction stays inside the subgroup. A consumer already
+/// spending the whole subgroup thus leaves the reduction dim lane-less and the
+/// reduction lane-local, which is where [16, 1] comes from.
+///
+/// maxReduceVectorSize is a per-lane budget shared across dims: the consumer's
+/// lane_data counts against it and only the remainder is placed here.
+///
+/// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
+/// lane_layout=[1,2,8], lane_data=[16,1,1], inst_data=[16,2,8].
+/// e.g. src=[16,16] reduce [1], consumer lane=[16] ->
+/// lane_layout=[16,1], lane_data=[1,16], inst_data=[16,16].
+static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+                  SmallVector<int64_t>>
+computeReductionLaneLayoutAndData(
+    ArrayRef<int64_t> srcShape, ArrayRef<int64_t> reductionDims,
+    int subgroupSize, int64_t maxReduceVectorSize,
+    xegpu::DistributeLayoutAttr consumerLayout = nullptr) {
   int srcRank = srcShape.size();
   SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
+  auto isReduced = [&](int dim) {
+    return llvm::is_contained(reductionDims, dim);
+  };
 
-  int innermost = srcRank - 1;
-  int secondInnermost = srcRank - 2;
+  // A consumer sliced over exactly the reduction dims already describes every
+  // source dim, so reuse its grid and packing instead of re-deriving something
+  // that could disagree with the IR. The reduce-vector size is raised on top.
+  auto consumerSlice = dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+  if (consumerSlice &&
+      consumerSlice.getDims().asArrayRef().equals(reductionDims)) {
+    auto parent = consumerSlice.getParent();
+    SmallVector<int64_t> parentLaneLayout =
+        parent.getEffectiveLaneLayoutAsInt();
+    SmallVector<int64_t> parentLaneData = parent.getEffectiveLaneDataAsInt();
+    // Only trust a parent that spends the whole subgroup. One that does not
+    // distributes less than the producer already does.
+    int64_t parentLanes = 1;
+    for (int64_t l : parentLaneLayout)
+      parentLanes *= l;
+    if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
+        static_cast<int>(parentLaneData.size()) == srcRank &&
+        parentLanes == subgroupSize) {
+      // Raise, never lower, what the parent asks for.
+      for (int i = 0; i < srcRank; ++i)
+        if (parentLaneLayout[i] == 1 && srcShape[i] > 1) {
+          parentLaneData[i] = std::max(
+              parentLaneData[i], std::min(maxReduceVectorSize, srcShape[i]));
+          break;
+        }
+      SmallVector<int64_t> parentInstData(srcRank);
+      for (int i = 0; i < srcRank; ++i)
+        parentInstData[i] = parentLaneLayout[i] * parentLaneData[i];
+      return {parentLaneLayout, parentLaneData, parentInstData};
+    }
+  }
 
-  if (verticalLaneLayout && secondInnermost >= 0) {
-    std::swap(innermost, secondInnermost);
+  // Otherwise the consumer covers only the surviving dims, at reduced rank.
+  SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
+  if (consumerLayout) {
+    consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+    consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
   }
-  int laneDim = innermost;
-  int vectorDim = secondInnermost; // negative for rank 1
 
-  laneLayout[laneDim] =
-      std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
+  // Lanes are confined to the innermost two dims.
+  int firstLaneDim = std::max(0, srcRank - 2);
+
+  // Surviving dims inherit the consumer's grid where lanes may live, and its
+  // packing as the base. Lanes it wants on a leading dim cannot be honored;
+  // convert_layout reconciles that.
+  int64_t lanesUsed = 1;
+  for (int i = 0, c = 0; i < srcRank; ++i) {
+    if (isReduced(i))
+      continue;
+    int64_t consumerLanes = c < static_cast<int>(consumerLaneLayout.size())
+                                ? consumerLaneLayout[c]
+                                : 1;
+    if (c < static_cast<int>(consumerLaneData.size()))
+      laneData[i] = consumerLaneData[c];
+    ++c;
+    if (i < firstLaneDim)
+      continue;
+    laneLayout[i] = std::min<int64_t>(consumerLanes, srcShape[i]);
+    lanesUsed *= laneLayout[i];
+  }
+
+  // Leftover lanes reduce within the subgroup, and only on the innermost dim:
+  // that is where producers spread theirs, so claiming them further out needs a
+  // redistribution convert_layout cannot lower.
+  int64_t spareLanes = subgroupSize / lanesUsed;
+  int innermost = srcRank - 1;
+  if (spareLanes > 1 && isReduced(innermost) && srcShape[innermost] > 1)
+    laneLayout[innermost] = std::min(spareLanes, srcShape[innermost]);
+
+  // Spend what is left of the budget on a lane-less, unpacked dim, since a lane
+  // owns such a dim whole. Prefer the dim just outside the innermost two: in
+  // the scale computation those two come from a shape_cast that split one dim,
+  // so the dim above them is what should carry the per-lane vector.
+  int64_t packedByConsumer = 1;
+  for (int64_t d : laneData)
+    packedByConsumer *= d;
+  int64_t vectorBudget =
+      maxReduceVectorSize / std::min(packedByConsumer, maxReduceVectorSize);
+
+  int vectorDim = -1;
+  int outsideLaneDims = srcRank - 3;
+  auto usable = [&](int i) {
+    return laneLayout[i] == 1 && laneData[i] == 1 && srcShape[i] > 1;
+  };
+  if (vectorBudget > 1) {
+    if (outsideLaneDims >= 0 && usable(outsideLaneDims))
+      vectorDim = outsideLaneDims;
+    else
+      for (int i = srcRank - 1; i >= 0; --i)
+        if (usable(i)) {
+          vectorDim = i;
+          break;
+        }
+  }
   if (vectorDim >= 0)
-    laneData[vectorDim] = std::min(maxReduceVectorSize, srcShape[vectorDim]);
+    laneData[vectorDim] = std::min(vectorBudget, srcShape[vectorDim]);
 
-  return {laneLayout, laneData};
+  SmallVector<int64_t> instData(srcRank);
+  for (int i = 0; i < srcRank; ++i)
+    instData[i] = laneLayout[i] * laneData[i];
+
+  return {laneLayout, laneData, instData};
 }
 
 //===----------------------------------------------------------------------===//
@@ -1767,7 +1861,8 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
 /// consumer layout (from its result's downstream uses) and validates it
 /// against uArch constraints; if valid, the consumer's `inst_data` /
 /// `sg_layout` are honored. Otherwise the helper falls back to defaults
-/// derived from uArch block parameters.
+/// derived from uArch block parameters. `lane_data` never comes from the
+/// consumer; see compute2DBlockIOLaneLayout.
 xegpu::DistributeLayoutAttr
 xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
                                VectorType resVecTy,
@@ -1925,6 +2020,20 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
   laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
   laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
 
+  // A matrix/gather load reads a single contiguous span of memory per lane, so
+  // a lane can only pack (lane_data > 1) elements along the innermost,
+  // contiguous dim. A consumer that requests lane_data > 1 on an outer, strided
+  // dim (e.g. a downstream reduction packing across mem_desc rows) is not
+  // lowerable: the contiguous load would fetch adjacent inner elements instead
+  // of the strided outer ones. Clamp the innermost dim to maxChunkSize and
+  // force outer dims to 1; a downstream convert_layout reconciles the
+  // difference with the consumer.
+  for (size_t i = 0; i + 1 < laneData.size(); ++i)
+    laneData[i] = 1;
+  if (!laneData.empty())
+    laneData.back() =
+        std::min(laneData.back(), static_cast<int64_t>(maxChunkSize));
+
   if (layoutKind == xegpu::LayoutKind::InstData) {
     SmallVector<int64_t> instData;
     instData.resize(resShape.size());
@@ -2436,7 +2545,7 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
   auto context = srcVecTy.getContext();
 
   const int subgroupSize = uArch->getSubgroupSize();
-  int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
+  int64_t maxReduceVectorSize = 16; // could extend to spirv vector Size
   xegpu::DistributeLayoutAttr srcLayout;
   if (layoutKind == xegpu::LayoutKind::Subgroup) {
     xegpu::SliceAttr consumerSliceLayout =
@@ -2506,39 +2615,14 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
                               /*laneData=*/{}, resOrderAttr);
     }
   } else if (layoutKind == xegpu::LayoutKind::InstData) {
-    xegpu::SliceAttr consumerSliceLayout =
-        dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
-    auto consumerReductionDims =
-        consumerSliceLayout
-            ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
-            : SmallVector<int64_t>({});
-    // A[i] reduced from A[i, j] is stored out directly, use vertical Lane
-    // layout like [16, 1]
-    bool verticalLaneLayout = consumerReductionDims.empty() &&
-                              reductionDims.size() == 1 &&
-                              reductionDims[0] == (srcRank - 1);
-    auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
-        srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
-        verticalLaneLayout);
-    // inst_data is the per-instruction data, i.e. the element-wise product of
-    // lane_layout and lane_data.
-    SmallVector<int64_t> instData(srcRank);
-    for (int i = 0; i < srcRank; i++)
-      instData[i] = laneLayout[i] * laneData[i];
+    auto [laneLayout, laneData, instData] =
+        computeReductionLaneLayoutAndData(srcShape, reductionDims, subgroupSize,
+                                          maxReduceVectorSize, consumerLayout);
     srcLayout =
         buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
   } else if (layoutKind == xegpu::LayoutKind::Lane) {
-    // Only the innermost two dimensions are distributed; all leading dimensions
-    // are assumed to be unit dimensions.
-    assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
-           "Lane reduction layout assumes all leading (non-innermost-two) "
-           "dimensions are unit dimensions");
     xegpu::SliceAttr consumerSliceLayout =
         dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
-    auto consumerReductionDims =
-        consumerSliceLayout
-            ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
-            : SmallVector<int64_t>({});
     if (consumerSliceLayout &&
         consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
       // at the lane level, the consumerSliceLayout can be directly reused
@@ -2546,12 +2630,17 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
       // the layout is not consistent
       srcLayout = consumerSliceLayout.getParent();
     } else {
-      bool verticalLaneLayout = consumerReductionDims.empty() &&
-                                reductionDims.size() == 1 &&
-                                reductionDims[0] == (srcRank - 1);
-      auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+      // The re-derivation path below only distributes the innermost two dims,
+      // so it requires the leading dims to be unit. The reused-consumer branch
+      // above has no such restriction (a consumer may legitimately keep the
+      // reduce-vector size as per-lane data on a non-unit leading dim).
+      assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
+             "Lane reduction layout assumes all leading (non-innermost-two) "
+             "dimensions are unit dimensions");
+      auto [laneLayout, laneData, instData] = computeReductionLaneLayoutAndData(
           srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
-          verticalLaneLayout);
+          consumerLayout);
+      (void)instData;
       srcLayout = buildLaneLayout(context, laneLayout, laneData);
     }
   }
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index c64defc471de8a..5ce54911214007 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -569,7 +569,7 @@ class MultiRed2dOpPattern
                        xegpu::DistributeLayoutAttr layout) const {
     assert(layout.isForSubgroup() && "Must know the lane layout");
     assert(reductionDims.size() == 2 && "Expected 2D reduction");
-    int64_t intra, cross = -1;
+    int64_t intra = -1, cross = -1;
     xegpu::LayoutAttr layoutAttr = dyn_cast<xegpu::LayoutAttr>(layout);
     if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
       layoutAttr =
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 841f06545eb197..d3bd493178c1f6 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -162,9 +162,9 @@ func.func @load_matrix_with_coalesce_chunksize(%arg0: memref<2048x512xf32>, %arg
 gpu.module @test {
 // CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<512xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} dense<true> : vector<16x32xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} dense<12> : vector<16x32xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<true> : vector<16x32xi1>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<12> : vector<16x32xindex>
+// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
 // CHECK-SAME: memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
 // CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
 // CHECK-SAME: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
@@ -183,10 +183,10 @@ func.func @scatter_ops_coalesce_chunksize(%src: memref<512xf32>) {
 gpu.module @test {
 // CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 1], lane_layout = [16, 1], lane_data = [1, 1]>}
 // CHECK-SAME:  dense<0> : vector<16x16xindex>
-// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 1], lane_layout = [16, 1], lane_data = [1, 1]>} dense<true> : vector<16x16xi1>
+// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 1], lane_layout = [16, 1], lane_data = [1, 1]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
 func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
   %c0 = arith.constant 0 : index
   %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
@@ -208,11 +208,11 @@ func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: m
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
 // CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
-// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
+// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
+// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
 func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
     %cst = arith.constant dense<true> : vector<1024xi1>
     %0 = vector.step : vector<1024xindex>
@@ -232,7 +232,7 @@ func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1
 // -----
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size(
-// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1], lane_layout = [1, 1, 1], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x1xf16> to vector<1xf16>
+// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16, 1], lane_layout = [1, 1, 1], lane_data = [1, 16, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x1xf16> to vector<1xf16>
 func.func @vector_2d_reduction_with_fractional_subgroup_size(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
     %cst = arith.constant dense<true> : vector<16xi1>
     %0 = vector.step : vector<16xindex>
@@ -250,7 +250,7 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size(%arg0: memref<1024x
 // -----
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(
-// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 1, 4], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x4xf16> to vector<1xf16>
+// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16, 4], lane_layout = [1, 1, 4], lane_data = [1, 16, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x4xf16> to vector<1xf16>
 func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
     %cst = arith.constant dense<true> : vector<64xi1>
     %0 = vector.step : vector<64xindex>
@@ -268,9 +268,9 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(%arg0: memref
 // -----
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_shape_cast_expand_and_merge(
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} dense<true> : vector<256xi1>
-// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} : vector<256xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>} dense<true> : vector<256xi1>
+// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>} : vector<256xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
 // CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 32], lane_layout = [1, 1, 16], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
 // CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} : vector<2x4x32xf16> to vector<1x256xf16>
 // CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} : vector<1x256xf16> to vector<256xf16>
@@ -700,3 +700,24 @@ gpu.module @test {
     gpu.return
   }
 }
+
+// -----
+// A 32-bit block load cannot pack along the innermost dim, so the load_nd
+// anchor insists on unit lane_data even though its multi_reduction consumer
+// asks to pack 16 elements per lane into lane_data.
+gpu.module @test {
+// CHECK-LABEL: func.func @f32_reduction_from_load_nd(
+// CHECK: %[[TD:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32, #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TD]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}>
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[LOAD]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 16]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
+func.func @f32_reduction_from_load_nd(%arg0: memref<16x16xf32>, %arg1: memref<256xf32>) {
+  %cst = arith.constant dense<0.000000e+00> : vector<16xf32>
+  %mask = arith.constant dense<true> : vector<16xi1>
+  %offset = vector.step : vector<16xindex>
+  %t = xegpu.create_nd_tdesc %arg0 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
+  %v = xegpu.load_nd %t[0, 0] : !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
+  %0 = vector.multi_reduction <add>, %v, %cst [1] : vector<16x16xf32> to vector<16xf32>
+  xegpu.store %0, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+  return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index d19c0ca17c0cf6..b3fe76c4da17cc 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -138,9 +138,9 @@ func.func @extf_truncf(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor
 gpu.module @test {
 // CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} dense<0> : vector<16x16xindex>
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
-// CHECK: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} dense<0> : vector<16x16xindex>
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} dense<true> : vector<16x16xi1>
+// CHECK: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
 func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
   %c0 = arith.constant 0 : index
   %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
@@ -517,7 +517,7 @@ func.func @if_multiple_uses(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.t
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_outer_reduction(
 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf32>) {
-// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<16x16xf32> to vector<16xf32>
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>, dims = [0]>} [0] : vector<16x16xf32> to vector<16xf32>
 func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf32>) {
   %cst = arith.constant dense<0.000000e+00> : vector<16xf32>
   %mask = arith.constant dense<true> : vector<16xi1>
@@ -531,7 +531,7 @@ func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf3
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_inner_reduction(
 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf32>) {
-// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 16]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
 func.func @vector_inner_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf32>) {
   %cst = arith.constant dense<0.000000e+00> : vector<16xf32>
   %mask = arith.constant dense<true> : vector<16xi1>
@@ -708,11 +708,11 @@ gpu.module @test{
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
 // CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<128xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<128xf16> to vector<1x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
-// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<1x8x16xf16> to vector<8x16xf16>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>} : vector<128xf16> to vector<1x8x16xf16>
+// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
+// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>, dims = [0]>} [0] : vector<1x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
 func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<128xf16>, %arg1: memref<16xf16>) {
     %cst = arith.constant dense<true> : vector<128xi1>
     %0 = vector.step : vector<128xindex>
@@ -862,9 +862,9 @@ func.func @vector_reduction_broadcast_transpose(%arg0: memref<1024x64xf32>, %arg
 // -----
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_shape_cast_expand_and_merge(
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} dense<true> : vector<256xi1>
-// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<256xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<256xi1>
+// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<256xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
 // CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
 // CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>} : vector<2x4x32xf16> to vector<1x256xf16>
 // CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<1x256xf16> to vector<256xf16>

>From 63126b2ace43572386482823ed65e858c20c2c1a Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 20 Sep 2026 03:38:16 +0000
Subject: [PATCH 2/6] [mlir][xegpu] Bound load and reduction layouts by what
 the shape can carry

Three places let a consumer's layout through unchecked, each producing a layout
the value it is attached to cannot distribute.

setupLoadNdAnchorLayout inferred the block variant from the consumer's
lane_data: a 32-bit consumer packing rows made it look for a VNNI transform
block, which exists only below 32 bits. getBlockWidthHeightCount then returned
nothing, the anchor gave up with a null layout, and the load kept the consumer's
lane_data unfiltered -- which requirePacked turned into pack_register on a
32-bit blockload2d, rejected by the XeVM verifier. Gate the transform variant on
the element being narrower than the general packed format size, so a wide load
takes the regular variant and derives unit lane_data as it should.

setupGenericLoadAnchorLayout copied the consumer's lane_layout verbatim. It
already receives resShape, so bound the lane count by it: a cross-subgroup
reduction reloading a 16x4 slab from SLM asked for the 16 lanes it uses on the
wider pre-SLM source, and the load_matrix verifier rejected 16 lanes over 4
columns.

computeReductionLaneLayoutAndData reused a slice consumer's parent layout whole.
A cross-subgroup reduction shares one result layout across both of its stages,
so that parent may describe the 16x32 source of the first stage rather than the
16x4 slab reduced by the second. Clamp the reused lane_layout to what this
source supplies; the parent's lane_data is still right, only the lane count was
borrowed, so clamping keeps agreement with the consumer where deriving afresh
would lose it.

Together these fix the remaining imex integration failures:
softmax_intra_sg_cross_lane and flash_attention_kernel_from_Tile_IR. All nine
previously failing cases now lower, and the full MLIR test suite is unchanged.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 28 ++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 85a4fcdd69168d..e3291e1fb756cd 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1315,6 +1315,15 @@ computeReductionLaneLayoutAndData(
     if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
         static_cast<int>(parentLaneData.size()) == srcRank &&
         parentLanes == subgroupSize) {
+      // Lanes cannot exceed what this source supplies. A cross-subgroup
+      // reduction reuses one result layout for both of its stages, so the
+      // parent may describe the wider pre-SLM source (16x32) rather than the
+      // slab being reduced here (16x4); spreading 16 lanes over 4 columns is
+      // not something the shape can carry.
+      for (int i = 0; i < srcRank; ++i)
+        parentLaneLayout[i] =
+            std::min(parentLaneLayout[i],
+                     std::max<int64_t>(1, srcShape[i] / parentLaneData[i]));
       // Raise, never lower, what the parent asks for.
       for (int i = 0; i < srcRank; ++i)
         if (parentLaneLayout[i] == 1 && srcShape[i] > 1) {
@@ -1903,7 +1912,16 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
   // attr
   bool hasTranspose =
       consumerLaneLayout[rank - 2] > 1 && consumerLaneLayout[rank - 1] == 1;
-  bool hasTransform = !hasTranspose && consumerLaneData[rank - 2] > 1 &&
+  // VNNI exists only below 32 bits, so a wider consumer packing rows into
+  // lane_data is not asking for a transform. Reading it as one sends us looking
+  // for a block variant the hardware does not have, and the whole layout is
+  // then abandoned -- leaving the consumer's lane_data on the load, which is
+  // what pack_register would wrongly be set from.
+  bool canTransform =
+      elemTy.isIntOrFloat() &&
+      elemTy.getIntOrFloatBitWidth() < uArch->getGeneralPackedFormatBitSize();
+  bool hasTransform = !hasTranspose && canTransform &&
+                      consumerLaneData[rank - 2] > 1 &&
                       consumerLaneData[rank - 1] == 1;
   assert((consumerLaneData[rank - 2] == 1 || consumerLaneData[rank - 1] == 1) &&
          "Expected consumer lane data to have at most one non-unit dim");
@@ -2034,6 +2052,14 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
     laneData.back() =
         std::min(laneData.back(), static_cast<int64_t>(maxChunkSize));
 
+  // Lanes are bounded by what this load's own tile can supply. A consumer may
+  // distribute a dim more finely than the tile is wide -- a cross-subgroup
+  // reduction reloading a 16x4 slab wants the 16 lanes it uses elsewhere -- and
+  // keeping that would leave the layout unable to distribute the result.
+  for (size_t i = 0; i < laneLayout.size() && i < resShape.size(); ++i)
+    laneLayout[i] = std::min(laneLayout[i],
+                             std::max<int64_t>(1, resShape[i] / laneData[i]));
+
   if (layoutKind == xegpu::LayoutKind::InstData) {
     SmallVector<int64_t> instData;
     instData.resize(resShape.size());

>From 76d9e765dcc3549c3381baf91f08ca52f6635afa Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 21 Sep 2026 15:35:19 +0000
Subject: [PATCH 3/6] [mlir][xegpu] Restructure reduction lane layout around
 three rules

computeReductionLaneLayoutAndData had two paths that each re-derived the
whole layout, so the shape bound and the reduce-vector placement were
written twice with slightly different meanings. Restate it as three
sequential rules -- take the layout from the consumer, bound lane_layout
by the source shape, then spend the vector budget -- with the reuse and
derive cases differing only in rule 1.

This drops the restriction that lanes live only in the innermost two
dims. A consumer's lane_layout is now honored on whatever dim it lands
on, leading dims included. Reducing src [16, 1, 4] over dim 2 previously
discarded the consumer's 16 lanes on dim 0, reassigned 4 of them to the
reduced dim, and left the feeding gather on lane_layout [4] -- 4 lanes of
16, and a cross-lane combine where a lane-local one was available. It now
yields lane_layout [16, 1, 1]. The assert requiring leading dims to be
unit guarded exactly this case and goes away with it.

Rule 3 also multiplies the chosen dim's lane_data up to the budget
instead of overwriting it, so a dim the consumer already packs can absorb
the remaining budget rather than being skipped.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 191 +++++++-----------
 mlir/test/Dialect/XeGPU/propagate-layout.mlir |  29 +++
 2 files changed, 106 insertions(+), 114 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index e3291e1fb756cd..657bcab4248b4c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -435,17 +435,6 @@ void xegpu::removeTemporaryLayoutAttrs(Operation *op) {
   });
 }
 
-/// Returns true if every dimension of `shape` except the innermost
-/// `numInnerDims` is a unit (size-1) dimension.
-[[maybe_unused]] static bool leadingDimsAreUnit(ArrayRef<int64_t> shape,
-                                                int numInnerDims) {
-  int numLeading = static_cast<int>(shape.size()) - numInnerDims;
-  if (numLeading <= 0)
-    return true;
-  return llvm::all_of(shape.take_front(numLeading),
-                      [](int64_t dim) { return dim == 1; });
-}
-
 static xegpu::LayoutAttr buildInstDataLayoutWithLane(
     mlir::MLIRContext *context, ArrayRef<int64_t> instData,
     ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
@@ -1269,17 +1258,27 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
   return {laneLayout, laneData, order};
 }
 
-/// Computes the (lane_layout, lane_data, inst_data) for a multi-reduction's
-/// source layout, with inst_data = lane_layout * lane_data on every dim.
+/// Computes a multi_reduction's source layout as (lane_layout, lane_data,
+/// inst_data), where inst_data = lane_layout * lane_data on every dim.
+///
+/// `maxReduceVectorSize` is the per-lane vector budget.
 ///
-/// Lanes live only in the innermost two dims. Surviving dims take the
-/// consumer's lane_layout; leftover lanes go to the innermost dim when it is
-/// reduced, so the reduction stays inside the subgroup. A consumer already
-/// spending the whole subgroup thus leaves the reduction dim lane-less and the
-/// reduction lane-local, which is where [16, 1] comes from.
+/// 1. Take lane_layout and lane_data from the consumer. A consumer sliced over
+///    exactly `reductionDims` covers every source dim, so its parent is reused
+///    as-is. Otherwise it covers only the surviving dims, one entry each, and
+///    each entry is honored on its own dim -- leading dims included, which is
+///    what keeps a reduction over a trailing dim lane-local. Lanes the consumer
+///    leaves unclaimed go to the innermost reduced dim rather than sit idle.
 ///
-/// maxReduceVectorSize is a per-lane budget shared across dims: the consumer's
-/// lane_data counts against it and only the remainder is placed here.
+/// 2. Limit lane_layout by `srcShape`: a dim can spread over no more lanes than
+///    it has elements to give.
+///
+/// 3. Raise lane_data until a lane packs `maxReduceVectorSize` elements, on one
+///    undistributed dim -- lane_layout 1, so every lane already spans its full
+///    extent and only inst_data widens. What lane_data packs elsewhere comes off
+///    the budget first, and the chosen dim is multiplied up, never overwritten.
+///    Prefer the dim just outside the innermost two, else the innermost
+///    undistributed dim.
 ///
 /// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
 /// lane_layout=[1,2,8], lane_data=[16,1,1], inst_data=[16,2,8].
@@ -1297,111 +1296,85 @@ computeReductionLaneLayoutAndData(
     return llvm::is_contained(reductionDims, dim);
   };
 
-  // A consumer sliced over exactly the reduction dims already describes every
-  // source dim, so reuse its grid and packing instead of re-deriving something
-  // that could disagree with the IR. The reduce-vector size is raised on top.
+  // Rule 1. Only a parent that covers the source rank and spends the whole
+  // subgroup is trusted; one spending less distributes less than the producer
+  // already does.
   auto consumerSlice = dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+  bool reusedParent = false;
   if (consumerSlice &&
       consumerSlice.getDims().asArrayRef().equals(reductionDims)) {
     auto parent = consumerSlice.getParent();
     SmallVector<int64_t> parentLaneLayout =
         parent.getEffectiveLaneLayoutAsInt();
     SmallVector<int64_t> parentLaneData = parent.getEffectiveLaneDataAsInt();
-    // Only trust a parent that spends the whole subgroup. One that does not
-    // distributes less than the producer already does.
-    int64_t parentLanes = 1;
-    for (int64_t l : parentLaneLayout)
-      parentLanes *= l;
     if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
         static_cast<int>(parentLaneData.size()) == srcRank &&
-        parentLanes == subgroupSize) {
-      // Lanes cannot exceed what this source supplies. A cross-subgroup
-      // reduction reuses one result layout for both of its stages, so the
-      // parent may describe the wider pre-SLM source (16x32) rather than the
-      // slab being reduced here (16x4); spreading 16 lanes over 4 columns is
-      // not something the shape can carry.
-      for (int i = 0; i < srcRank; ++i)
-        parentLaneLayout[i] =
-            std::min(parentLaneLayout[i],
-                     std::max<int64_t>(1, srcShape[i] / parentLaneData[i]));
-      // Raise, never lower, what the parent asks for.
-      for (int i = 0; i < srcRank; ++i)
-        if (parentLaneLayout[i] == 1 && srcShape[i] > 1) {
-          parentLaneData[i] = std::max(
-              parentLaneData[i], std::min(maxReduceVectorSize, srcShape[i]));
-          break;
-        }
-      SmallVector<int64_t> parentInstData(srcRank);
-      for (int i = 0; i < srcRank; ++i)
-        parentInstData[i] = parentLaneLayout[i] * parentLaneData[i];
-      return {parentLaneLayout, parentLaneData, parentInstData};
+        computeProduct(parentLaneLayout) == subgroupSize) {
+      laneLayout = std::move(parentLaneLayout);
+      laneData = std::move(parentLaneData);
+      reusedParent = true;
     }
   }
 
-  // Otherwise the consumer covers only the surviving dims, at reduced rank.
-  SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
-  if (consumerLayout) {
-    consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
-    consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+  if (!reusedParent) {
+    SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
+    if (consumerLayout) {
+      consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+      consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+    }
+    int64_t lanesUsed = 1;
+    for (int i = 0, c = 0; i < srcRank; ++i) {
+      if (isReduced(i))
+        continue;
+      if (c < static_cast<int>(consumerLaneLayout.size())) {
+        laneLayout[i] = consumerLaneLayout[c];
+        lanesUsed *= laneLayout[i];
+      }
+      if (c < static_cast<int>(consumerLaneData.size()))
+        laneData[i] = consumerLaneData[c];
+      ++c;
+    }
+
+    // The innermost reduced dim, not an outer one: claiming lanes further out
+    // needs a redistribution convert_layout cannot lower.
+    int64_t spareLanes = std::max<int64_t>(1, subgroupSize / lanesUsed);
+    int innermost = srcRank - 1;
+    if (spareLanes > 1 && isReduced(innermost))
+      laneLayout[innermost] = spareLanes;
   }
 
-  // Lanes are confined to the innermost two dims.
-  int firstLaneDim = std::max(0, srcRank - 2);
+  // Rule 2. This also bounds a reused parent, which may describe a wider source
+  // than the one reduced here: a cross-subgroup reduction shares one result
+  // layout across both stages, so the parent can ask for 16 lanes on a dim this
+  // slab has only 4 columns of.
+  for (int i = 0; i < srcRank; ++i)
+    laneLayout[i] = std::min(laneLayout[i],
+                             std::max<int64_t>(1, srcShape[i] / laneData[i]));
 
-  // Surviving dims inherit the consumer's grid where lanes may live, and its
-  // packing as the base. Lanes it wants on a leading dim cannot be honored;
-  // convert_layout reconciles that.
-  int64_t lanesUsed = 1;
-  for (int i = 0, c = 0; i < srcRank; ++i) {
-    if (isReduced(i))
-      continue;
-    int64_t consumerLanes = c < static_cast<int>(consumerLaneLayout.size())
-                                ? consumerLaneLayout[c]
-                                : 1;
-    if (c < static_cast<int>(consumerLaneData.size()))
-      laneData[i] = consumerLaneData[c];
-    ++c;
-    if (i < firstLaneDim)
-      continue;
-    laneLayout[i] = std::min<int64_t>(consumerLanes, srcShape[i]);
-    lanesUsed *= laneLayout[i];
-  }
-
-  // Leftover lanes reduce within the subgroup, and only on the innermost dim:
-  // that is where producers spread theirs, so claiming them further out needs a
-  // redistribution convert_layout cannot lower.
-  int64_t spareLanes = subgroupSize / lanesUsed;
-  int innermost = srcRank - 1;
-  if (spareLanes > 1 && isReduced(innermost) && srcShape[innermost] > 1)
-    laneLayout[innermost] = std::min(spareLanes, srcShape[innermost]);
-
-  // Spend what is left of the budget on a lane-less, unpacked dim, since a lane
-  // owns such a dim whole. Prefer the dim just outside the innermost two: in
-  // the scale computation those two come from a shape_cast that split one dim,
-  // so the dim above them is what should carry the per-lane vector.
-  int64_t packedByConsumer = 1;
-  for (int64_t d : laneData)
-    packedByConsumer *= d;
+  // Rule 3. The dim just outside the innermost two is preferred because in the
+  // scale computation those two come from a shape_cast that split one dim, so
+  // the dim above them is what should carry the per-lane vector.
   int64_t vectorBudget =
-      maxReduceVectorSize / std::min(packedByConsumer, maxReduceVectorSize);
-
-  int vectorDim = -1;
-  int outsideLaneDims = srcRank - 3;
-  auto usable = [&](int i) {
-    return laneLayout[i] == 1 && laneData[i] == 1 && srcShape[i] > 1;
+      maxReduceVectorSize /
+      std::min(computeProduct(laneData), maxReduceVectorSize);
+  auto canCarryVector = [&](int i) {
+    return laneLayout[i] == 1 && srcShape[i] > 1;
   };
+  int vectorDim = -1;
   if (vectorBudget > 1) {
-    if (outsideLaneDims >= 0 && usable(outsideLaneDims))
-      vectorDim = outsideLaneDims;
+    int preferredDim = srcRank - 3;
+    if (preferredDim >= 0 && canCarryVector(preferredDim))
+      vectorDim = preferredDim;
     else
       for (int i = srcRank - 1; i >= 0; --i)
-        if (usable(i)) {
+        if (canCarryVector(i)) {
           vectorDim = i;
           break;
         }
   }
   if (vectorDim >= 0)
-    laneData[vectorDim] = std::min(vectorBudget, srcShape[vectorDim]);
+    laneData[vectorDim] =
+        std::min(vectorBudget * laneData[vectorDim], srcShape[vectorDim]);
 
   SmallVector<int64_t> instData(srcRank);
   for (int i = 0; i < srcRank; ++i)
@@ -2466,13 +2439,10 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
 /// (`ResolveLayoutConflicts`), which inserts a `convert_layout` op - this
 /// function never has to give up.
 ///
-/// For the InstData and Lane layout kinds only the innermost two dimensions
-/// are distributed; all leading dimensions are assumed to be unit dimensions.
-/// This assumption is checked via `leadingDimsAreUnit`. The lane_layout and
-/// lane_data are computed by `computeReductionLaneLayoutAndData`, which picks
-/// a layout that minimizes cross-lane reduction (reducing within a lane when
-/// only one of the innermost two dims is a reduction dim). The inst_data is
-/// simply the element-wise product lane_layout * lane_data.
+/// For the InstData and Lane layout kinds the lane_layout and lane_data are
+/// computed by `computeReductionLaneLayoutAndData`, which prefers a lane-local
+/// reduction and spends only otherwise-idle lanes on a cross-lane one. The
+/// inst_data is simply the element-wise product lane_layout * lane_data.
 ///
 /// The function returns the *result* layout (the SliceAttr). The *source*
 /// layout it decides on is the parent of that slice; both are listed below so
@@ -2656,13 +2626,6 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
       // the layout is not consistent
       srcLayout = consumerSliceLayout.getParent();
     } else {
-      // The re-derivation path below only distributes the innermost two dims,
-      // so it requires the leading dims to be unit. The reused-consumer branch
-      // above has no such restriction (a consumer may legitimately keep the
-      // reduce-vector size as per-lane data on a non-unit leading dim).
-      assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
-             "Lane reduction layout assumes all leading (non-innermost-two) "
-             "dimensions are unit dimensions");
       auto [laneLayout, laneData, instData] = computeReductionLaneLayoutAndData(
           srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
           consumerLayout);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index b3fe76c4da17cc..0aba1faeb1cba4 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -1201,3 +1201,32 @@ func.func @broadcast_slice_expanded_dim(%dst: !xegpu.mem_desc<16x4xbf16>) {
   return
 }
 }
+
+// -----
+// A reduction whose only non-unit surviving dim is a *leading* dim: src
+// [16, 1, 4] reducing dim 2. The 16 lanes belong on dim 0, which leaves the
+// reduction lane-local and lets the gather feeding it use the whole subgroup.
+gpu.module @test {
+// CHECK-LABEL: func.func @reduce_innermost_lanes_on_leading(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<64xi1>
+// CHECK: %[[IDX:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<64xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[IDX]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<64xf16>, vector<64xindex>, vector<64xi1> -> vector<64xf16>
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>} : vector<64xf16> to vector<16x1x4xf16>
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>, dims = [2]>} dense<0.000000e+00> : vector<16x1xf16>
+// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[SC]], %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>, dims = [2]>} [2] : vector<16x1x4xf16> to vector<16x1xf16>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[RED]] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16x1xf16> to vector<16xf16>
+// CHECK: xegpu.store %[[FLAT]]
+func.func @reduce_innermost_lanes_on_leading(%arg0: memref<64xf16>, %arg1: memref<16xf16>) {
+    %cst = arith.constant dense<true> : vector<64xi1>
+    %0 = vector.step : vector<64xindex>
+    %1 = xegpu.load %arg0[%0], %cst : memref<64xf16>, vector<64xindex>, vector<64xi1> -> vector<64xf16>
+    %2 = vector.shape_cast %1 : vector<64xf16> to vector<16x1x4xf16>
+    %cst_0 = arith.constant dense<0.000000e+00> : vector<16x1xf16>
+    %3 = vector.multi_reduction <add>, %2, %cst_0 [2] : vector<16x1x4xf16> to vector<16x1xf16>
+    %4 = vector.shape_cast %3 : vector<16x1xf16> to vector<16xf16>
+    %cst_2 = arith.constant dense<true> : vector<16xi1>
+    %cst_3 = arith.constant dense<1> : vector<16xindex>
+    xegpu.store %4, %arg1[%cst_3], %cst_2 : vector<16xf16>, memref<16xf16>, vector<16xindex>, vector<16xi1>
+    return
+  }
+}

>From eb0a5f9cd84992c821a3daf8e615d49ff6b8bb56 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 21 Sep 2026 21:20:18 +0000
Subject: [PATCH 4/6] [mlir][xegpu] Keep the reduce vector off reshaped-from-2D
 sources

Rule 3 of computeReductionLaneLayoutAndData placed the per-lane reduce
vector on any dim the lanes did not split, preferring the dim just
outside the innermost two. On a source reshaped from 2D -- the scale
computation splits one dim in two -- that dim is exactly where the vector
must not go: reducing src [16, 32, 32] over dim 2 produced
lane_data = [16, 1, 1], a non-unit leading dim that lane distribution
cannot consume.

The rule now skips entirely when the dim just outside the innermost two
is non-unit, and otherwise places the vector only on a reduced dim within
the innermost two that has more than one element. The budget is
multiplied onto that dim's existing lane_data rather than assigned, so a
reused parent's packing is built on instead of discarded.

For the scale computation this gives lane_data = [1, 1, 1], clearing one
of the two reasons its reduction could not be distributed; the remaining
one is that lane_layout [1, 2, 8] spreads lanes over two dims, which
SgToLaneVectorExtractStridedSlice does not support.

This supersedes the leading-dim handling from the previous commit: the
leadingDimsAreUnit assert is restored and reduce_innermost_lanes_on_leading
is dropped, since [16, 1, 4] states a 16x4 reduction with the unit dim in
the middle, which that assert's convention excludes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 105 ++++++++++++------
 .../XeGPU/propagate-layout-inst-data.mlir     |   6 +-
 mlir/test/Dialect/XeGPU/propagate-layout.mlir |  28 -----
 3 files changed, 72 insertions(+), 67 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 657bcab4248b4c..07a60b413acc8a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -435,6 +435,17 @@ void xegpu::removeTemporaryLayoutAttrs(Operation *op) {
   });
 }
 
+/// Returns true if every dimension of `shape` except the innermost
+/// `numInnerDims` is a unit (size-1) dimension.
+[[maybe_unused]] static bool leadingDimsAreUnit(ArrayRef<int64_t> shape,
+                                                int numInnerDims) {
+  int numLeading = static_cast<int>(shape.size()) - numInnerDims;
+  if (numLeading <= 0)
+    return true;
+  return llvm::all_of(shape.take_front(numLeading),
+                      [](int64_t dim) { return dim == 1; });
+}
+
 static xegpu::LayoutAttr buildInstDataLayoutWithLane(
     mlir::MLIRContext *context, ArrayRef<int64_t> instData,
     ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
@@ -1266,24 +1277,33 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
 /// 1. Take lane_layout and lane_data from the consumer. A consumer sliced over
 ///    exactly `reductionDims` covers every source dim, so its parent is reused
 ///    as-is. Otherwise it covers only the surviving dims, one entry each, and
-///    each entry is honored on its own dim -- leading dims included, which is
-///    what keeps a reduction over a trailing dim lane-local. Lanes the consumer
-///    leaves unclaimed go to the innermost reduced dim rather than sit idle.
+///    each entry is honored on its own dim. Lanes the consumer leaves unclaimed
+///    go to the innermost reduced dim rather than sit idle.
 ///
 /// 2. Limit lane_layout by `srcShape`: a dim can spread over no more lanes than
 ///    it has elements to give.
 ///
-/// 3. Raise lane_data until a lane packs `maxReduceVectorSize` elements, on one
-///    undistributed dim -- lane_layout 1, so every lane already spans its full
-///    extent and only inst_data widens. What lane_data packs elsewhere comes off
-///    the budget first, and the chosen dim is multiplied up, never overwritten.
-///    Prefer the dim just outside the innermost two, else the innermost
-///    undistributed dim.
+/// 3. Give one dim a per-lane vector, so a lane reduces up to
+///    `maxReduceVectorSize` elements per instruction instead of one.
+///
+///    Which dim: it must be reduced, hold more than one element, and have
+///    lane_layout 1, so a lane owns the whole run it reduces. Among the
+///    innermost two dims, the innermost eligible one is used.
+///
+///    How much: `maxReduceVectorSize / product(lane_data)`, multiplied onto the
+///    dim's existing lane_data rather than replacing it.
+///
+///    Skipped when the dim just outside the innermost two is non-unit: the
+///    source was likely reshaped from 2D -- the scale computation splits one
+///    dim in two -- and no dim of such a source can carry the vector.
 ///
 /// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
-/// lane_layout=[1,2,8], lane_data=[16,1,1], inst_data=[16,2,8].
+/// lane_layout=[1,2,8], lane_data=[1,1,1], inst_data=[1,2,8].
 /// e.g. src=[16,16] reduce [1], consumer lane=[16] ->
 /// lane_layout=[16,1], lane_data=[1,16], inst_data=[16,16].
+/// e.g. src=[16,32] reduce [1], consumer sliced over [1] with parent
+/// lane=[16,1] data=[1,2] -> the parent is reused, then the remaining 16/2 is
+/// multiplied onto dim 1: lane_data=[1,16], inst_data=[16,16].
 static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
                   SmallVector<int64_t>>
 computeReductionLaneLayoutAndData(
@@ -1335,46 +1355,47 @@ computeReductionLaneLayoutAndData(
       ++c;
     }
 
-    // The innermost reduced dim, not an outer one: claiming lanes further out
-    // needs a redistribution convert_layout cannot lower.
+    // The innermost dim only, and only when reduced: an outer one would take
+    // lane_layout [1,2] to [8,2] against a producer still supplying [1,2], and
+    // a lane_layout-only convert_layout has no lowering.
+    // TODO: spend these lanes on an outer reduced dim once that is lowerable --
+    // they sit idle today. Widening it now, even just to srcRank-2, breaks
+    // reduction_2D_scalar_cross_sg_intra_lane.
     int64_t spareLanes = std::max<int64_t>(1, subgroupSize / lanesUsed);
     int innermost = srcRank - 1;
     if (spareLanes > 1 && isReduced(innermost))
       laneLayout[innermost] = spareLanes;
   }
 
-  // Rule 2. This also bounds a reused parent, which may describe a wider source
-  // than the one reduced here: a cross-subgroup reduction shares one result
-  // layout across both stages, so the parent can ask for 16 lanes on a dim this
-  // slab has only 4 columns of.
+  // Rule 2, applied to a reused parent as much as a derived layout.
   for (int i = 0; i < srcRank; ++i)
     laneLayout[i] = std::min(laneLayout[i],
                              std::max<int64_t>(1, srcShape[i] / laneData[i]));
 
-  // Rule 3. The dim just outside the innermost two is preferred because in the
-  // scale computation those two come from a shape_cast that split one dim, so
-  // the dim above them is what should carry the per-lane vector.
+  // Rule 3. A non-unit dim just outside the innermost two means the source was
+  // likely reshaped from 2D -- the scale computation splits one dim in two --
+  // where no dim is a sound place for the vector, so lane_data is left alone.
+  bool mayReshapedFrom2D = srcRank >= 3 && srcShape[srcRank - 3] != 1;
   int64_t vectorBudget =
-      maxReduceVectorSize /
-      std::min(computeProduct(laneData), maxReduceVectorSize);
+      mayReshapedFrom2D
+          ? 1
+          : maxReduceVectorSize /
+                std::min(computeProduct(laneData), maxReduceVectorSize);
+
+  // Exactly one dim carries the vector: it must be unsplit by lanes (so the
+  // lane owns it whole), reduced (that is what the vector feeds), and non-unit.
   auto canCarryVector = [&](int i) {
-    return laneLayout[i] == 1 && srcShape[i] > 1;
+    return laneLayout[i] == 1 && isReduced(i) && srcShape[i] > 1;
   };
+  int innermost = srcRank - 1;
   int vectorDim = -1;
-  if (vectorBudget > 1) {
-    int preferredDim = srcRank - 3;
-    if (preferredDim >= 0 && canCarryVector(preferredDim))
-      vectorDim = preferredDim;
-    else
-      for (int i = srcRank - 1; i >= 0; --i)
-        if (canCarryVector(i)) {
-          vectorDim = i;
-          break;
-        }
-  }
-  if (vectorDim >= 0)
-    laneData[vectorDim] =
-        std::min(vectorBudget * laneData[vectorDim], srcShape[vectorDim]);
+  if (canCarryVector(innermost))
+    vectorDim = innermost;
+  else if (srcRank >= 2 && canCarryVector(innermost - 1))
+    vectorDim = innermost - 1;
+
+  if (vectorBudget > 1 && vectorDim >= 0)
+    laneData[vectorDim] = std::min(vectorBudget, srcShape[vectorDim]);
 
   SmallVector<int64_t> instData(srcRank);
   for (int i = 0; i < srcRank; ++i)
@@ -2444,6 +2465,13 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
 /// reduction and spends only otherwise-idle lanes on a cross-lane one. The
 /// inst_data is simply the element-wise product lane_layout * lane_data.
 ///
+/// At the Lane kind a consumer sliced over exactly `reductionDims` is reused
+/// verbatim: the InstData run already settled the packing and inserted any
+/// `convert_layout` needed, so this stage adopts that decision rather than
+/// re-deriving one. Re-deriving is reserved for the other consumers, and
+/// assumes all leading (non-innermost-two) dims are unit, as
+/// `leadingDimsAreUnit` asserts.
+///
 /// The function returns the *result* layout (the SliceAttr). The *source*
 /// layout it decides on is the parent of that slice; both are listed below so
 /// the relationship is explicit.
@@ -2626,6 +2654,11 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
       // the layout is not consistent
       srcLayout = consumerSliceLayout.getParent();
     } else {
+      // Only the innermost two dimensions are distributed; all leading
+      // dimensions are assumed to be unit dimensions.
+      assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
+             "Lane reduction layout assumes all leading (non-innermost-two) "
+             "dimensions are unit dimensions");
       auto [laneLayout, laneData, instData] = computeReductionLaneLayoutAndData(
           srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
           consumerLayout);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index d3bd493178c1f6..a7f9ff34ca1aeb 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -208,10 +208,10 @@ func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: m
 gpu.module @test {
 // CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
 // CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 8, 16], lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
+// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 8, 16], lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
 // CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 8, 16], lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
 // CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
 func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
     %cst = arith.constant dense<true> : vector<1024xi1>
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 0aba1faeb1cba4..b0f9c0a929e3b7 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -1202,31 +1202,3 @@ func.func @broadcast_slice_expanded_dim(%dst: !xegpu.mem_desc<16x4xbf16>) {
 }
 }
 
-// -----
-// A reduction whose only non-unit surviving dim is a *leading* dim: src
-// [16, 1, 4] reducing dim 2. The 16 lanes belong on dim 0, which leaves the
-// reduction lane-local and lets the gather feeding it use the whole subgroup.
-gpu.module @test {
-// CHECK-LABEL: func.func @reduce_innermost_lanes_on_leading(
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<64xi1>
-// CHECK: %[[IDX:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<64xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[IDX]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<64xf16>, vector<64xindex>, vector<64xi1> -> vector<64xf16>
-// CHECK: %[[SC:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>} : vector<64xf16> to vector<16x1x4xf16>
-// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>, dims = [2]>} dense<0.000000e+00> : vector<16x1xf16>
-// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[SC]], %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>, dims = [2]>} [2] : vector<16x1x4xf16> to vector<16x1xf16>
-// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[RED]] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16x1xf16> to vector<16xf16>
-// CHECK: xegpu.store %[[FLAT]]
-func.func @reduce_innermost_lanes_on_leading(%arg0: memref<64xf16>, %arg1: memref<16xf16>) {
-    %cst = arith.constant dense<true> : vector<64xi1>
-    %0 = vector.step : vector<64xindex>
-    %1 = xegpu.load %arg0[%0], %cst : memref<64xf16>, vector<64xindex>, vector<64xi1> -> vector<64xf16>
-    %2 = vector.shape_cast %1 : vector<64xf16> to vector<16x1x4xf16>
-    %cst_0 = arith.constant dense<0.000000e+00> : vector<16x1xf16>
-    %3 = vector.multi_reduction <add>, %2, %cst_0 [2] : vector<16x1x4xf16> to vector<16x1xf16>
-    %4 = vector.shape_cast %3 : vector<16x1xf16> to vector<16xf16>
-    %cst_2 = arith.constant dense<true> : vector<16xi1>
-    %cst_3 = arith.constant dense<1> : vector<16xindex>
-    xegpu.store %4, %arg1[%cst_3], %cst_2 : vector<16xf16>, memref<16xf16>, vector<16xindex>, vector<16xi1>
-    return
-  }
-}

>From cddf8c738267303872e4ede5f9e6549ee8848a33 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 23 Sep 2026 16:58:03 +0000
Subject: [PATCH 5/6] [mlir][xegpu] Drop the redundant lane_layout clamp on
 generic loads

setupGenericLoadAnchorLayout bounded the consumer's lane_layout by
resShape / lane_data on every dim, added alongside the reduction-side
bound for a cross-subgroup reduction reloading a narrow slab. Removing it
leaves both suites green, including the flash_attention case it cited:
rule 2 of computeReductionLaneLayoutAndData now applies that bound to a
reused parent as well, which covers the same ground.

The lane_data clamp stays -- both of its halves are load-bearing. Forcing
outer dims to 1 is what three cross-subgroup reduction tests need, and
capping the innermost dim at maxChunkSize is what
scatter_ops_coalesce_chunksize needs. It is now built directly rather than
copied from the consumer and overwritten, since only the innermost element
survives that copy.

Also corrects the Lane examples on setupMultiReductionResultLayout, which
still showed laneData=[1,1] from before maxReduceVectorSize was raised to
16. Each is an instance of the reduce vector landing on whichever dim the
lanes vacate: [32,64] reducing dim 0 gives laneData=[16,1], reducing dim 1
gives [1,16], reducing both gives [16,1].

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 81 +++++++------------
 1 file changed, 29 insertions(+), 52 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 07a60b413acc8a..0576096cc7d988 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1864,8 +1864,7 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
 /// consumer layout (from its result's downstream uses) and validates it
 /// against uArch constraints; if valid, the consumer's `inst_data` /
 /// `sg_layout` are honored. Otherwise the helper falls back to defaults
-/// derived from uArch block parameters. `lane_data` never comes from the
-/// consumer; see compute2DBlockIOLaneLayout.
+/// derived from uArch block parameters.
 xegpu::DistributeLayoutAttr
 xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
                                VectorType resVecTy,
@@ -1906,11 +1905,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
   // attr
   bool hasTranspose =
       consumerLaneLayout[rank - 2] > 1 && consumerLaneLayout[rank - 1] == 1;
-  // VNNI exists only below 32 bits, so a wider consumer packing rows into
-  // lane_data is not asking for a transform. Reading it as one sends us looking
-  // for a block variant the hardware does not have, and the whole layout is
-  // then abandoned -- leaving the consumer's lane_data on the load, which is
-  // what pack_register would wrongly be set from.
+  // VNNI exists only below 32 bits.
   bool canTransform =
       elemTy.isIntOrFloat() &&
       elemTy.getIntOrFloatBitWidth() < uArch->getGeneralPackedFormatBitSize();
@@ -2025,34 +2020,15 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
   SmallVector<int64_t> consumerLaneData =
       consumerLayout.getEffectiveLaneDataAsInt();
 
-  SmallVector<int64_t> laneLayout;
-  SmallVector<int64_t> laneData;
   assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
          "Expected consumer layout to have lane_layout and lane_data");
-  laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
-  laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
-
-  // A matrix/gather load reads a single contiguous span of memory per lane, so
-  // a lane can only pack (lane_data > 1) elements along the innermost,
-  // contiguous dim. A consumer that requests lane_data > 1 on an outer, strided
-  // dim (e.g. a downstream reduction packing across mem_desc rows) is not
-  // lowerable: the contiguous load would fetch adjacent inner elements instead
-  // of the strided outer ones. Clamp the innermost dim to maxChunkSize and
-  // force outer dims to 1; a downstream convert_layout reconciles the
-  // difference with the consumer.
-  for (size_t i = 0; i + 1 < laneData.size(); ++i)
-    laneData[i] = 1;
-  if (!laneData.empty())
-    laneData.back() =
-        std::min(laneData.back(), static_cast<int64_t>(maxChunkSize));
-
-  // Lanes are bounded by what this load's own tile can supply. A consumer may
-  // distribute a dim more finely than the tile is wide -- a cross-subgroup
-  // reduction reloading a 16x4 slab wants the 16 lanes it uses elsewhere -- and
-  // keeping that would leave the layout unable to distribute the result.
-  for (size_t i = 0; i < laneLayout.size() && i < resShape.size(); ++i)
-    laneLayout[i] = std::min(laneLayout[i],
-                             std::max<int64_t>(1, resShape[i] / laneData[i]));
+  SmallVector<int64_t> laneLayout(consumerLaneLayout);
+
+  // A matrix/gather load reads one contiguous span per lane, so a lane can pack
+  // only along the innermost dim, and only up to maxChunkSize.
+  SmallVector<int64_t> laneData(consumerLaneData.size(), 1);
+  laneData.back() =
+      std::min(consumerLaneData.back(), static_cast<int64_t>(maxChunkSize));
 
   if (layoutKind == xegpu::LayoutKind::InstData) {
     SmallVector<int64_t> instData;
@@ -2460,16 +2436,13 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
 /// (`ResolveLayoutConflicts`), which inserts a `convert_layout` op - this
 /// function never has to give up.
 ///
-/// For the InstData and Lane layout kinds the lane_layout and lane_data are
-/// computed by `computeReductionLaneLayoutAndData`, which prefers a lane-local
-/// reduction and spends only otherwise-idle lanes on a cross-lane one. The
-/// inst_data is simply the element-wise product lane_layout * lane_data.
+/// InstData kind: `computeReductionLaneLayoutAndData` derives lane_layout and
+/// lane_data, and inst_data is their element-wise product.
 ///
-/// At the Lane kind a consumer sliced over exactly `reductionDims` is reused
-/// verbatim: the InstData run already settled the packing and inserted any
-/// `convert_layout` needed, so this stage adopts that decision rather than
-/// re-deriving one. Re-deriving is reserved for the other consumers, and
-/// assumes all leading (non-innermost-two) dims are unit, as
+/// Lane kind: a consumer sliced over exactly `reductionDims` is reused verbatim
+/// rather than re-derived, so this stage does not re-decide packing an earlier
+/// run already committed to. Any other consumer goes through the same
+/// derivation, which assumes all leading (non-innermost-two) dims are unit, as
 /// `leadingDimsAreUnit` asserts.
 ///
 /// The function returns the *result* layout (the SliceAttr). The *source*
@@ -2522,28 +2495,32 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
 ///   3. Lane layout - Default (lanes on innermost dim):
 ///      srcShape=[32, 64], reductionDims=[0], subgroupSize=16
 ///      * Source Layout (decided by this function):
-///        laneLayout=[1, 16], laneData=[1, 1] (returned sliced over dim 0).
-///      The innermost dim is not reduced, so lanes stay on it.
+///        laneLayout=[1, 16], laneData=[16, 1] (returned sliced over dim 0).
+///      The innermost dim is not reduced, so lanes stay on it. Reduced dim 0 is
+///      then lane-free and takes the per-lane reduce vector:
+///      laneData[0] = min(maxReduceVectorSize, 32) = 16.
 ///
 ///   4. Lane layout - Switch (lanes moved off the reduction dim):
 ///      srcShape=[32, 64], reductionDims=[1], subgroupSize=16
 ///      * Source Layout (decided by this function):
-///        laneLayout=[16, 1], laneData=[1, 1] (returned sliced over dim 1).
+///        laneLayout=[16, 1], laneData=[1, 16] (returned sliced over dim 1).
 ///      The innermost dim is the sole reduction dim, so lanes move to the
-///      non-reduction dim to reduce within a lane. This switch only happens
-///      when the consumer has no reduction dims to broadcast the result back
-///      along (i.e. the consumer layout is not a slice over this reduction);
-///      otherwise the default (example 3) is used.
+///      non-reduction dim to reduce within a lane, and the now lane-free
+///      reduction dim takes the per-lane reduce vector (16 of its 64 elements).
+///      This switch only happens when the consumer has no reduction dims to
+///      broadcast the result back along (i.e. the consumer layout is not a
+///      slice over this reduction); otherwise the default (example 3) is used.
 ///
 ///   5. Lane layout - No switch when both inner dims are reduced (reduction to
 ///   scalar):
 ///      srcShape=[32, 64], reductionDims=[0, 1], subgroupSize=16
 ///      * Source Layout (decided by this function):
-///        laneLayout=[1, 16], laneData=[1, 1] (returned sliced over dims
+///        laneLayout=[1, 16], laneData=[16, 1] (returned sliced over dims
 ///        [0,1]).
 ///      Both dims are reduced, so this is not a *sole* innermost reduction; the
 ///      switch condition (example 4) does not apply and lanes stay on the
-///      innermost dim. The cross-lane reduction here is unavoidable.
+///      innermost dim. Reduced dim 0 is lane-free, so it takes the reduce
+///      vector. The cross-lane reduction over dim 1 here is unavoidable.
 ///
 ///   6. Lane layout - No switch when the consumer slices the reduction dim:
 ///      srcShape=[32, 64], reductionDims=[1], subgroupSize=16
@@ -2569,7 +2546,7 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
   auto context = srcVecTy.getContext();
 
   const int subgroupSize = uArch->getSubgroupSize();
-  int64_t maxReduceVectorSize = 16; // could extend to spirv vector Size
+  int64_t maxReduceVectorSize = 16; // extend to spirv vector Size
   xegpu::DistributeLayoutAttr srcLayout;
   if (layoutKind == xegpu::LayoutKind::Subgroup) {
     xegpu::SliceAttr consumerSliceLayout =

>From 32ecfc5f736d2a009f09d13550d1b83cc88b35cb Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 23 Sep 2026 18:45:10 +0000
Subject: [PATCH 6/6] [mlir][xegpu] Redistribute lanes freed by the reduction
 shape bound

The spare-lane step ran before the srcShape bound, so lanes the bound
freed were never handed out: bounding lane_layout on one dim leaves the
grid narrower than the subgroup, and nothing noticed. It now runs after
the bound, over all dims rather than only the derived path, and takes the
leftover from computeProduct(laneLayout) instead of an accumulator. Since
a reused parent can already carry lanes on the innermost dim, the leftover
is multiplied on rather than assigned, then re-capped: for src [16, 4]
reducing dim 1 the bound gives lane_layout [1, 4], and 4 spare lanes would
otherwise put 16 lanes on a 4-wide dim.

Leaving lanes unused is not just wasteful. Dropping the step entirely
fails eight imex tests, both flash-attention kernels among them, with an
unlowerable lane_layout-only convert_layout ([1,16] -> [1,1]): the
producer distributes over the whole subgroup, so a narrower grid here has
nothing to bridge it.

Rule 3 spends the budget by multiplying onto the dim's existing lane_data
again, so a dim a reused parent already packs reaches maxReduceVectorSize
instead of being overwritten below it.

Drops the mayReshapedFrom2D guard added in the previous commit. Its job
was to keep the reduce vector off a source reshaped from 2D, but rule 3's
eligibility test already does that: for src [16, 32, 32] reducing dim 2,
the innermost dim carries the spare lanes so lane_layout is not 1, and the
dim beside it is not reduced. The scale computation still gets
lane_data = [1, 1, 1].

Also rewrites the rule list for two review comments: "entry" and "it" in
rule 1 never said what they referred to, and rule 3 opened with "give one
dim a per-lane vector", naming neither the field it changes nor a defined
term.

Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
 .../XeGPU/Transforms/XeGPULayoutImpl.cpp      | 78 +++++++++----------
 1 file changed, 36 insertions(+), 42 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 0576096cc7d988..52991a2d576c10 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1275,35 +1275,34 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
 /// `maxReduceVectorSize` is the per-lane vector budget.
 ///
 /// 1. Take lane_layout and lane_data from the consumer. A consumer sliced over
-///    exactly `reductionDims` covers every source dim, so its parent is reused
-///    as-is. Otherwise it covers only the surviving dims, one entry each, and
-///    each entry is honored on its own dim. Lanes the consumer leaves unclaimed
-///    go to the innermost reduced dim rather than sit idle.
+///    exactly `reductionDims` is a view of a source-rank parent, and that
+///    parent is reused as-is. Any other consumer has the result's rank: its
+///    lane_layout and lane_data hold one element per surviving dim, taken in
+///    order. Reduced dims keep lane_layout 1 for now.
 ///
 /// 2. Limit lane_layout by `srcShape`: a dim can spread over no more lanes than
 ///    it has elements to give.
 ///
-/// 3. Give one dim a per-lane vector, so a lane reduces up to
-///    `maxReduceVectorSize` elements per instruction instead of one.
+///    Then hand any unused lanes to the innermost dim, if that dim is reduced.
+///    This is a heuristic: the innermost dim of the source typically already
+///    carries several lanes, whether the source is loaded from memory or is a
+///    dpas result, so putting them there matches what the producer did.
+///
+/// 3. Raise lane_data on exactly one dim, so the lane holds a run of elements
+///    along that dim and reduces the run with one vector op instead of one
+///    element at a time. The run is at most `maxReduceVectorSize` long.
 ///
 ///    Which dim: it must be reduced, hold more than one element, and have
 ///    lane_layout 1, so a lane owns the whole run it reduces. Among the
 ///    innermost two dims, the innermost eligible one is used.
 ///
 ///    How much: `maxReduceVectorSize / product(lane_data)`, multiplied onto the
-///    dim's existing lane_data rather than replacing it.
-///
-///    Skipped when the dim just outside the innermost two is non-unit: the
-///    source was likely reshaped from 2D -- the scale computation splits one
-///    dim in two -- and no dim of such a source can carry the vector.
+///    dim's existing lane_data and capped by its extent.
 ///
 /// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
 /// lane_layout=[1,2,8], lane_data=[1,1,1], inst_data=[1,2,8].
 /// e.g. src=[16,16] reduce [1], consumer lane=[16] ->
 /// lane_layout=[16,1], lane_data=[1,16], inst_data=[16,16].
-/// e.g. src=[16,32] reduce [1], consumer sliced over [1] with parent
-/// lane=[16,1] data=[1,2] -> the parent is reused, then the remaining 16/2 is
-/// multiplied onto dim 1: lane_data=[1,16], inst_data=[16,16].
 static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
                   SmallVector<int64_t>>
 computeReductionLaneLayoutAndData(
@@ -1342,52 +1341,46 @@ computeReductionLaneLayoutAndData(
       consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
       consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
     }
-    int64_t lanesUsed = 1;
     for (int i = 0, c = 0; i < srcRank; ++i) {
       if (isReduced(i))
         continue;
-      if (c < static_cast<int>(consumerLaneLayout.size())) {
+      if (c < static_cast<int>(consumerLaneLayout.size()))
         laneLayout[i] = consumerLaneLayout[c];
-        lanesUsed *= laneLayout[i];
-      }
       if (c < static_cast<int>(consumerLaneData.size()))
         laneData[i] = consumerLaneData[c];
       ++c;
     }
-
-    // The innermost dim only, and only when reduced: an outer one would take
-    // lane_layout [1,2] to [8,2] against a producer still supplying [1,2], and
-    // a lane_layout-only convert_layout has no lowering.
-    // TODO: spend these lanes on an outer reduced dim once that is lowerable --
-    // they sit idle today. Widening it now, even just to srcRank-2, breaks
-    // reduction_2D_scalar_cross_sg_intra_lane.
-    int64_t spareLanes = std::max<int64_t>(1, subgroupSize / lanesUsed);
-    int innermost = srcRank - 1;
-    if (spareLanes > 1 && isReduced(innermost))
-      laneLayout[innermost] = spareLanes;
   }
 
   // Rule 2, applied to a reused parent as much as a derived layout.
-  for (int i = 0; i < srcRank; ++i)
-    laneLayout[i] = std::min(laneLayout[i],
-                             std::max<int64_t>(1, srcShape[i] / laneData[i]));
+  SmallVector<int64_t> maxLanes(srcRank);
+  for (int i = 0; i < srcRank; ++i) {
+    maxLanes[i] = std::max<int64_t>(1, srcShape[i] / laneData[i]);
+    laneLayout[i] = std::min(laneLayout[i], maxLanes[i]);
+  }
+
+  // Lanes the consumer never claimed, plus any the bound above freed, go to the
+  // innermost dim -- and only when it is reduced.
+  // TODO: spend these lanes on an outer reduced dim once that is lowerable --
+  // they sit idle today. Widening it to srcRank-2 breaks
+  // reduction_2D_scalar_cross_sg_intra_lane.
+  int innermost = srcRank - 1;
+  int64_t spareLanes =
+      std::max<int64_t>(1, subgroupSize / computeProduct(laneLayout));
+  if (spareLanes > 1 && isReduced(innermost))
+    laneLayout[innermost] =
+        std::min(laneLayout[innermost] * spareLanes, maxLanes[innermost]);
 
-  // Rule 3. A non-unit dim just outside the innermost two means the source was
-  // likely reshaped from 2D -- the scale computation splits one dim in two --
-  // where no dim is a sound place for the vector, so lane_data is left alone.
-  bool mayReshapedFrom2D = srcRank >= 3 && srcShape[srcRank - 3] != 1;
+  // Rule 3. What lane_data already packs comes off the budget first.
   int64_t vectorBudget =
-      mayReshapedFrom2D
-          ? 1
-          : maxReduceVectorSize /
-                std::min(computeProduct(laneData), maxReduceVectorSize);
+      maxReduceVectorSize /
+      std::min(computeProduct(laneData), maxReduceVectorSize);
 
   // Exactly one dim carries the vector: it must be unsplit by lanes (so the
   // lane owns it whole), reduced (that is what the vector feeds), and non-unit.
   auto canCarryVector = [&](int i) {
     return laneLayout[i] == 1 && isReduced(i) && srcShape[i] > 1;
   };
-  int innermost = srcRank - 1;
   int vectorDim = -1;
   if (canCarryVector(innermost))
     vectorDim = innermost;
@@ -1395,7 +1388,8 @@ computeReductionLaneLayoutAndData(
     vectorDim = innermost - 1;
 
   if (vectorBudget > 1 && vectorDim >= 0)
-    laneData[vectorDim] = std::min(vectorBudget, srcShape[vectorDim]);
+    laneData[vectorDim] =
+        std::min(vectorBudget * laneData[vectorDim], srcShape[vectorDim]);
 
   SmallVector<int64_t> instData(srcRank);
   for (int i = 0; i < srcRank; ++i)