[Mlir-commits] [mlir] 894dad7 - [mlir][xegpu] Raise the reduction vector size to match maxReduceVectorSize (#225122)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 3 22:52:37 PDT 2026
Author: Jianhui Li
Date: 2026-10-03T22:52:30-07:00
New Revision: 894dad791c0124915dcf6b5093744c9307cffaf2
URL: https://github.com/llvm/llvm-project/commit/894dad791c0124915dcf6b5093744c9307cffaf2
DIFF: https://github.com/llvm/llvm-project/commit/894dad791c0124915dcf6b5093744c9307cffaf2.diff
LOG: [mlir][xegpu] Raise the reduction vector size to match maxReduceVectorSize (#225122)
`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; // extend to spirv vector Size
```
The PR only raises the vector size for the innermost dim, and only when
that dim is reduced and carries a single lane.
The reason we don't raise vector size for outer dim: Raising the vector
size for the outer dim means the vector may be strided in flat element
order. In case that the producer is a shape_cast op, it causes a
potential convert_layout after shape_cast which assumes all lane data to
be contiguous.
As we simplify propagation rule, the trade-off is that the rule is
conservative since we don't know the information about producer due to
the backward propagation.
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
mlir/test/Dialect/XeGPU/propagate-layout.mlir
mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 4dda10c6a7a26..031d971cebe6c 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,122 @@ 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` 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.
+///
+/// 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 the innermost dim, so the lane holds a run of elements
+/// along it and reduces the run with one vector op instead of one element at
+/// a time. The run is at most `maxReduceVectorSize` long. The dim must also
+/// be reduced and have lane_layout 1.
+///
+/// How much: `maxReduceVectorSize / product(lane_data)`, multiplied onto the
+/// 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].
+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;
+ // 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;
+ }
+ }
- if (verticalLaneLayout && secondInnermost >= 0) {
- std::swap(innermost, secondInnermost);
+ if (!reusedParent) {
+ SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
+ if (consumerLayout) {
+ consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+ }
+ for (int i = 0, c = 0; i < srcRank; ++i) {
+ if (isReduced(i))
+ continue;
+ if (c < static_cast<int>(consumerLaneLayout.size()))
+ laneLayout[i] = consumerLaneLayout[c];
+ if (c < static_cast<int>(consumerLaneData.size()))
+ laneData[i] = consumerLaneData[c];
+ ++c;
+ }
}
- int laneDim = innermost;
- int vectorDim = secondInnermost; // negative for rank 1
- laneLayout[laneDim] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
- if (vectorDim >= 0)
- laneData[vectorDim] = std::min(maxReduceVectorSize, srcShape[vectorDim]);
+ // Rule 2, applied to a reused parent as much as a derived layout.
+ 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]);
+ }
- return {laneLayout, laneData};
+ // 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. What lane_data already packs comes off the budget first.
+ int64_t vectorBudget =
+ maxReduceVectorSize /
+ std::min(computeProduct(laneData), maxReduceVectorSize);
+
+ bool canCarryVector = laneLayout[innermost] == 1 && isReduced(innermost) &&
+ srcShape[innermost] > 1;
+
+ if (vectorBudget > 1 && canCarryVector)
+ laneData[innermost] =
+ std::min(vectorBudget * laneData[innermost], srcShape[innermost]);
+
+ SmallVector<int64_t> instData(srcRank);
+ for (int i = 0; i < srcRank; ++i)
+ instData[i] = laneLayout[i] * laneData[i];
+
+ return {laneLayout, laneData, instData};
}
//===----------------------------------------------------------------------===//
@@ -1808,7 +1877,12 @@ 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.
+ 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");
@@ -1918,12 +1992,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());
+ 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;
@@ -2331,13 +2408,13 @@ 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.
+/// InstData kind: `computeReductionLaneLayoutAndData` derives lane_layout and
+/// lane_data, and inst_data is their element-wise product.
+///
+/// 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.
///
/// 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
@@ -2390,17 +2467,20 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
/// 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.
+/// The innermost dim is not reduced, so lanes stay on it. Reduced dim 0 is
+/// then lane-free, but it is not the innermost dim, so it takes no reduce
+/// vector and the reduction runs one element at a time.
///
/// 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):
@@ -2410,7 +2490,7 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
/// [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. The cross-lane reduction over dim 1 is unavoidable.
///
/// 6. Lane layout - No switch when the consumer slices the reduction dim:
/// srcShape=[32, 64], reductionDims=[1], subgroupSize=16
@@ -2436,7 +2516,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; // extend to spirv vector Size
xegpu::DistributeLayoutAttr srcLayout;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
xegpu::SliceAttr consumerSliceLayout =
@@ -2531,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
@@ -2571,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 27f33cdb92c57..b430a5430c3c7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -564,7 +564,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/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 744606c3a553b..9533d7bd5d3ae 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2723,7 +2723,8 @@ void XeGPUSgToLaneDistributePass::runOnOperation() {
xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
typeConverter, patterns, target, root);
target.addLegalOp<UnrealizedConversionCastOp>();
- (void)applyPartialConversion(root, target, std::move(patterns));
+ if (failed(applyPartialConversion(root, target, std::move(patterns))))
+ return signalPassFailure();
}
// Fold cancelling cast chains and erase dead casts.
xegpu::cleanupUnrealizedConversionCasts(root, existingCasts);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 05c5ec0e4b036..30e02cc42095c 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>
@@ -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>
@@ -736,3 +736,39 @@ func.func @shape_cast_split_fills_undistributed_inner_dim(%arg0: memref<16x256xb
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
+}
+}
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @shape_cast_split_fills_strided_inner_dim(
+// CHECK: %[[CAST:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 2, 32], lane_layout = [1, 2, 8], lane_data = [1, 1, 4]>} : vector<16x1024xbf16> to vector<16x32x32xbf16>
+func.func @shape_cast_split_fills_strided_inner_dim(%arg0: memref<16x1024xbf16>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x1024xbf16> -> !xegpu.tensor_desc<16x1024xbf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x1024xbf16> -> vector<16x1024xbf16>
+ %2 = vector.shape_cast %1 : vector<16x1024xbf16> to vector<16x32x32xbf16>
+ %3 = xegpu.convert_layout %2
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 8], lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>}>
+ : vector<16x32x32xbf16>
+ return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index d19c0ca17c0cf..4d1d28c4e0a5c 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>
@@ -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>
@@ -829,6 +829,32 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size_1x1x4(%arg0: memref
}
}
+// -----
+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
+ }
+}
+
// -----
gpu.module @test {
// CHECK: func.func @vector_reduction_broadcast_transpose(%[[ARG0:.*]]: memref<1024x64xf32>, %[[ARG1:.*]]: memref<1024x64xf32>)
@@ -862,9 +888,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>
@@ -1201,3 +1227,4 @@ func.func @broadcast_slice_expanded_dim(%dst: !xegpu.mem_desc<16x4xbf16>) {
return
}
}
+
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
index 8ed375dac0c9d..4ccea7115eb17 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
@@ -26,7 +26,8 @@ module @gemm attributes {gpu.container_module} {
// along the K dimension with block size 32. B and its scale are passed in
// pre-quantized (packed ui8 fp4 and f8E8M0). The quantized values are then
// consumed by xegpu.dpas_mx.
- gpu.func @gemm_mxfp(%arg0: memref<256x4096xbf16>, %arg1: memref<2048x256xi8>, %arg3: memref<128x256xf8E8M0FNU>, %arg4: memref<256x256xf32>) kernel {
+ gpu.func @gemm_mxfp(%arg0: memref<256x4096xbf16>, %arg1: memref<2048x256xi8>, %arg3: memref<128x256xf8E8M0FNU>, %arg4: memref<256x256xf32>) kernel attributes
+ {known_block_size = array<i32: 64, 1, 1>} {
%c0 = arith.constant 0 : index
%mstep = arith.constant 32 : index
%nstep = arith.constant 32 : index
More information about the Mlir-commits
mailing list