[Mlir-commits] [mlir] [mlir][xegpu] Raise the reduction vector size to match maxReduceVectorSize (PR #225122)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 21 08:51:16 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
`maxReduceVectorSize` was pinned at 1, so a lane reduced one element at a time no
matter how much data it held. This raises it to 16 and makes the layout
computation actually spend it, so a lane's reduce vector matches the per-lane
budget.
```cpp
- int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
+ int64_t maxReduceVectorSize = 16; // could extend to spirv vector Size
```
---
Patch is 38.61 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/225122.diff
4 Files Affected:
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+160-82)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp (+1-1)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+37-16)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+42-13)
``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 05b3f24d660b4..657bcab4248b4 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,42 +1258,129 @@ 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 a multi_reduction's source layout as (lane_layout, lane_data,
+/// inst_data), where 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) {
+/// `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 -- 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.
+///
+/// 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].
+/// 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;
-
- if (verticalLaneLayout && secondInnermost >= 0) {
- std::swap(innermost, secondInnermost);
+ // 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();
+ if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
+ static_cast<int>(parentLaneData.size()) == srcRank &&
+ computeProduct(parentLaneLayout) == subgroupSize) {
+ laneLayout = std::move(parentLaneLayout);
+ laneData = std::move(parentLaneData);
+ reusedParent = true;
+ }
}
- int laneDim = innermost;
- int vectorDim = secondInnermost; // negative for rank 1
- laneLayout[laneDim] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
+ 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;
+ }
+
+ // 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]));
+
+ // 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(computeProduct(laneData), maxReduceVectorSize);
+ auto canCarryVector = [&](int i) {
+ return laneLayout[i] == 1 && srcShape[i] > 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(maxReduceVectorSize, srcShape[vectorDim]);
+ laneData[vectorDim] =
+ std::min(vectorBudget * laneData[vectorDim], 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 +1843,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,
@@ -1808,7 +1885,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");
@@ -1925,6 +2011,28 @@ 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));
+
+ // 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());
@@ -2331,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
@@ -2436,7 +2541,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 +2611,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 +2626,10 @@ 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(
+ 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 c64defc471de8..5ce5491121400 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 841f06545eb19..d3bd493178c1f 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...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/225122
More information about the Mlir-commits
mailing list