[Mlir-commits] [mlir] [MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (PR #203156)
Jianhui Li
llvmlistbot at llvm.org
Thu Jun 18 15:50:51 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/203156
>From 77c58e0c7500d0347c14dd2a59f9a98dfe307946 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 1 Jun 2026 18:58:11 +0000
Subject: [PATCH 01/42] [mlir][xegpu] Refactor reduction/bitcast/interleave
layout setup
Three related cleanups in the layout-setup helpers:
* setupMultiReductionResultLayout (InstData/Lane): replace the hard-coded
innermost-two-dims pattern with two helpers. leadingDimsAreUnit asserts
that all leading dims really are size-1 (previously assumed silently).
computeReductionLaneLayoutAndData becomes the single source of truth for
lane_layout/lane_data, and inst_data is derived as the elementwise
product. SliceAttr consumers contribute their slice dims to override the
reduction dims. Resulting inst_data layouts now also carry
lane_layout/lane_data.
* setupBitCastResultLayout / setupInterleaveResultLayout: extract the
shared "double the innermost data field until divisible by ratio" logic
into adjustInnermostDimForDivisibility, parameterized by (layoutKind,
innerMostDim, ratio, bound). Both call sites collapse into a few lines
each; bitcast gains an explicit early-return for the same-or-larger
element-type case.
* setupInsertStridedSliceResultLayout: disable the InstData clamping
path; it now stubs out the same way Subgroup does.
Tests in propagate-layout.mlir and propagate-layout-inst-data.mlir
updated to reflect the new richer inst-data layouts (which now carry
lane_layout/lane_data) and the disabled insert_strided_slice clamp.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 314 +++++++++++-------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 10 +
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 2 -
.../XeGPU/propagate-layout-inst-data.mlir | 26 +-
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 18 +-
5 files changed, 234 insertions(+), 136 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 4b821fce3e40c..96c08f9880a72 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -718,6 +718,21 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// Use case 2: Dim split from source to result, for multi-stage reduction
SmallVector<SmallVector<int64_t>> splitDimGroups;
if (xegpu::matchSplitDimExpansion(srcShape, resShape, splitDimGroups)) {
+ llvm::dbgs() << "[DEBUG inferShapeCastSourceLayout] case2 split. resLayout="
+ << resLayout << " srcShape=[";
+ for (auto s : srcShape)
+ llvm::dbgs() << s << ",";
+ llvm::dbgs() << "] resShape=[";
+ for (auto s : resShape)
+ llvm::dbgs() << s << ",";
+ llvm::dbgs() << "] groups=";
+ for (auto &g : splitDimGroups) {
+ llvm::dbgs() << "{";
+ for (auto d : g)
+ llvm::dbgs() << d << ",";
+ llvm::dbgs() << "} ";
+ }
+ llvm::dbgs() << "\n";
auto srcLayout = resLayout;
for (const auto &dimGroup : splitDimGroups)
srcLayout = srcLayout.collapseDims(dimGroup);
@@ -776,6 +791,82 @@ xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
return payloadLayout;
}
+/// Returns true if every dimension of `shape` except the innermost
+/// `numInnerDims` is a unit (size-1) dimension.
+///
+/// Several reduction layout-setup paths (InstData, Lane) only distribute the
+/// innermost one or two dimensions and rely on all the leading dimensions
+/// being degenerate. This helper makes that assumption explicit and checkable
+/// instead of silently leaving leading dimensions undistributed.
+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; });
+}
+
+/// Computes the lane_layout and lane_data for a multi-reduction's source
+/// layout. Only the innermost two dimensions are distributed; all leading
+/// dimensions are assumed to be unit (the caller verifies this via
+/// `leadingDimsAreUnit`).
+///
+/// The layout is chosen to minimize cross-lane reduction: whenever possible a
+/// reduction dimension is reduced *within* a lane (lane_layout == 1, with up to
+/// `maxReduceVectorSize` elements packed into lane_data), and the subgroup's
+/// lanes are spread across a non-reduction dimension instead.
+///
+/// - Exactly one of the innermost two dims is a reduction dim: place
+/// `subgroupSize` lanes on the non-reduction dim and keep lane_layout == 1
+/// on the reduction dim, packing up to `maxReduceVectorSize` reduced
+/// elements into lane_data along that reduction dim.
+/// - Both innermost dims are reduction dims (or the source is rank 1): fall
+/// back to the default of `subgroupSize` lanes on the innermost dim,
+/// packing `maxReduceVectorSize` elements on the second-to-innermost dim.
+///
+/// Returns the (lane_layout, lane_data) pair. The corresponding inst_data is
+/// simply the element-wise product lane_layout * lane_data.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
+ ArrayRef<int64_t> reductionDims,
+ int subgroupSize,
+ int64_t maxReduceVectorSize) {
+ int srcRank = srcShape.size();
+ SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
+ llvm::dbgs() << "[DEBUG computeReductionLaneLayoutAndData] srcRank=" << srcRank
+ << " srcShape.size()=" << srcShape.size() << " srcShape=[";
+ for (size_t i = 0; i < srcShape.size(); ++i) {
+ if (i > 0) llvm::dbgs() << ", ";
+ llvm::dbgs() << srcShape[i];
+ }
+ llvm::dbgs() << "]\n";
+
+ int innermost = srcRank - 1;
+ int secondInnermost = srcRank - 2;
+
+ // `laneDim` carries the subgroupSize lanes; `vectorDim` packs the reduced
+ // elements into lane_data. Default: lanes on the innermost dim, reduced
+ // vector on the second-to-innermost dim.
+ int laneDim = innermost;
+ int vectorDim = secondInnermost; // negative for rank 1
+
+ // If only the innermost dim is reduced, spread the lanes across the
+ // non-reduction (second-to-innermost) dim and reduce the innermost dim
+ // within each lane instead.
+ // if (srcRank >= 2 && isReduction(innermost) &&
+ // !isReduction(secondInnermost)) {
+ // laneDim = secondInnermost;
+ // vectorDim = innermost;
+ // }
+
+ laneLayout[laneDim] =
+ std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
+ if (vectorDim >= 0)
+ laneData[vectorDim] = std::min(maxReduceVectorSize, srcShape[vectorDim]);
+
+ return {laneLayout, laneData};
+}
+
/// Sets up layout for reduction operations by creating a SliceAttr for the
/// result.
///
@@ -792,9 +883,13 @@ xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
/// reuse the slice layout's parent layout for the source to further minimize
/// potential data redistribution.
///
-/// InstData requries {1, ..., min(maxReduceVectorSize, srcShape),subgroupSize}
-/// Lane Layout requires {1, ..., 1, subgroupSize}
-/// Lane data requires {1, ..., min(maxReduceVectorSize, srcShape), 1}
+/// 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.
///
/// Examples:
/// 1. Subgroup layout - Row reduction on 2D tensor:
@@ -922,22 +1017,35 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
(!orderAttr || orderAttr.empty()) ? nullptr : toInt32Attr(order));
}
} else if (layoutKind == xegpu::LayoutKind::InstData) {
-
- SmallVector<int64_t> instData(srcRank, 1);
- if (srcRank >= 2)
- instData[srcRank - 2] =
- std::min(maxReduceVectorSize, srcShape[srcRank - 2]);
- instData[srcRank - 1] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[srcRank - 1]);
- srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(instData));
+ xegpu::SliceAttr consumerSliceLayout =
+ dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ auto reductionDimsOverrideConsumer = consumerSliceLayout?
+ SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef()): reductionDims;
+ auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ srcShape, reductionDimsOverrideConsumer, subgroupSize, maxReduceVectorSize);
+ // 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];
+ srcLayout = xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr,
+ /*inst_data=*/toInt32Attr(instData),
+ /*lane_layout=*/toInt32Attr(laneLayout),
+ /*lane_data=*/toInt32Attr(laneData),
+ /*order=*/nullptr);
} else if (layoutKind == xegpu::LayoutKind::Lane) {
-
- SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
- laneLayout[srcRank - 1] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[srcRank - 1]);
- if (srcRank >= 2)
- laneData[srcRank - 2] =
- std::min(maxReduceVectorSize, srcShape[srcRank - 2]);
+ // 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 reductionDimsOverrideConsumer = consumerSliceLayout?
+ SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef()): reductionDims;
+ auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ srcShape, reductionDimsOverrideConsumer, subgroupSize, maxReduceVectorSize);
srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(laneLayout),
toInt32Attr(laneData));
}
@@ -978,6 +1086,58 @@ xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
return result;
}
+/// Adjusts `consumerLayout`'s innermost-dim data field selected by
+/// `layoutKind` so that the source layout can be safely inferred by dividing
+/// that value by `ratio`. Doubles the value until the divisibility constraint
+/// is met, bounded above by result-shape.
+///
+/// Used by ops whose source relates to the result by a fixed factor along the
+/// innermost dim (e.g., bitcast: bitwidth ratio; interleave: 2x).
+///
+/// Divisibility constraints per LayoutKind:
+/// - Subgroup: sgData[innermost] % ratio == 0
+/// - InstData: instData[innermost] % (laneLayout[innermost] * ratio) == 0
+/// (laneLayout falls back to subgroupSize if absent)
+/// - Lane: laneData[innermost] % ratio == 0
+static xegpu::DistributeLayoutAttr
+adjustInnermostDimForDivisibility(xegpu::DistributeLayoutAttr consumerLayout,
+ xegpu::LayoutKind layoutKind,
+ size_t innerMostDim, int ratio, int64_t bound,
+ const xegpu::uArch::uArch *uArch) {
+ SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+
+ int64_t sgDataValue = -1;
+ int64_t instDataValue = -1;
+ int64_t laneDataValue = -1;
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ sgDataValue = sgData[innerMostDim];
+ while ((sgDataValue <= bound) && (sgDataValue % ratio) != 0)
+ sgDataValue *= 2;
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
+ while ((instDataValue <= bound) &&
+ (instDataValue % (innermostDimLaneLayout * ratio) != 0))
+ instDataValue *= 2;
+ assert((bound % instDataValue) == 0 &&
+ "bound, instData, and laneLayout for innermost must be 2^n!");
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ laneDataValue = laneData[innerMostDim];
+ while ((laneDataValue <= bound) && (laneDataValue % ratio) != 0)
+ laneDataValue *= 2;
+ }
+
+ return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
+ laneDataValue);
+}
+
/// Sets up the result layout for a bitcast operation.
/// When casting to a smaller bitwidth, adjusts the layout dimensions (sgData,
/// instData, or laneData) by multiplying by the bitwidth ratio to ensure the
@@ -1009,53 +1169,23 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
ArrayRef<int64_t> resShape = resVecTy.getShape();
- SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> laneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
"laneData must be available for all dimensions");
+
+ // Casting to same/larger element type: result has fewer (or equal) elements
+ // along the innermost dim, no adjustment needed.
+ if (srcElemTyBitWidth <= resElemTyBitWidth)
+ return consumerLayout;
+
+ // Casting to smaller element type: result has more elements along innermost
+ // dim. Adjust the innermost data field upward so the source layout can be
+ // recovered by dividing by bitWidthRatio.
size_t innerMostDim = srcShape.size() - 1;
- int64_t sgDataValue = -1;
- int64_t instDataValue = -1;
- int64_t laneDataValue = -1;
- if (srcElemTyBitWidth > resElemTyBitWidth) {
- // When casting to a smaller bitwidth, multiply the result layout
- // accordingly to ensure it can be divided by the ratio back to the
- // source layout.
- int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[innerMostDim];
- while ((sgDataValue <= resShape[innerMostDim]) &&
- (sgDataValue % bitWidthRatio) != 0)
- sgDataValue *= 2;
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[innerMostDim];
- const int innermostDimLaneLayout = laneLayout.empty()
- ? uArch->getSubgroupSize()
- : laneLayout[innerMostDim];
- // Adjust instDataValue so it still fits within an instruction after
- // dividing by bitWidthRatio
- while ((instDataValue <= resShape[innerMostDim]) &&
- (instDataValue % (innermostDimLaneLayout * bitWidthRatio) != 0))
- instDataValue *= 2;
- assert((resShape[innerMostDim] % instDataValue) == 0 &&
- "resShape, instData, and lanelayout for innermost must be 2^n !");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[innerMostDim];
- while ((laneDataValue <= resShape[innerMostDim]) &&
- (laneDataValue % bitWidthRatio != 0))
- laneDataValue *= 2;
- }
- // Now set only instData and laneData, preserving sgData
- xegpu::DistributeLayoutAttr resLayout;
- resLayout = consumerLayout.setDimData(innerMostDim, sgDataValue,
- instDataValue, laneDataValue);
- return resLayout;
- }
- return consumerLayout;
+ int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
+ return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
+ innerMostDim, bitWidthRatio,
+ resShape[innerMostDim], uArch);
}
/// Sets up the result layout for an interleave operation to ensure the source
@@ -1077,52 +1207,18 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
- ArrayRef<int64_t> srcShape = srcVecTy.getShape();
- SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> laneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
-
- assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
+ assert(consumerLayout.getRank() == static_cast<int64_t>(resShape.size()) &&
"consumer layout rank must match source shape rank");
- const size_t innerMostDim = srcShape.size() - 1;
- int64_t sgDataValue = -1;
- int64_t instDataValue = -1;
- int64_t laneDataValue = -1;
- // Interleave doubles the innermost dimension (ratio = 2)
+ // Interleave doubles the innermost dimension (ratio = 2). Adjust the
+ // innermost data field so the source layout can be recovered by dividing
+ // by 2.
+ const size_t innerMostDim = resShape.size() - 1;
constexpr int ratio = 2;
-
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[innerMostDim];
- // Ensure sgDataValue is divisible by ratio so source sgData can be inferred
- while ((sgDataValue <= srcShape[innerMostDim]) &&
- (sgDataValue % ratio != 0))
- sgDataValue *= ratio;
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[innerMostDim];
- const int innermostDimLaneLayout = laneLayout.empty()
- ? uArch->getSubgroupSize()
- : laneLayout[innerMostDim];
- // Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
- // ratio) when inferring the source layout
- while ((instDataValue <= srcShape[innerMostDim]) &&
- (instDataValue % (innermostDimLaneLayout * ratio) != 0))
- instDataValue *= ratio;
- assert((srcShape[innerMostDim] % instDataValue) == 0 &&
- "srcShape, instData, and laneLayout for innermost must be 2^n!");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[innerMostDim];
- // Ensure laneDataValue is at least 2 and divisible by ratio
- // so that source laneData = laneDataValue/2 is valid
- while ((laneDataValue <= srcShape[innerMostDim]) &&
- (laneDataValue % ratio != 0))
- laneDataValue *= ratio;
- }
-
- return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
- laneDataValue);
+ return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
+ innerMostDim, ratio,
+ resShape[innerMostDim], uArch);
}
/// Sets up the result layout for an insert strided slice operation.
@@ -1141,21 +1237,15 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
SmallVector<int64_t> consumerLaneLayout =
consumerLayout.getEffectiveLaneLayoutAsInt();
ArrayRef<int64_t> srcShape = srcVectorTy.getShape();
- int64_t instDataValue = -1;
int64_t laneDataValue = -1;
requiredResLayout = consumerLayout;
int srcRank = srcShape.size();
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ if (layoutKind == xegpu::LayoutKind::Subgroup ||
+ layoutKind == xegpu::LayoutKind::InstData) {
assert(true &&
"subgroup layout assignment not supported for insertStridedSlice.");
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- for (int dim = 0; dim < srcRank; dim++) {
- instDataValue = std::min(srcShape[dim], consumerInstData[dim]);
- requiredResLayout =
- requiredResLayout.setDimData(dim, -1, instDataValue, -1);
- }
} else if (layoutKind == xegpu::LayoutKind::Lane) {
for (int dim = 0; dim < srcRank; dim++) {
assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 6a37ae6502b2d..9d7b99c1e95a7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -684,12 +684,18 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
auto requiredResLayoutAttr = xegpu::setupMultiReductionResultLayout(
layoutKind, sourceTy, consumerLayoutAttr, reductionDims, numSg, uArch);
+ llvm::dbgs() << "[DEBUG visitMultiRed] op=" << *reduction
+ << " consumer=" << consumerLayoutAttr
+ << " required=" << requiredResLayoutAttr << "\n";
+
xegpu::setTemporaryLayout(reduction->getResult(0), requiredResLayoutAttr);
// derive the source layout from the dominant layout and reduction dims
auto srcLayoutAttr = xegpu::inferMultiReductionSourceLayout(
requiredResLayoutAttr, reductionDims);
+ llvm::dbgs() << "[DEBUG visitMultiRed] srcLayout=" << srcLayoutAttr << "\n";
+
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
// Accumulator should have the same layout as the result.
propagateIfChanged(operands[1],
@@ -758,6 +764,10 @@ void LayoutInfoPropagation::visitShapeCastOp(
xegpu::DistributeLayoutAttr srcLayoutAttr =
xegpu::inferShapeCastSourceLayout(resultLayoutAttr, resShape, srcShape);
+ llvm::dbgs() << "[DEBUG visitShapeCast] op=" << *shapeCast
+ << " consumer=" << resultLayoutAttr
+ << " src=" << srcLayoutAttr << "\n";
+
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index fc7c3b170dd3b..55bfc317f89fe 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -969,8 +969,6 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
if (valType.isIntOrFloat()) {
rewriter.replaceOp(op, op.getSource());
- assert(!inputLayout.dropInstData() && !targetLayout.dropInstData() &&
- "unexpected layout attributes for scalar type");
return success();
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 5f493c8ca0df6..98703c1beb9d3 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -220,9 +220,9 @@ func.func @scatter_ops_chunksize_slice(%src: memref<1024xf32>) {
gpu.module @test {
// CHECK-LABEL: func.func @insert_strided_slice_inst_data_no_packing(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf32>) {
-// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>} dense<1.000000e+00> : vector<4x16xf32>
-// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>} dense<0.000000e+00> : vector<8x32xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>, offsets = [0, 0], strides = [1, 1]} : vector<4x16xf32> into vector<8x32xf32>
+// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<1.000000e+00> : vector<4x16xf32>
+// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<8x32xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>, offsets = [0, 0], strides = [1, 1]} : vector<4x16xf32> into vector<8x32xf32>
// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
// CHECK: xegpu.store_nd %[[INSERT]], %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
func.func @insert_strided_slice_inst_data_no_packing(%arg0: memref<8x32xf32>) {
@@ -240,9 +240,9 @@ func.func @insert_strided_slice_inst_data_no_packing(%arg0: memref<8x32xf32>) {
gpu.module @test {
// CHECK-LABEL: func.func @insert_strided_slice_inst_data_with_packing(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x64xi8>) {
-// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 64]>} dense<1> : vector<4x64xi8>
-// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [4, 64]>} dense<0> : vector<8x64xi8>
-// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [4, 64]>, offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
+// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 64]>} dense<1> : vector<4x64xi8>
+// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 64]>} dense<0> : vector<8x64xi8>
+// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [8, 64]>, offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
%c0 = arith.constant 0 : index
%cst_small = arith.constant dense<1> : vector<4x64xi8>
@@ -258,11 +258,11 @@ func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
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]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 16]>} : vector<1024xf16> to vector<8x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
-// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16]>, 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]>, 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]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
+// 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>
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>
@@ -282,7 +282,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]>, 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, 1, 1], lane_layout = [1, 1, 1], lane_data = [1, 1, 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>
@@ -300,7 +300,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]>, 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, 1, 4], lane_layout = [1, 1, 4], lane_data = [1, 1, 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>
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index bad956d45d186..14579606bcbed 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -656,20 +656,20 @@ 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<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
+// 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<8x8x16xf16> to vector<8x16xf16>
+// 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>
-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>
- %1 = xegpu.load %arg0[%0], %cst : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
- %2 = vector.shape_cast %1 : vector<1024xf16> to vector<8x8x16xf16>
+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>
+ %1 = xegpu.load %arg0[%0], %cst : memref<128xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
+ %2 = vector.shape_cast %1 : vector<128xf16> to vector<1x8x16xf16>
%cst_0 = arith.constant dense<0.000000e+00> : vector<8x16xf16>
%cst_1 = arith.constant dense<0.000000e+00> : vector<16xf16>
- %3 = vector.multi_reduction <add>, %2, %cst_0 [0] : vector<8x8x16xf16> to vector<8x16xf16>
+ %3 = vector.multi_reduction <add>, %2, %cst_0 [0] : vector<1x8x16xf16> to vector<8x16xf16>
%4 = vector.multi_reduction <add>, %3, %cst_1 [0] : vector<8x16xf16> to vector<16xf16>
%cst_2 = arith.constant dense<true> : vector<16xi1>
%cst_3 = arith.constant dense<1> : vector<16xindex>
>From ece64f9935eefd0b6dd3aad38ebd04b824c55f0b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 1 Jun 2026 22:38:31 +0000
Subject: [PATCH 02/42] remove isChunkedStore/Load handling in the
setUpGenericLoad/store rule
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 3 +-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 127 ++++++++----------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 10 +-
.../XeGPU/propagate-layout-inst-data.mlir | 107 +++++----------
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 61 +++------
5 files changed, 118 insertions(+), 190 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 94d1d5aecbe60..253dcb18af67c 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -216,7 +216,7 @@ setupLoadGatherAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
/// Sets up the anchor layout for load matrix operation.
DistributeLayoutAttr
setupLoadMatrixAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
- DistributeLayoutAttr consumerLayout,
+ int chunkSize, DistributeLayoutAttr consumerLayout,
const uArch::uArch *uArch);
/// Sets up the anchor layout for a store scatter operation.
@@ -228,6 +228,7 @@ DistributeLayoutAttr setupStoreScatterAnchorLayout(LayoutKind layoutKind,
/// Sets up the anchor layout for a store matrix operation.
DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
+ int chunkSize,
const uArch::uArch *uArch);
/// Sets up the anchor layouts for a dpas operands (A, B, and C/D).
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 96c08f9880a72..427515ea75b83 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -833,10 +833,12 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
int64_t maxReduceVectorSize) {
int srcRank = srcShape.size();
SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
- llvm::dbgs() << "[DEBUG computeReductionLaneLayoutAndData] srcRank=" << srcRank
- << " srcShape.size()=" << srcShape.size() << " srcShape=[";
+ llvm::dbgs() << "[DEBUG computeReductionLaneLayoutAndData] srcRank="
+ << srcRank << " srcShape.size()=" << srcShape.size()
+ << " srcShape=[";
for (size_t i = 0; i < srcShape.size(); ++i) {
- if (i > 0) llvm::dbgs() << ", ";
+ if (i > 0)
+ llvm::dbgs() << ", ";
llvm::dbgs() << srcShape[i];
}
llvm::dbgs() << "]\n";
@@ -1019,10 +1021,13 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
} else if (layoutKind == xegpu::LayoutKind::InstData) {
xegpu::SliceAttr consumerSliceLayout =
dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- auto reductionDimsOverrideConsumer = consumerSliceLayout?
- SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef()): reductionDims;
+ auto reductionDimsOverrideConsumer =
+ consumerSliceLayout
+ ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
+ : reductionDims;
auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDimsOverrideConsumer, subgroupSize, maxReduceVectorSize);
+ srcShape, reductionDimsOverrideConsumer, subgroupSize,
+ maxReduceVectorSize);
// inst_data is the per-instruction data, i.e. the element-wise product of
// lane_layout and lane_data.
SmallVector<int64_t> instData(srcRank);
@@ -1042,10 +1047,13 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
"dimensions are unit dimensions");
xegpu::SliceAttr consumerSliceLayout =
dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- auto reductionDimsOverrideConsumer = consumerSliceLayout?
- SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef()): reductionDims;
+ auto reductionDimsOverrideConsumer =
+ consumerSliceLayout
+ ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
+ : reductionDims;
auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDimsOverrideConsumer, subgroupSize, maxReduceVectorSize);
+ srcShape, reductionDimsOverrideConsumer, subgroupSize,
+ maxReduceVectorSize);
srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(laneLayout),
toInt32Attr(laneData));
}
@@ -1273,8 +1281,8 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
/// lane_data={1,min(consumer, maxLaneLoadSize)}
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
- xegpu::DistributeLayoutAttr consumerLayout, bool isChunkedLoad,
- int maxChunkSize, ArrayRef<int64_t> resShape, int subgroupSize) {
+ xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
+ ArrayRef<int64_t> resShape, int subgroupSize) {
if (layoutKind == xegpu::LayoutKind::Subgroup)
return consumerLayout;
@@ -1288,31 +1296,16 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
SmallVector<int> laneLayout(resShape.size(), 1);
SmallVector<int> laneData(resShape.size(), 1);
- if (!isChunkedLoad) {
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData.back() = std::min(static_cast<int>(consumerInstData.back()),
- maxChunkSize * subgroupSize);
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneData.back() =
- std::min(static_cast<int>(consumerLaneData.back()), maxChunkSize);
- laneLayout.back() = std::min(static_cast<int64_t>(subgroupSize),
- resShape.back() / laneData.back());
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
- } else {
- assert(resShape.size() == 2 && "Chunked Store must access 2D tensor tile.");
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData[0] = subgroupSize;
- instData[1] =
- std::min(static_cast<int>(consumerInstData[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[0] = subgroupSize;
- laneData[1] =
- std::min(static_cast<int>(consumerLaneData[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ instData.back() = std::min(static_cast<int>(consumerInstData.back()),
+ maxChunkSize * subgroupSize);
+ return xegpu::LayoutAttr::get(context, instData);
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ laneData.back() =
+ std::min(static_cast<int>(consumerLaneData.back()), maxChunkSize);
+ laneLayout.back() = std::min(static_cast<int64_t>(subgroupSize),
+ resShape.back() / laneData.back());
+ return xegpu::LayoutAttr::get(context, laneLayout, laneData);
}
return nullptr;
}
@@ -1333,15 +1326,14 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
- (chunkSize > 1), maxChunkSize, resShape,
- subgroupSize);
+ maxChunkSize, resShape, subgroupSize);
}
/// Sets up the anchor layout for load matrix operation.
/// TODO: enhance load matrix to indicate lowering to chunked load or not.
xegpu::DistributeLayoutAttr
xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType resVecTy,
+ VectorType resVecTy, int chunkSize,
xegpu::DistributeLayoutAttr consumerLayout,
const xegpu::uArch::uArch *uArch) {
@@ -1353,10 +1345,10 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneLoadSize(elemBitWidth), chunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
- false, maxChunkSize, resShape,
- subgroupSize);
+ maxChunkSize, resShape, subgroupSize);
}
/// Sets up the anchor layout for store scatter and store matrix operation.
@@ -1372,9 +1364,8 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
/// lane_data={1,min(srcVec, maxLaneStoreSize)}
static xegpu::DistributeLayoutAttr
setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
- mlir::MLIRContext *context, bool isChunkedStore,
- int maxChunkSize, ArrayRef<int64_t> srcShape,
- int subgroupSize) {
+ mlir::MLIRContext *context, int maxChunkSize,
+ ArrayRef<int64_t> srcShape, int subgroupSize) {
int srcShapeSize = srcShape.size();
SmallVector<int> instData(srcShapeSize, 1);
@@ -1386,28 +1377,16 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
"subgroup layout assignment not supported for storeScatter.");
return nullptr;
}
-
- if (!isChunkedStore) {
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
- } else {
- assert(srcShapeSize == 2 && "Chunked Store must access 2D tensor tile.");
- if (layoutKind == xegpu::LayoutKind::InstData) {
- instData[0] = subgroupSize;
- instData[1] = std::min(static_cast<int>(srcShape[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[0] = subgroupSize;
- laneData[1] = std::min(static_cast<int>(srcShape[1]), maxChunkSize);
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
- }
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ laneLayout[srcShapeSize - 1] =
+ std::min(subgroupSize, static_cast<int>(srcShape.back()));
+ laneData[srcShapeSize - 1] =
+ std::min(maxChunkSize, static_cast<int>(srcShape.back()));
+ return xegpu::LayoutAttr::get(context, instData);
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ laneLayout[srcShapeSize - 1] =
+ std::min(subgroupSize, static_cast<int>(srcShape.back()));
+ return xegpu::LayoutAttr::get(context, laneLayout, laneData);
}
return nullptr;
}
@@ -1426,15 +1405,16 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
- return setupGenericStoreAnchorLayout(layoutKind, context, (chunkSize > 1),
- maxChunkSize, srcShape, subgroupSize);
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneStoreSize(elemBitWidth), chunkSize);
+ return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
+ srcShape, subgroupSize);
}
/// Sets up the anchor layout for a store matrix operation.
xegpu::DistributeLayoutAttr
xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy,
+ VectorType srcVecTy, int chunkSize,
const xegpu::uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1445,9 +1425,10 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+ int maxChunkSize =
+ std::min(uArchInstruction->getMaxLaneStoreSize(elemBitWidth), chunkSize);
- return setupGenericStoreAnchorLayout(layoutKind, context, false, maxChunkSize,
+ return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize);
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 9d7b99c1e95a7..c1ad3c4ee8e2e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1341,8 +1341,10 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
const uArch *uArch = getUArch(getChipStr(loadMatrixOp).value_or(""));
if (!uArch)
return;
+ int chunkSize =
+ 1; // placeHolder for future use when LoadMatrix supports coalescing
auto requiredAnchorLayoutAttr = xegpu::setupLoadMatrixAnchorLayout(
- layoutKind, resVecTy, consumerLayoutAttr, uArch);
+ layoutKind, resVecTy, chunkSize, consumerLayoutAttr, uArch);
loadMatrixOp.setLayoutAttr(requiredAnchorLayoutAttr);
}
}
@@ -1360,8 +1362,10 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
if (!uArch)
return;
- auto requiredAnchorLayoutAttr =
- xegpu::setupStoreMatrixAnchorLayout(layoutKind, srcVecTy, uArch);
+ int chunkSize =
+ 1; // placeHolder for future use when StoreMatrix supports coalescing
+ auto requiredAnchorLayoutAttr = xegpu::setupStoreMatrixAnchorLayout(
+ layoutKind, srcVecTy, chunkSize, uArch);
storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
layout = LayoutInfo(requiredAnchorLayoutAttr);
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 98703c1beb9d3..3ea06d52690dc 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -115,26 +115,6 @@ gpu.module @test_kernel {
}
}
-// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 8 : i64, layout = #xegpu.layout<inst_data = [16, 8]>}>
-// CHECK-SAME: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 8 : i64, layout = #xegpu.layout<inst_data = [16, 8]>}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize(%src: memref<256xf16>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}>
- : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
// -----
gpu.module @test {
// CHECK-LABEL: func.func @store_matrix(
@@ -148,45 +128,24 @@ func.func @store_matrix(%arg0: !xegpu.mem_desc<16x64xf16>) {
}
}
-// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize_excessive(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<1024xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize_excessive(%src: memref<1024xf32>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=32}>
- : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=32}>
- : vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize_excessive_anchor(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<1024xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 32 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize_excessive_anchor(%src: memref<1024xf32>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=32}>
- : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16x32xf32>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=32, layout = #xegpu.layout<inst_data = [16, 16]>}>
- : vector<16x32xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
+// 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 = [16]>} dense<true> : vector<16x32xi1>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16x32xindex>
+// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16]>}> :
+// CHECK-SAME: memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16]>}> :
+// CHECK-SAME: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+func.func @scatter_ops_coalesce_chunksize(%src: memref<512xf32>) {
+ %1 = arith.constant dense<1>: vector<16x32xi1>
+ %offset = arith.constant dense<12> : vector<16x32xindex>
+ %3 = xegpu.load %src[%offset], %1
+ : memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+ xegpu.store %3, %src[%offset], %1 <{layout = #xegpu.layout<inst_data = [16, 16]>}>
+ : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
return
}
}
@@ -194,24 +153,24 @@ func.func @scatter_ops_chunksize_excessive_anchor(%src: memref<1024xf32>) {
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize_slice(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<1024xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<true> : vector<16xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16xindex>
-// CHECK: %[[LOADED:.*]] = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16]>}> :
-// CHECK-SAME: memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[LOADED]] {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} : vector<16xf32> to vector<16x16xf32>
-// CHECK: xegpu.store %[[BCAST]], %[[ARG0]][%{{.*}}], %{{.*}} <{chunk_size = 16 : i64, layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: vector<16x16xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize_slice(%src: memref<1024xf32>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1
- : memref<1024xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
-
- %4 = vector.broadcast %3 : vector<16xf32> to vector<16x16xf32>
- xegpu.store %4, %src[%offset], %1 <{chunk_size=16, layout = #xegpu.layout<inst_data = [16, 16]>}>
- : vector<16x16xf32>, memref<1024xf32>, vector<16xindex>, vector<16xi1>
+// 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], lane_data = [1]>}
+// CHECK-SAME: dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
+// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
+// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16xindex>, vector<16xi1> -> 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>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %offset = arith.constant dense<0> : vector<16x16xindex>
+ %mask = arith.constant dense<true> : vector<16x16xi1>
+ %3 = xegpu.load %arg1[%offset], %mask
+ : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+ %4 = vector.transpose %3, [1, 0] : vector<16x16xf16> to vector<16x16xf16>
+ %5 = xegpu.dpas %1, %4 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
+ %6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+ xegpu.store_nd %5, %6[0, 0] : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
return
}
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 14579606bcbed..c04e949015368 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -133,20 +133,20 @@ func.func @extf_truncf(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @load_gather_with_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-LABEL: func.func @load_gather_with_coalesce_chunksize(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x16xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
// CHECK-SAME: dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{chunk_size = 16 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
-func.func @load_gather_with_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
+// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<16x16xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
+func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %offset = arith.constant dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
- %mask = arith.constant dense<true> : vector<16xi1>
- %3 = xegpu.load %arg1[%offset], %mask <{chunk_size=16}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
+ %offset = arith.constant dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16x16xindex>
+ %mask = arith.constant dense<true> : vector<16x16xi1>
+ %3 = xegpu.load %arg1[%offset], %mask
+ : memref<16x16xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
%4 = vector.transpose %3, [1, 0] : vector<16x16xf16> to vector<16x16xf16>
%5 = xegpu.dpas %1, %4 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
%6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
@@ -157,37 +157,20 @@ func.func @load_gather_with_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @store_scatter_with_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<128xf32>) {
-// CHECK-NEXT: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>} dense<1.000000e+00> : vector<16x8xf32>
-// CHECK-NEXT: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK-NEXT: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
-// CHECK-NEXT: xegpu.store %[[CST]], %[[ARG0]][%[[CST_1]]], %[[CST_0]] <{chunk_size = 8 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}> : vector<16x8xf32>, memref<128xf32>, vector<16xindex>, vector<16xi1>
-func.func @store_scatter_with_chunksize(%arg0: memref<128xf32>) {
- %val = arith.constant dense<1.000000e+00> : vector<16x8xf32>
- %mask = arith.constant dense<true> : vector<16xi1>
- %offset = arith.constant dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
- xegpu.store %val, %arg0[%offset], %mask <{chunk_size = 8}>: vector<16x8xf32>, memref<128xf32>, vector<16xindex>, vector<16xi1>
- return
-}
-}
-
-// -----
-gpu.module @test {
-// CHECK-LABEL: func.func @scatter_ops_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
-// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{chunk_size = 8 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}>
-// CHECK-SAME: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
-// CHECK: xegpu.store %[[LOAD_VEC]], %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{chunk_size = 8 : i64, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
-func.func @scatter_ops_chunksize(%src: memref<256xf16>) {
- %1 = arith.constant dense<1>: vector<16xi1>
- %offset = arith.constant dense<12> : vector<16xindex>
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}>
- : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}>
- : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
+// CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x8xf16>) {
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16x8xi1>
+// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16x8xindex>
+// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}>
+// CHECK-SAME: memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
+// CHECK: xegpu.store %[[LOAD_VEC]], %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}> : vector<16x8xf16>, memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1>
+func.func @scatter_ops_coalesce_chunksize(%src: memref<16x8xf16>) {
+ %1 = arith.constant dense<1>: vector<16x8xi1>
+ %offset = arith.constant dense<12> : vector<16x8xindex>
+ %3 = xegpu.load %src[%offset], %1
+ : memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
+ xegpu.store %3, %src[%offset], %1
+ : vector<16x8xf16>, memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1>
return
}
}
>From 5c19d85d65056ac2d6ab94005a2f2d2129ca056f Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 2 Jun 2026 04:51:43 +0000
Subject: [PATCH 03/42] refactor generic load/store setup rule
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 155 ++++++++++--------
1 file changed, 84 insertions(+), 71 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 427515ea75b83..d461caebeb716 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -833,15 +833,6 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
int64_t maxReduceVectorSize) {
int srcRank = srcShape.size();
SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
- llvm::dbgs() << "[DEBUG computeReductionLaneLayoutAndData] srcRank="
- << srcRank << " srcShape.size()=" << srcShape.size()
- << " srcShape=[";
- for (size_t i = 0; i < srcShape.size(); ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << srcShape[i];
- }
- llvm::dbgs() << "]\n";
int innermost = srcRank - 1;
int secondInnermost = srcRank - 2;
@@ -852,15 +843,6 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
int laneDim = innermost;
int vectorDim = secondInnermost; // negative for rank 1
- // If only the innermost dim is reduced, spread the lanes across the
- // non-reduction (second-to-innermost) dim and reduce the innermost dim
- // within each lane instead.
- // if (srcRank >= 2 && isReduction(innermost) &&
- // !isReduction(secondInnermost)) {
- // laneDim = secondInnermost;
- // vectorDim = innermost;
- // }
-
laneLayout[laneDim] =
std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
if (vectorDim >= 0)
@@ -1097,7 +1079,7 @@ xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
/// Adjusts `consumerLayout`'s innermost-dim data field selected by
/// `layoutKind` so that the source layout can be safely inferred by dividing
/// that value by `ratio`. Doubles the value until the divisibility constraint
-/// is met, bounded above by result-shape.
+/// is met, bounded above by `bound` like result-shape.
///
/// Used by ops whose source relates to the result by a fixed factor along the
/// innermost dim (e.g., bitcast: bitwidth ratio; interleave: 2x).
@@ -1267,18 +1249,36 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
return requiredResLayout;
}
+/// Computes lane_layout and lane_data for scatter-style load anchor layouts
+/// (load gather, load matrix). Lanes and the per-lane vector both live on the
+/// innermost dim; the per-lane vector width is hinted by the consumer's
+/// lane_data[innermost]:
+/// - laneData[innermost] = min(consumerLaneData[innermost], maxChunkSize)
+/// - laneLayout[innermost] = min(subgroupSize,
+/// resShape[innermost] / laneData[innermost])
+/// All other entries are 1.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeScatterLoadLaneLayoutAndData(ArrayRef<int64_t> resShape,
+ ArrayRef<int64_t> consumerLaneData,
+ int subgroupSize, int64_t maxChunkSize) {
+ int rank = resShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int innermost = rank - 1;
+ laneData[innermost] = std::min(consumerLaneData.back(), maxChunkSize);
+ laneLayout[innermost] = std::min(static_cast<int64_t>(subgroupSize),
+ resShape[innermost] / laneData[innermost]);
+ return {laneLayout, laneData};
+}
+
/// Sets up the anchor layout for load gather and load matrix operation.
/// load matrix lowers to load gather and 1d block load. All of them share the
/// same layout setup logic.
/// For Subgroup layout, uses the consumer layout directly.
-/// non-chunked loads (1D or 2D):
-/// InstData = {1, ..., min(consumer, maxLaneLoadSize * subgroupSize)}
-/// LaneLayout = {1, ..., subgroupSize}
-/// lane_data = {1, ..., min(consumer, maxLaneLoadSize)}
-/// chunked loads (2D only):
-/// InstData = {subgroupSize, min(consumer, maxLaneLoadSize)}
-/// LaneLayout = {subgroupSize, 1}
-/// lane_data={1,min(consumer, maxLaneLoadSize)}
+/// For InstData layout, the innermost inst_data is taken directly from the
+/// consumer's inst_data, capped by `maxChunkSize * subgroupSize`.
+/// For Lane layout, lane_layout/lane_data are derived via
+/// `computeScatterLoadLaneLayoutAndData` using the consumer's
+/// lane_data[innermost] as the per-lane vector hint.
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
@@ -1292,20 +1292,18 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int> instData(resShape.size(), 1);
- SmallVector<int> laneLayout(resShape.size(), 1);
- SmallVector<int> laneData(resShape.size(), 1);
-
if (layoutKind == xegpu::LayoutKind::InstData) {
+ SmallVector<int> instData(resShape.size(), 1);
instData.back() = std::min(static_cast<int>(consumerInstData.back()),
maxChunkSize * subgroupSize);
return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneData.back() =
- std::min(static_cast<int>(consumerLaneData.back()), maxChunkSize);
- laneLayout.back() = std::min(static_cast<int64_t>(subgroupSize),
- resShape.back() / laneData.back());
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
+ }
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ auto [laneLayout, laneData] = computeScatterLoadLaneLayoutAndData(
+ resShape, consumerLaneData, subgroupSize, maxChunkSize);
+ SmallVector<int> laneLayout32(laneLayout.begin(), laneLayout.end());
+ SmallVector<int> laneData32(laneData.begin(), laneData.end());
+ return xegpu::LayoutAttr::get(context, laneLayout32, laneData32);
}
return nullptr;
}
@@ -1351,42 +1349,57 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
maxChunkSize, resShape, subgroupSize);
}
+/// Computes lane_layout and lane_data for scatter-style store anchor layouts
+/// (store scatter, store matrix). Lanes and the per-lane vector both live on
+/// the innermost dim:
+/// - laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
+/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
+/// maxChunkSize)
+/// All other entries are 1.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeScatterStoreLaneLayoutAndData(ArrayRef<int64_t> srcShape,
+ int subgroupSize, int64_t maxChunkSize) {
+ int rank = srcShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int innermost = rank - 1;
+ laneLayout[innermost] =
+ std::min(static_cast<int64_t>(subgroupSize), srcShape[innermost]);
+ laneData[innermost] =
+ std::min(srcShape[innermost] / laneLayout[innermost], maxChunkSize);
+ return {laneLayout, laneData};
+}
+
/// Sets up the anchor layout for store scatter and store matrix operation.
-/// store matrix lowers to store scatter and 1d block store. All of them share
-/// the same layout setup logic. For Subgroup layout, not supported yet.
-/// non-chunked stores (1D or 2D):
-/// InstData = {1, ..., subgroupSize}
-/// LaneLayout = {1, ..., subgroupSize}
-/// lane_data = {1, ..., 1}
-/// chunked stores (2D only):
-/// InstData = {subgroupSize, min(srcVec, maxLaneStoreSize)}
-/// LaneLayout = {subgroupSize, 1}
-/// lane_data={1,min(srcVec, maxLaneStoreSize)}
+/// store matrix lowers to store scatter and 1d block store. All of them
+/// share the same layout setup logic. For Subgroup layout, not supported
+/// yet.
+///
+/// Lane layout is derived first via `computeScatterStoreLaneLayoutAndData`;
+/// inst_data is then the element-wise product lane_layout * lane_data.
static xegpu::DistributeLayoutAttr
setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
mlir::MLIRContext *context, int maxChunkSize,
ArrayRef<int64_t> srcShape, int subgroupSize) {
- int srcShapeSize = srcShape.size();
- SmallVector<int> instData(srcShapeSize, 1);
- SmallVector<int> laneLayout(srcShapeSize, 1);
- SmallVector<int> laneData(srcShapeSize, 1);
-
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(true &&
"subgroup layout assignment not supported for storeScatter.");
return nullptr;
}
+
+ auto [laneLayout, laneData] = computeScatterStoreLaneLayoutAndData(
+ srcShape, subgroupSize, maxChunkSize);
+
if (layoutKind == xegpu::LayoutKind::InstData) {
- laneLayout[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
- laneData[srcShapeSize - 1] =
- std::min(maxChunkSize, static_cast<int>(srcShape.back()));
+ SmallVector<int> instData(srcShape.size());
+ for (size_t i = 0; i < srcShape.size(); ++i)
+ instData[i] = static_cast<int>(laneLayout[i] * laneData[i]);
return xegpu::LayoutAttr::get(context, instData);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneLayout[srcShapeSize - 1] =
- std::min(subgroupSize, static_cast<int>(srcShape.back()));
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
+ }
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ SmallVector<int> laneLayout32(laneLayout.begin(), laneLayout.end());
+ SmallVector<int> laneData32(laneData.begin(), laneData.end());
+ return xegpu::LayoutAttr::get(context, laneLayout32, laneData32);
}
return nullptr;
}
@@ -1557,8 +1570,8 @@ getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
return std::make_tuple(instDataA, instDataB, instDataCD);
}
-/// Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.
-/// Returns the three layouts if successful, nullopt otherwise.
+/// Helper function to set up subgroup layouts for DPAS operands A, B, and
+/// C/D. Returns the three layouts if successful, nullopt otherwise.
static std::optional<
std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
xegpu::DistributeLayoutAttr>>
@@ -1688,9 +1701,9 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
}
/// Helper to create a scale layout derived from a matrix operand layout.
-/// The scale layout is computed by mapping each dimension of the matrix layout
-/// to the corresponding scale tensor dimension using the ratio between the
-/// matrix and scale shapes.
+/// The scale layout is computed by mapping each dimension of the matrix
+/// layout to the corresponding scale tensor dimension using the ratio
+/// between the matrix and scale shapes.
static xegpu::DistributeLayoutAttr
createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
VectorType scaleTy, xegpu::DistributeLayoutAttr matrixLayout,
@@ -1781,8 +1794,8 @@ createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
}
/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
-/// B_scale). The numSg and consumerLayout (optional) are only used by sg layout
-/// creation.
+/// B_scale). The numSg and consumerLayout (optional) are only used by sg
+/// layout creation.
std::optional<
std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
@@ -1908,8 +1921,8 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
shapeCast.getSourceVectorType().getShape());
}
- // For vector::InsertStridedSliceOp, infer source layout from result layout.
- // Dest vector must have the same layout as the result.
+ // For vector::InsertStridedSliceOp, infer source layout from result
+ // layout. Dest vector must have the same layout as the result.
if (auto insertSlice = dyn_cast<vector::InsertStridedSliceOp>(op)) {
if (idx == 0) {
return xegpu::inferInsertStridedSliceSourceLayout(
@@ -1978,8 +1991,8 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
if (dyn_cast<vector::ExtractStridedSliceOp>(op))
return resLayout;
- // For elementwise operations, all operands must have the same layout as the
- // result.
+ // For elementwise operations, all operands must have the same layout as
+ // the result.
if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)
return resLayout;
>From edd49273fe0151b505cb74fe5c0139dec7fcfc00 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 2 Jun 2026 06:41:09 +0000
Subject: [PATCH 04/42] add lane layout to dpas, refactor load rules to be more
generic. still missing nd ops andcomplete if user only set inst_data.
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 209 ++++++++++++------
1 file changed, 137 insertions(+), 72 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d461caebeb716..fc2aad83b92a6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -806,6 +806,41 @@ static bool leadingDimsAreUnit(ArrayRef<int64_t> shape, int numInnerDims) {
[](int64_t dim) { return dim == 1; });
}
+/// Builds a LayoutAttr carrying inst_data, lane_layout, and lane_data (no
+/// sg_layout / sg_data / order). Used by InstData-kind setup paths so the
+/// result layout can later be distributed without re-deriving the lane
+/// layout. `instData`, `laneLayout`, and `laneData` may have different
+/// element types; they are normalized to int32 entries.
+static xegpu::LayoutAttr buildInstDataLayoutWithLane(
+ mlir::MLIRContext *context, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr,
+ /*inst_data=*/toI32Attr(instData),
+ /*lane_layout=*/toI32Attr(laneLayout),
+ /*lane_data=*/toI32Attr(laneData),
+ /*order=*/nullptr);
+}
+
+static xegpu::LayoutAttr buildLaneLayout(mlir::MLIRContext *context,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr,
+ /*inst_data=*/nullptr,
+ /*lane_layout=*/toI32Attr(laneLayout),
+ /*lane_data=*/toI32Attr(laneData),
+ /*order=*/nullptr);
+}
+
/// Computes the lane_layout and lane_data for a multi-reduction's source
/// layout. Only the innermost two dimensions are distributed; all leading
/// dimensions are assumed to be unit (the caller verifies this via
@@ -1015,12 +1050,8 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
SmallVector<int64_t> instData(srcRank);
for (int i = 0; i < srcRank; i++)
instData[i] = laneLayout[i] * laneData[i];
- srcLayout = xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
- /*sg_data=*/nullptr,
- /*inst_data=*/toInt32Attr(instData),
- /*lane_layout=*/toInt32Attr(laneLayout),
- /*lane_data=*/toInt32Attr(laneData),
- /*order=*/nullptr);
+ 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.
@@ -1249,26 +1280,28 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
return requiredResLayout;
}
-/// Computes lane_layout and lane_data for scatter-style load anchor layouts
-/// (load gather, load matrix). Lanes and the per-lane vector both live on the
-/// innermost dim; the per-lane vector width is hinted by the consumer's
-/// lane_data[innermost]:
-/// - laneData[innermost] = min(consumerLaneData[innermost], maxChunkSize)
-/// - laneLayout[innermost] = min(subgroupSize,
-/// resShape[innermost] / laneData[innermost])
-/// All other entries are 1.
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-computeScatterLoadLaneLayoutAndData(ArrayRef<int64_t> resShape,
- ArrayRef<int64_t> consumerLaneData,
- int subgroupSize, int64_t maxChunkSize) {
- int rank = resShape.size();
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- int innermost = rank - 1;
- laneData[innermost] = std::min(consumerLaneData.back(), maxChunkSize);
- laneLayout[innermost] = std::min(static_cast<int64_t>(subgroupSize),
- resShape[innermost] / laneData[innermost]);
- return {laneLayout, laneData};
-}
+// /// Computes lane_layout and lane_data for scatter-style load anchor layouts
+// /// (load gather, load matrix). Lanes and the per-lane vector both live on
+// the
+// /// innermost dim; the per-lane vector width is hinted by the consumer's
+// /// lane_data[innermost]:
+// /// - laneData[innermost] = min(consumerLaneData[innermost],
+// maxChunkSize)
+// /// - laneLayout[innermost] = min(subgroupSize,
+// /// resShape[innermost] /
+// laneData[innermost])
+// /// All other entries are 1.
+// static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+// computeScatterLoadLaneLayoutAndData(ArrayRef<int64_t> resShape,
+// ArrayRef<int64_t> consumerLaneData,
+// int subgroupSize, int64_t maxChunkSize) {
+// int rank = resShape.size();
+// SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+// int innermost = rank - 1;
+// laneData[innermost] = std::min(consumerLaneData.back(), maxChunkSize);
+// laneLayout[innermost] = consumerLaneLayout.back();
+// return {laneLayout, laneData};
+// }
/// Sets up the anchor layout for load gather and load matrix operation.
/// load matrix lowers to load gather and 1d block load. All of them share the
@@ -1291,19 +1324,25 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneLayout(resShape.size(), 1);
+ SmallVector<int64_t> laneData(resShape.size(), 1);
if (layoutKind == xegpu::LayoutKind::InstData) {
- SmallVector<int> instData(resShape.size(), 1);
- instData.back() = std::min(static_cast<int>(consumerInstData.back()),
- maxChunkSize * subgroupSize);
- return xegpu::LayoutAttr::get(context, instData);
+ SmallVector<int64_t> instData(resShape.size(), 1);
+ laneData.back() = std::min(static_cast<int64_t>(consumerLaneData.back()),
+ int64_t(maxChunkSize));
+ laneLayout.back() = consumerLaneLayout.back();
+ instData.back() = laneData.back() * laneLayout.back();
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- auto [laneLayout, laneData] = computeScatterLoadLaneLayoutAndData(
- resShape, consumerLaneData, subgroupSize, maxChunkSize);
- SmallVector<int> laneLayout32(laneLayout.begin(), laneLayout.end());
- SmallVector<int> laneData32(laneData.begin(), laneData.end());
- return xegpu::LayoutAttr::get(context, laneLayout32, laneData32);
+
+ laneData.back() = std::min(static_cast<int64_t>(consumerLaneData.back()),
+ int64_t(maxChunkSize));
+ laneLayout.back() = consumerLaneLayout.back();
+ return buildLaneLayout(context, laneLayout, laneData);
}
return nullptr;
}
@@ -1356,16 +1395,17 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
/// maxChunkSize)
/// All other entries are 1.
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+static std::pair<SmallVector<int>, SmallVector<int>>
computeScatterStoreLaneLayoutAndData(ArrayRef<int64_t> srcShape,
int subgroupSize, int64_t maxChunkSize) {
int rank = srcShape.size();
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ SmallVector<int> laneLayout(rank, 1), laneData(rank, 1);
int innermost = rank - 1;
- laneLayout[innermost] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[innermost]);
+ laneLayout[innermost] = std::min(static_cast<int>(subgroupSize),
+ static_cast<int>(srcShape[innermost]));
laneData[innermost] =
- std::min(srcShape[innermost] / laneLayout[innermost], maxChunkSize);
+ std::min(static_cast<int>(srcShape[innermost] / laneLayout[innermost]),
+ static_cast<int>(maxChunkSize));
return {laneLayout, laneData};
}
@@ -1394,12 +1434,16 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
SmallVector<int> instData(srcShape.size());
for (size_t i = 0; i < srcShape.size(); ++i)
instData[i] = static_cast<int>(laneLayout[i] * laneData[i]);
- return xegpu::LayoutAttr::get(context, instData);
+ return xegpu::LayoutAttr::get(
+ context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr,
+ /*inst_data=*/DenseI32ArrayAttr::get(context, instData),
+ /*lane_layout=*/DenseI32ArrayAttr::get(context, laneLayout),
+ /*lane_data=*/DenseI32ArrayAttr::get(context, laneData),
+ /*order=*/nullptr);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- SmallVector<int> laneLayout32(laneLayout.begin(), laneLayout.end());
- SmallVector<int> laneData32(laneData.begin(), laneData.end());
- return xegpu::LayoutAttr::get(context, laneLayout32, laneData32);
+ return xegpu::LayoutAttr::get(context, laneLayout, laneData);
}
return nullptr;
}
@@ -1445,13 +1489,15 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
srcShape, subgroupSize);
}
-// This function returns the default lane layout for a given vector type.
+// Returns the default (lane_layout, lane_data) pair for a given 1D/2D vector
+// type used by 2D block IO ops.
// - `packingSize` means multiple consecutive elements can be accessed
-// together as a single unit.
+// together as a single unit.
// - `vnni` means data packing is column-wise (i.e., 2x1xf16 with vnni vs.
-// 1x2xf16 w/o vnni).
+// 1x2xf16 w/o vnni).
template <typename RankedTy>
-static xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo(
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+getDefaultLaneLayoutAndData2DBlockIo(
RankedTy ty, const xegpu::uArch::uArch *uArch,
std::optional<unsigned> packingSize = std::nullopt, bool vnni = false) {
// Expecting at least 1D vector. For rank > 2, leading dims are batch dims.
@@ -1461,17 +1507,27 @@ static xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo(
assert(ty.getElementType().isIntOrFloat() &&
"Expected int or float element type.");
- auto context = ty.getContext();
auto rank = ty.getRank();
- SmallVector<int> laneLayout(rank, 1);
- SmallVector<int> laneData(rank, 1);
+ SmallVector<int64_t> laneLayout(rank, 1);
+ SmallVector<int64_t> laneData(rank, 1);
if (packingSize.has_value()) {
unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
- int &laneDataPos = vnni ? laneData[rank - 2] : laneData.back();
+ int64_t &laneDataPos = vnni ? laneData[rank - 2] : laneData.back();
laneDataPos = bitwidth < *packingSize ? *packingSize / bitwidth : 1;
}
laneLayout.back() = uArch->getSubgroupSize();
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
+ return {laneLayout, laneData};
+}
+
+// Convenience wrapper: returns a LayoutAttr carrying only the default lane
+// layout / lane data for a 2D block IO vector type.
+template <typename RankedTy>
+static xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo(
+ RankedTy ty, const xegpu::uArch::uArch *uArch,
+ std::optional<unsigned> packingSize = std::nullopt, bool vnni = false) {
+ auto [laneLayout, laneData] =
+ getDefaultLaneLayoutAndData2DBlockIo(ty, uArch, packingSize, vnni);
+ return buildLaneLayout(ty.getContext(), laneLayout, laneData);
}
// This function returns all layouts for the given sgCount, whose sgData:
@@ -1670,7 +1726,12 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
-
+ auto [laneLayoutA, laneDataA] = getDefaultLaneLayoutAndData2DBlockIo(
+ aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = getDefaultLaneLayoutAndData2DBlockIo(
+ bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
+ auto [laneLayoutCD, laneDataCD] =
+ getDefaultLaneLayoutAndData2DBlockIo(cdTy, uArch);
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
@@ -1682,19 +1743,14 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
return std::make_tuple(
- xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataA.begin(), instDataA.end())),
- xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataB.begin(), instDataB.end())),
- xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataCD.begin(), instDataCD.end())));
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
+ buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
+ laneDataCD));
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- auto aLayout = getDefaultLaneLayout2DBlockIo(
- aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
- auto bLayout = getDefaultLaneLayout2DBlockIo(
- bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
- auto cdLayout = getDefaultLaneLayout2DBlockIo(
- cdTy, uArch /*, packingSize = std::nullopt */);
+ auto aLayout = buildLaneLayout(context, laneLayoutA, laneDataA);
+ auto bLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
+ auto cdLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
return std::make_tuple(aLayout, bLayout, cdLayout);
}
return std::nullopt;
@@ -1833,12 +1889,21 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- auto dpasALayout = xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataA.begin(), instDataA.end()));
- auto dpasBLayout = xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataB.begin(), instDataB.end()));
- auto dpasCDLayout = xegpu::LayoutAttr::get(
- context, SmallVector<int>(instDataCD.begin(), instDataCD.end()));
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ auto [laneLayoutA, laneDataA] = getDefaultLaneLayoutAndData2DBlockIo(
+ aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = getDefaultLaneLayoutAndData2DBlockIo(
+ bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
+ auto [laneLayoutCD, laneDataCD] =
+ getDefaultLaneLayoutAndData2DBlockIo(cdTy, uArch);
+ auto dpasALayout =
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA);
+ auto dpasBLayout =
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB);
+ auto dpasCDLayout = buildInstDataLayoutWithLane(context, instDataCD,
+ laneLayoutCD, laneDataCD);
// Create scale layouts
auto aScaleLayout =
>From d917853241560b9d6ecc73d5a38e37f81b9104a0 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 9 Jun 2026 23:22:35 +0000
Subject: [PATCH 05/42] add setupStoreNdAnchorLayout and
setupPrefetchNdAnchorLayout. shared generic helper sets inst_data and lane
info together for nd ops.
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 15 ++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 149 +++++++++++++++++-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 125 ++++-----------
3 files changed, 194 insertions(+), 95 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 253dcb18af67c..759d23b2d40e6 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -231,6 +231,21 @@ DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
int chunkSize,
const uArch::uArch *uArch);
+/// Sets up the anchor layout for a store_nd operation. StoreNd does not
+/// consider a consumer layout (it is a data sink), and picks its layout from
+/// uArch block parameters. `numSg` is only used for Subgroup-kind layouts.
+DistributeLayoutAttr setupStoreNdAnchorLayout(LayoutKind layoutKind,
+ VectorType vectorTy, int numSg,
+ const uArch::uArch *uArch);
+
+/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
+/// value result and thus no consumer; it picks its layout from uArch block
+/// parameters. `numSg` is only used for Subgroup-kind layouts.
+DistributeLayoutAttr setupPrefetchNdAnchorLayout(LayoutKind layoutKind,
+ TensorDescType tdescTy,
+ int numSg,
+ const uArch::uArch *uArch);
+
/// Sets up the anchor layouts for a dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout creation.
std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index fc2aad83b92a6..1b63f29571115 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1489,6 +1489,153 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
srcShape, subgroupSize);
}
+// Forward declaration: defined later in the file.
+using LayoutRepresentation = std::pair<int64_t, int64_t>;
+static SmallVector<LayoutRepresentation>
+getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
+ int64_t sgCount);
+
+/// Generic anchor-layout setup for ND ops (store_nd, prefetch_nd) that pick
+/// their own layout without considering a consumer.
+///
+/// Given hardware-supported block widths/heights, picks the largest divisor
+/// of the trailing two dims of `dataShape` and uses that as `inst_data`. The
+/// lane layout is the standard 2D-block-IO default (subgroupSize lanes on the
+/// innermost dim, optional packing on the innermost or second-to-innermost
+/// for vnni).
+///
+/// For Lane kind: returns just the lane layout / lane data.
+/// For InstData kind: returns inst_data + the lane layout it must be a
+/// multiple of (Category A: inst_data = k * lane_layout * lane_data, k>=1).
+/// For Subgroup kind: requires `numSg`, picks the most balanced sg_layout
+/// that evenly divides the shape and where sg_data is a multiple of the
+/// chosen inst_data.
+static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
+ xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
+ ArrayRef<int64_t> dataShape, Type elemTy, ArrayRef<int> bWidths,
+ ArrayRef<int> bHeights, unsigned packingSize, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ int rank = dataShape.size();
+ assert(rank >= 1 && "Expected at least 1D shape for ND op");
+
+ // Compute the default 2D block IO lane layout / lane data.
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
+ int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
+ SmallVector<int64_t> laneLayout(rank, 1);
+ SmallVector<int64_t> laneData(rank, 1);
+ laneLayout.back() = uArch->getSubgroupSize();
+ laneData.back() = packingFactor;
+
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return buildLaneLayout(context, laneLayout, laneData);
+
+ // Compute inst_data from hardware block params: pick the largest supported
+ // block width/height that divides the corresponding tensor dim.
+ SmallVector<int64_t> instData(rank, 1);
+ int instWidth = xegpu::getLargestDivisor(
+ static_cast<int>(dataShape.back()), bWidths);
+ if (instWidth == -1)
+ return nullptr;
+ instData.back() = instWidth;
+ if (rank >= 2) {
+ int instHeight = xegpu::getLargestDivisor(
+ static_cast<int>(dataShape[rank - 2]), bHeights);
+ if (instHeight == -1)
+ return nullptr;
+ instData[rank - 2] = instHeight;
+ }
+
+ // Validate Category A invariant: inst_data must be a multiple of
+ // lane_layout * lane_data on each dim (k >= 1).
+ for (int dim = 0; dim < rank; ++dim) {
+ int64_t laneProduct = laneLayout[dim] * laneData[dim];
+ assert(instData[dim] % laneProduct == 0 &&
+ "inst_data must be a multiple of lane_layout * lane_data for ND op");
+ (void)laneProduct;
+ }
+
+ if (layoutKind == xegpu::LayoutKind::InstData)
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ // Subgroup-kind layout creation currently only supports rank-2 shapes
+ // (mirrors getValidLayouts).
+ if (rank != 2)
+ return nullptr;
+ auto sgLayouts = getValidLayouts(dataShape, instData, numSg);
+ if (sgLayouts.empty())
+ return nullptr;
+ SmallVector<int> sgLayout = {static_cast<int>(sgLayouts[0].first),
+ static_cast<int>(sgLayouts[0].second)};
+ SmallVector<int> sgData = {
+ static_cast<int>(dataShape[0]) / sgLayout[0],
+ static_cast<int>(dataShape[1]) / sgLayout[1]};
+ return xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgData),
+ /*inst_data=*/nullptr, /*lane_layout=*/nullptr,
+ /*lane_data=*/nullptr, /*order=*/nullptr);
+ }
+
+ return nullptr;
+}
+
+/// Sets up the anchor layout for a store_nd operation. StoreNd picks its
+/// own layout based on uArch block parameters (it does not take a consumer
+/// layout, since it is a data sink).
+xegpu::DistributeLayoutAttr
+xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = srcVecTy.getContext();
+ Type elemTy = srcVecTy.getElementType();
+
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
+ if (!uArchInstruction)
+ return nullptr;
+ auto blockWHC = uArchInstruction->getBlockWidthHeightCount(elemTy);
+ if (!blockWHC)
+ return nullptr;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
+ unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
+
+ return setupGenericNdAnchorLayout(layoutKind, context, srcVecTy.getShape(),
+ elemTy, bWidths, bHeights, packingSize,
+ numSg, uArch);
+}
+
+/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
+/// consumer (it produces no value), so it picks its own layout from uArch
+/// block parameters.
+xegpu::DistributeLayoutAttr
+xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ xegpu::TensorDescType tdescTy, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = tdescTy.getContext();
+ Type elemTy = tdescTy.getElementType();
+
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
+ if (!uArchInstruction)
+ return nullptr;
+ auto blockWHC = uArchInstruction->getBlockWidthHeightCount(elemTy);
+ if (!blockWHC)
+ return nullptr;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
+ unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
+
+ return setupGenericNdAnchorLayout(layoutKind, context, tdescTy.getShape(),
+ elemTy, bWidths, bHeights, packingSize,
+ numSg, uArch);
+}
+
// Returns the default (lane_layout, lane_data) pair for a given 1D/2D vector
// type used by 2D block IO ops.
// - `packingSize` means multiple consecutive elements can be accessed
@@ -1537,7 +1684,7 @@ static xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo(
// wgShape = [128, 64], instData = [8, 16], sgCount = 32
// Returns layouts:
// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
-using LayoutRepresentation = std::pair<int64_t, int64_t>;
+// Definition (forward-declared above).
static SmallVector<LayoutRepresentation>
getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
int64_t sgCount) {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index c1ad3c4ee8e2e..0e371d98f61f7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -599,49 +599,31 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
if (hasParamsOfLayoutKind(anchorLayout)) {
prefetchLayout = LayoutInfo(anchorLayout);
} else {
- // Here we assign the default layout to the tensor descriptor operand of
- // prefetch.
auto tdescTy = prefetch.getTensorDescType();
-
const uArch *uArch = getUArch(getChipStr(prefetch).value_or(""));
if (!uArch)
return;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
-
- auto blockWHC =
- uArchInstruction->getBlockWidthHeightCount(tdescTy.getElementType());
- if (!blockWHC)
- prefetch.emitWarning("No known block params found for the element type.");
- auto [bWidth, bHeight, bCount] = blockWHC.value();
- SmallVector<int> instData;
- int instWidth = xegpu::getLargestDivisor(
- static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 1)), bWidth);
- if (instWidth == -1)
- prefetch.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- if (tdescTy.getRank() == 1)
- instData = {instWidth};
- else {
- int instHeight = xegpu::getLargestDivisor(
- static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 2)), bHeight);
- if (instHeight == -1)
+
+ int numSg = 0;
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ auto numSgOrErr = getNumSg(prefetch, uArch->getSubgroupSize());
+ if (failed(numSgOrErr)) {
prefetch.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- instData = {instHeight, instWidth};
+ "Unable to determine the number of subgroups for the operation.");
+ return;
+ }
+ numSg = numSgOrErr.value();
}
- if (layoutKind == xegpu::LayoutKind::InstData)
- prefetchLayout =
- LayoutInfo(xegpu::LayoutAttr::get(tdescTy.getContext(), instData));
- else
- prefetchLayout = getSIMTLayoutInfoBlockIO(
- tdescTy, uArch, uArchInstruction->getPackedFormatBitSize());
-
- prefetch.setLayoutAttr(
- dyn_cast<xegpu::DistributeLayoutAttr>(prefetchLayout.get()));
+ auto layoutAttr =
+ xegpu::setupPrefetchNdAnchorLayout(layoutKind, tdescTy, numSg, uArch);
+ if (!layoutAttr) {
+ prefetch.emitWarning(
+ "Failed to determine required layout for prefetch_nd.");
+ return;
+ }
+ prefetchLayout = LayoutInfo(layoutAttr);
+ prefetch.setLayoutAttr(layoutAttr);
}
// Propagate the layout to the source tensor descriptor.
propagateIfChanged(operands[0], operands[0]->meet(prefetchLayout));
@@ -980,71 +962,26 @@ void LayoutInfoPropagation::visitStoreNdOp(
const uArch *uArch = getUArch(getChipStr(store).value_or(""));
if (!uArch)
return;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
- VectorType dataTy = store.getValueType();
- auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
- store.getValueType().getElementType());
- if (!blockWHC)
- store.emitWarning("No known block params found for the element type.");
- auto [bWidth, bHeight, bCount] = blockWHC.value();
- // Default to 1 for any leading batch dims; rank-1 and rank>=2 cases
- // overwrite the trailing entries below.
- SmallVector<int> instData(dataTy.getRank(), 1);
- int instWidth = xegpu::getLargestDivisor(
- static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 1)), bWidth);
- if (instWidth == -1)
- store.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- if (dataTy.getRank() == 1) {
- instData = {instWidth};
- } else {
- int instHeight = xegpu::getLargestDivisor(
- static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 2)), bHeight);
- if (instHeight == -1)
- store.emitWarning(
- "No suitable instruction multiple found for the given shape.");
- instData[dataTy.getRank() - 2] = instHeight;
- instData[dataTy.getRank() - 1] = instWidth;
- }
- if (layoutKind == xegpu::LayoutKind::InstData)
- storeLayout =
- LayoutInfo(xegpu::LayoutAttr::get(dataTy.getContext(), instData));
- else if (layoutKind == xegpu::LayoutKind::Lane)
- storeLayout =
- getSIMTLayoutInfoBlockIO(store.getValueType(), uArch,
- uArchInstruction->getPackedFormatBitSize());
- else { // xegpu::LayoutKind::Subgroup
- auto sgSize = uArch->getSubgroupSize();
- auto numSgOrErr = getNumSg(store, sgSize);
+ int numSg = 0;
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ auto numSgOrErr = getNumSg(store, uArch->getSubgroupSize());
if (failed(numSgOrErr)) {
store.emitWarning(
"Unable to determine the number of subgroups for the operation.");
return;
}
- auto sgLayouts = getValidLayouts(store.getValueType().getShape(),
- instData, numSgOrErr.value());
- if (sgLayouts.empty()) {
- store.emitWarning(
- "Unable to determine suitable subgroup layout for store value.");
- return;
- }
- SmallVector<int> sgLayout = {sgLayouts[0].first, sgLayouts[0].second};
- SmallVector<int> sgData = {
- static_cast<int>(dataTy.getShape()[0]) / sgLayout[0],
- static_cast<int>(dataTy.getShape()[1]) / sgLayout[1]};
- storeLayout = LayoutInfo(xegpu::LayoutAttr::get(
- dataTy.getContext(),
- DenseI32ArrayAttr::get(dataTy.getContext(), sgLayout),
- DenseI32ArrayAttr::get(dataTy.getContext(), sgData),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/nullptr));
+ numSg = numSgOrErr.value();
+ }
+
+ auto layoutAttr = xegpu::setupStoreNdAnchorLayout(
+ layoutKind, store.getValueType(), numSg, uArch);
+ if (!layoutAttr) {
+ store.emitWarning("Failed to determine required layout for store_nd.");
+ return;
}
- store.setLayoutAttr(
- dyn_cast<xegpu::DistributeLayoutAttr>(storeLayout.get()));
+ storeLayout = LayoutInfo(layoutAttr);
+ store.setLayoutAttr(layoutAttr);
}
// Propagate the layout to the value operand.
// Both operands should have the same layout
>From e3b6e22971263e8197decd52fcb5d58b1e284cfd Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 01:43:06 +0000
Subject: [PATCH 06/42] add setupLoadNdAnchorLayout. validates consumer
inst_data/sg_layout against uArch; falls back to defaults otherwise.
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 12 ++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 149 +++++++++++++++---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 47 +++++-
3 files changed, 178 insertions(+), 30 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 759d23b2d40e6..2d2ec1c33b0f4 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -246,6 +246,18 @@ DistributeLayoutAttr setupPrefetchNdAnchorLayout(LayoutKind layoutKind,
int numSg,
const uArch::uArch *uArch);
+/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
+/// (downstream) consumer layout and validates it against uArch constraints;
+/// when valid, the consumer's `inst_data` / `sg_layout` are honored.
+/// Otherwise defaults derived from uArch block parameters are used.
+/// `consumerLayout` may be null. `numSg` is only used for Subgroup-kind
+/// layouts when the consumer does not already provide an sg_layout.
+DistributeLayoutAttr setupLoadNdAnchorLayout(LayoutKind layoutKind,
+ VectorType vectorTy,
+ DistributeLayoutAttr consumerLayout,
+ int numSg,
+ const uArch::uArch *uArch);
+
/// Sets up the anchor layouts for a dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout creation.
std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 1b63f29571115..da8863547bfae 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1495,25 +1495,72 @@ static SmallVector<LayoutRepresentation>
getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
int64_t sgCount);
-/// Generic anchor-layout setup for ND ops (store_nd, prefetch_nd) that pick
-/// their own layout without considering a consumer.
+/// Validates whether `instData` is a hardware-viable inst_data for an ND op
+/// with the given block params and lane factor. Specifically:
+/// - leading batch dims must be 1
+/// - innermost dim must be a divisor of `dataShape.back()`, a multiple of
+/// `bWidths.front()` (smallest supported block width), and ≤ a small
+/// multiple of the largest supported block width
+/// - second-to-innermost dim must be in the supported heights for rank >= 2
+/// - each dim must be a multiple of `lane_layout[dim] * lane_data[dim]`
+static bool isValidNdInstData(ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> dataShape,
+ ArrayRef<int> bWidths, ArrayRef<int> bHeights,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
+ int rank = dataShape.size();
+ if (static_cast<int>(instData.size()) != rank)
+ return false;
+
+ for (int dim = 0; dim < rank - 2; ++dim)
+ if (instData[dim] != 1)
+ return false;
+
+ int64_t inner = instData.back();
+ if (inner <= 0 || dataShape.back() % inner != 0)
+ return false;
+ int minWidth = *llvm::min_element(bWidths);
+ int maxWidth = *llvm::max_element(bWidths);
+ if (inner % minWidth != 0 || inner > maxWidth * /*maxBlockCount*/ 4)
+ return false;
+
+ if (rank >= 2) {
+ int64_t height = instData[rank - 2];
+ if (!llvm::is_contained(bHeights, static_cast<int>(height)))
+ return false;
+ }
+
+ for (int dim = 0; dim < rank; ++dim) {
+ int64_t laneProduct = laneLayout[dim] * laneData[dim];
+ if (laneProduct == 0 || instData[dim] % laneProduct != 0)
+ return false;
+ }
+ return true;
+}
+
+/// Generic anchor-layout setup for ND ops (load_nd, store_nd, prefetch_nd).
///
/// Given hardware-supported block widths/heights, picks the largest divisor
-/// of the trailing two dims of `dataShape` and uses that as `inst_data`. The
+/// of the trailing two dims of `dataShape` as the default `inst_data`. The
/// lane layout is the standard 2D-block-IO default (subgroupSize lanes on the
-/// innermost dim, optional packing on the innermost or second-to-innermost
-/// for vnni).
+/// innermost dim, optional packing on the innermost dim).
+///
+/// `consumerLayout` (optional) is honored when its parameters are valid w.r.t.
+/// the uArch constraints; otherwise the helper falls back to defaults.
///
-/// For Lane kind: returns just the lane layout / lane data.
-/// For InstData kind: returns inst_data + the lane layout it must be a
-/// multiple of (Category A: inst_data = k * lane_layout * lane_data, k>=1).
-/// For Subgroup kind: requires `numSg`, picks the most balanced sg_layout
-/// that evenly divides the shape and where sg_data is a multiple of the
-/// chosen inst_data.
+/// For Lane kind: returns just the lane layout / lane data (consumer ignored;
+/// lane layout is fully determined by hardware).
+/// For InstData kind: returns inst_data + lane_layout/lane_data (Category A:
+/// inst_data = k * lane_layout * lane_data, k >= 1). Honors consumer's
+/// inst_data when it is uArch-valid.
+/// For Subgroup kind: if the consumer specifies a workgroup-level layout,
+/// reuses it directly; otherwise picks the most balanced sg_layout via
+/// `getValidLayouts` (requires `numSg`).
static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
ArrayRef<int64_t> dataShape, Type elemTy, ArrayRef<int> bWidths,
- ArrayRef<int> bHeights, unsigned packingSize, int numSg,
+ ArrayRef<int> bHeights, unsigned packingSize,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
const xegpu::uArch::uArch *uArch) {
int rank = dataShape.size();
assert(rank >= 1 && "Expected at least 1D shape for ND op");
@@ -1529,8 +1576,15 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
- // Compute inst_data from hardware block params: pick the largest supported
- // block width/height that divides the corresponding tensor dim.
+ // Subgroup-kind fast path: if the consumer already specifies a
+ // workgroup-level layout, reuse it directly. Skip the inst_data
+ // computation, which can fail for very small shapes (e.g. dpas_mx scale
+ // operands like 128x16 where no supported block width divides 16).
+ if (layoutKind == xegpu::LayoutKind::Subgroup && consumerLayout &&
+ consumerLayout.isForWorkgroup())
+ return consumerLayout;
+
+ // Compute the default inst_data from hardware block params.
SmallVector<int64_t> instData(rank, 1);
int instWidth = xegpu::getLargestDivisor(
static_cast<int>(dataShape.back()), bWidths);
@@ -1545,8 +1599,17 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
instData[rank - 2] = instHeight;
}
- // Validate Category A invariant: inst_data must be a multiple of
- // lane_layout * lane_data on each dim (k >= 1).
+ // Honor the consumer's inst_data if it is uArch-valid.
+ if (consumerLayout) {
+ SmallVector<int64_t> consumerInstData =
+ consumerLayout.getEffectiveInstDataAsInt();
+ if (!consumerInstData.empty() &&
+ isValidNdInstData(consumerInstData, dataShape, bWidths, bHeights,
+ laneLayout, laneData))
+ instData.assign(consumerInstData.begin(), consumerInstData.end());
+ }
+
+ // Category A invariant: inst_data is a multiple of lane_layout * lane_data.
for (int dim = 0; dim < rank; ++dim) {
int64_t laneProduct = laneLayout[dim] * laneData[dim];
assert(instData[dim] % laneProduct == 0 &&
@@ -1558,12 +1621,12 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(numSg > 0 &&
- "Number of subgroups must be provided for sg layout creation.");
- // Subgroup-kind layout creation currently only supports rank-2 shapes
- // (mirrors getValidLayouts).
if (rank != 2)
return nullptr;
+ // The consumer-fast-path above already returned for the case where the
+ // consumer carries a workgroup-level layout, so here numSg must be set.
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
auto sgLayouts = getValidLayouts(dataShape, instData, numSg);
if (sgLayouts.empty())
return nullptr;
@@ -1606,7 +1669,7 @@ xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
return setupGenericNdAnchorLayout(layoutKind, context, srcVecTy.getShape(),
elemTy, bWidths, bHeights, packingSize,
- numSg, uArch);
+ /*consumerLayout=*/nullptr, numSg, uArch);
}
/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
@@ -1633,7 +1696,49 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
return setupGenericNdAnchorLayout(layoutKind, context, tdescTy.getShape(),
elemTy, bWidths, bHeights, packingSize,
- numSg, uArch);
+ /*consumerLayout=*/nullptr, numSg, uArch);
+}
+
+/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
+/// 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.
+xegpu::DistributeLayoutAttr xegpu::setupLoadNdAnchorLayout(
+ xegpu::LayoutKind layoutKind, VectorType resVecTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = resVecTy.getContext();
+ Type elemTy = resVecTy.getElementType();
+
+ // Subgroup-kind fast path: if the consumer already specifies a complete
+ // workgroup-level layout, reuse it directly. We don't need the uArch block
+ // params at all (which may be unavailable for unusual element types like
+ // sub-byte floats used in dpas_mx scale operands).
+ if (layoutKind == xegpu::LayoutKind::Subgroup && consumerLayout &&
+ consumerLayout.isForWorkgroup())
+ return consumerLayout;
+
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
+ if (!uArchInstruction)
+ return nullptr;
+ // Transform / transpose / upConv are lane-level concerns; treat them as
+ // no-op at the propagation stage (consistent with the existing
+ // visitLoadNdOp, which warns on transpose).
+ auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
+ elemTy, /*hasTransform=*/false, /*hasTranspose=*/false,
+ /*upConv=*/false);
+ if (!blockWHC)
+ return nullptr;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
+ unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
+
+ return setupGenericNdAnchorLayout(layoutKind, context, resVecTy.getShape(),
+ elemTy, bWidths, bHeights, packingSize,
+ consumerLayout, numSg, uArch);
}
// Returns the default (lane_layout, lane_data) pair for a given 1D/2D vector
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 0e371d98f61f7..df5d1409b8395 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -999,21 +999,52 @@ void LayoutInfoPropagation::visitLoadNdOp(
if (hasParamsOfLayoutKind(anchorLayout)) {
loadLayout = LayoutInfo(anchorLayout);
} else {
-
LayoutInfo valueLayout = results[0]->getValue();
- // Need the layout of the value to propagate to the tensor descriptor.
if (!valueLayout.isAssigned())
return;
- loadLayout = valueLayout;
- // LoadNdOp has the transpose effect. However, at the stage of this analysis
- // this effect is not expected and should be abstracted away. Emit a
- // warning.
+
+ // LoadNdOp has a transpose effect. At this stage of analysis it is not
+ // expected and should be abstracted away; emit a warning and pre-transpose
+ // the consumer hint so it lines up with the load's source.
+ auto consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(valueLayout.get());
if (auto transpose = load.getTranspose()) {
load.emitWarning("Transpose effect is not expected for LoadNdOp at "
"LayoutInfoPropagation stage.");
- loadLayout = valueLayout.transpose(transpose.value());
+ LayoutInfo transposed = valueLayout.transpose(transpose.value());
+ consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(transposed.get());
+ }
+
+ const uArch *uArch = getUArch(getChipStr(load).value_or(""));
+ if (!uArch)
+ return;
+
+ int numSg = 0;
+ // numSg is only needed when the helper has to derive a fresh sg_layout
+ // for Subgroup kind. If the consumer already provides a complete
+ // sg_layout, the helper will reuse it without consulting numSg.
+ bool consumerHasSgLayout =
+ consumerLayoutAttr &&
+ !consumerLayoutAttr.getEffectiveSgLayoutAsInt().empty();
+ if (layoutKind == xegpu::LayoutKind::Subgroup && !consumerHasSgLayout) {
+ auto numSgOrErr = getNumSg(load, uArch->getSubgroupSize());
+ if (failed(numSgOrErr)) {
+ load.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
+ }
+ numSg = numSgOrErr.value();
+ }
+
+ auto layoutAttr = xegpu::setupLoadNdAnchorLayout(
+ layoutKind, load.getType(), consumerLayoutAttr, numSg, uArch);
+ if (!layoutAttr) {
+ load.emitWarning("Failed to determine required layout for load_nd.");
+ return;
}
- load.setLayoutAttr(dyn_cast<xegpu::DistributeLayoutAttr>(loadLayout.get()));
+ loadLayout = LayoutInfo(layoutAttr);
+ load.setLayoutAttr(layoutAttr);
}
// Propagate the new layout to the tensor descriptor operand.
propagateIfChanged(operands[0], operands[0]->meet(loadLayout));
>From e671f4a1fc4a0929f4140f57518102930faddc66 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 02:28:49 +0000
Subject: [PATCH 07/42] complete inst_data-only anchor layouts on scatter ops;
preserve lane info in shape_cast collapse inference. eliminates assertion
crashes when user only specifies inst_data.
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 17 +++++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 76 +++++++++++++++++++
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 28 +++++--
3 files changed, 113 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 2d2ec1c33b0f4..f3cbf6ae72f5d 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -231,6 +231,23 @@ DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
int chunkSize,
const uArch::uArch *uArch);
+/// If the consumer layout has only inst_data (no lane_layout/lane_data),
+/// completes it by running the corresponding scatter-style Lane-kind setup
+/// rule with inst_data as the destination shape. The resulting lane info is
+/// merged with the consumer's inst_data so downstream setup* paths see a
+/// fully-populated layout. If the consumer already has lane info (or no
+/// inst_data), returns it unchanged. The "load-side" version uses
+/// setupLoadGatherAnchorLayout under the hood; the "store-side" version uses
+/// setupStoreScatterAnchorLayout.
+DistributeLayoutAttr
+completeLoadGatherLayoutFromInstData(DistributeLayoutAttr consumerLayout,
+ Type elemTy, const uArch::uArch *uArch);
+
+DistributeLayoutAttr
+completeStoreScatterLayoutFromInstData(DistributeLayoutAttr consumerLayout,
+ Type elemTy,
+ const uArch::uArch *uArch);
+
/// Sets up the anchor layout for a store_nd operation. StoreNd does not
/// consider a consumer layout (it is a data sink), and picks its layout from
/// uArch block parameters. `numSg` is only used for Subgroup-kind layouts.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index da8863547bfae..51b17b3993ceb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1489,6 +1489,82 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
srcShape, subgroupSize);
}
+/// If `consumerLayout` has inst_data set but no lane_layout/lane_data,
+/// derive a lane factorization by re-running the load-side Lane setup with
+/// inst_data as the destination shape, and merge the result back so the
+/// returned LayoutAttr carries inst_data + lane_layout + lane_data. This
+/// guarantees the lane factorization the downstream load setup will see is
+/// the same one its own Lane-kind setup would produce.
+xegpu::DistributeLayoutAttr xegpu::completeLoadGatherLayoutFromInstData(
+ xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::uArch *uArch) {
+ if (!consumerLayout)
+ return consumerLayout;
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ if (instData.empty())
+ return consumerLayout;
+ if (!consumerLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !consumerLayout.getEffectiveLaneDataAsInt().empty())
+ return consumerLayout;
+
+ // Reuse the load-side setup with inst_data as the destination shape.
+ const int subgroupSize = uArch->getSubgroupSize();
+ auto *context = consumerLayout.getContext();
+ auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ if (!uArchInstruction)
+ return consumerLayout;
+ int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+
+ auto laneOnly = setupGenericLoadAnchorLayout(
+ xegpu::LayoutKind::Lane, context, /*consumerLayout=*/nullptr,
+ maxChunkSize, instData, subgroupSize);
+ if (!laneOnly)
+ return consumerLayout;
+
+ SmallVector<int64_t> laneLayout = laneOnly.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = laneOnly.getEffectiveLaneDataAsInt();
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+}
+
+/// If `consumerLayout` has inst_data set but no lane_layout/lane_data,
+/// derive a lane factorization by re-running the store-side Lane setup with
+/// inst_data as the destination shape, and merge the result back. Returned
+/// LayoutAttr carries inst_data + lane_layout + lane_data.
+xegpu::DistributeLayoutAttr xegpu::completeStoreScatterLayoutFromInstData(
+ xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::uArch *uArch) {
+ if (!consumerLayout)
+ return consumerLayout;
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ if (instData.empty())
+ return consumerLayout;
+ if (!consumerLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !consumerLayout.getEffectiveLaneDataAsInt().empty())
+ return consumerLayout;
+
+ const int subgroupSize = uArch->getSubgroupSize();
+ auto *context = consumerLayout.getContext();
+ auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
+ if (!uArchInstruction)
+ return consumerLayout;
+ int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+
+ auto laneOnly = setupGenericStoreAnchorLayout(
+ xegpu::LayoutKind::Lane, context, maxChunkSize, instData, subgroupSize);
+ if (!laneOnly)
+ return consumerLayout;
+
+ SmallVector<int64_t> laneLayout = laneOnly.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = laneOnly.getEffectiveLaneDataAsInt();
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+}
+
// Forward declaration: defined later in the file.
using LayoutRepresentation = std::pair<int64_t, int64_t>;
static SmallVector<LayoutRepresentation>
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index df5d1409b8395..5a660101879a2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1221,7 +1221,11 @@ void LayoutInfoPropagation::visitLoadGatherOp(
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
- requiredAnchorLayoutAttr = anchorLayoutAttr;
+ // The user-provided anchor may carry only inst_data; complete it with
+ // a lane factorization derived from the load-side Lane setup so
+ // downstream paths see a fully-populated layout.
+ requiredAnchorLayoutAttr = xegpu::completeLoadGatherLayoutFromInstData(
+ anchorLayoutAttr, resVecTy.getElementType(), uArch);
} else {
if (!resVecTy) {
load.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1261,7 +1265,11 @@ void LayoutInfoPropagation::visitStoreScatterOp(
int chunkSize = storeScatter.getChunkSize().value_or(1);
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
- requiredAnchorLayoutAttr = anchorLayoutAttr;
+ // The user-provided anchor may carry only inst_data; complete it with
+ // a lane factorization derived from the store-side Lane setup so
+ // downstream paths see a fully-populated layout.
+ requiredAnchorLayoutAttr = xegpu::completeStoreScatterLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArch);
} else {
if (!srcVecTy) {
storeScatter.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1322,14 +1330,18 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
ArrayRef<const LayoutInfoLattice *> results) {
xegpu::DistributeLayoutAttr anchorLayout = storeMatrix.getLayoutAttr();
LayoutInfo layout;
+ VectorType srcVecTy =
+ llvm::cast<VectorType>(storeMatrix.getData().getType());
+ const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
+ if (!uArch)
+ return;
if (hasParamsOfLayoutKind(anchorLayout)) {
- layout = LayoutInfo(anchorLayout);
+ // The user-provided anchor may carry only inst_data; complete it with
+ // a lane factorization derived from the store-side Lane setup.
+ auto completed = xegpu::completeStoreScatterLayoutFromInstData(
+ anchorLayout, srcVecTy.getElementType(), uArch);
+ layout = LayoutInfo(completed);
} else {
- VectorType srcVecTy =
- llvm::cast<VectorType>(storeMatrix.getData().getType());
- const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
- if (!uArch)
- return;
int chunkSize =
1; // placeHolder for future use when StoreMatrix supports coalescing
auto requiredAnchorLayoutAttr = xegpu::setupStoreMatrixAnchorLayout(
>From 557b76b683291563dea5cd1c684b39d10d776792 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 02:34:59 +0000
Subject: [PATCH 08/42] honor consumer lane_layout/lane_data in nd anchor
setup. dpas operands with vnni packing now propagate correctly to load_nd.
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 24 +++++++++++++++++++
1 file changed, 24 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 51b17b3993ceb..83892079da7b0 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1649,6 +1649,30 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
laneLayout.back() = uArch->getSubgroupSize();
laneData.back() = packingFactor;
+ // Honor the consumer's lane info when valid. The consumer (e.g. dpas) may
+ // require VNNI-style packing on a non-innermost dim that the load's own
+ // packing default doesn't capture.
+ auto honorConsumerLane = [&]() {
+ if (!consumerLayout)
+ return;
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ if (consumerLaneLayout.size() != static_cast<size_t>(rank) ||
+ consumerLaneData.size() != static_cast<size_t>(rank))
+ return;
+ // Validate: lane_layout * lane_data must divide each dim of dataShape.
+ for (int dim = 0; dim < rank; ++dim) {
+ int64_t product = consumerLaneLayout[dim] * consumerLaneData[dim];
+ if (product == 0 || dataShape[dim] % product != 0)
+ return;
+ }
+ laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
+ laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
+ };
+ honorConsumerLane();
+
if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
>From d8a007c988ec8d5ae90c8e27b562e610c183c57d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 03:53:47 +0000
Subject: [PATCH 09/42] lane-kind nd anchor setup skips block-WHC lookup. lets
sub-byte float types (e.g. f4E2M1FN) get a default lane layout via packing
factor alone. also propagate consumer order through nd anchor setup.
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 48 +++++++++++++------
1 file changed, 34 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 83892079da7b0..ad591625754ae 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -811,9 +811,12 @@ static bool leadingDimsAreUnit(ArrayRef<int64_t> shape, int numInnerDims) {
/// result layout can later be distributed without re-deriving the lane
/// layout. `instData`, `laneLayout`, and `laneData` may have different
/// element types; they are normalized to int32 entries.
-static xegpu::LayoutAttr buildInstDataLayoutWithLane(
- mlir::MLIRContext *context, ArrayRef<int64_t> instData,
- ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData) {
+static xegpu::LayoutAttr
+buildInstDataLayoutWithLane(mlir::MLIRContext *context,
+ ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
auto toI32Attr = [&](auto range) {
SmallVector<int32_t> v(range.begin(), range.end());
return DenseI32ArrayAttr::get(context, v);
@@ -823,12 +826,13 @@ static xegpu::LayoutAttr buildInstDataLayoutWithLane(
/*inst_data=*/toI32Attr(instData),
/*lane_layout=*/toI32Attr(laneLayout),
/*lane_data=*/toI32Attr(laneData),
- /*order=*/nullptr);
+ /*order=*/orderAttr);
}
-static xegpu::LayoutAttr buildLaneLayout(mlir::MLIRContext *context,
- ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData) {
+static xegpu::LayoutAttr
+buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
auto toI32Attr = [&](auto range) {
SmallVector<int32_t> v(range.begin(), range.end());
return DenseI32ArrayAttr::get(context, v);
@@ -838,7 +842,7 @@ static xegpu::LayoutAttr buildLaneLayout(mlir::MLIRContext *context,
/*inst_data=*/nullptr,
/*lane_layout=*/toI32Attr(laneLayout),
/*lane_data=*/toI32Attr(laneData),
- /*order=*/nullptr);
+ /*order=*/orderAttr);
}
/// Computes the lane_layout and lane_data for a multi-reduction's source
@@ -1673,8 +1677,12 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
};
honorConsumerLane();
+ // If the consumer carries an explicit `order`, propagate it through.
+ DenseI32ArrayAttr orderAttr =
+ consumerLayout ? consumerLayout.getOrder() : nullptr;
+
if (layoutKind == xegpu::LayoutKind::Lane)
- return buildLaneLayout(context, laneLayout, laneData);
+ return buildLaneLayout(context, laneLayout, laneData, orderAttr);
// Subgroup-kind fast path: if the consumer already specifies a
// workgroup-level layout, reuse it directly. Skip the inst_data
@@ -1718,7 +1726,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
}
if (layoutKind == xegpu::LayoutKind::InstData)
- return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData,
+ orderAttr);
if (layoutKind == xegpu::LayoutKind::Subgroup) {
if (rank != 2)
@@ -1825,16 +1834,27 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadNdAnchorLayout(
xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
if (!uArchInstruction)
return nullptr;
- // Transform / transpose / upConv are lane-level concerns; treat them as
- // no-op at the propagation stage (consistent with the existing
- // visitLoadNdOp, which warns on transpose).
+ unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
+
+ // Lane kind only needs subgroupSize + packingSize. Skip the block-WHC
+ // lookup, which can fail for element types without a uArch entry (e.g.
+ // sub-byte floats like f4E2M1FN), and let the generic helper produce a
+ // default lane layout from packingSize alone.
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return setupGenericNdAnchorLayout(
+ layoutKind, context, resVecTy.getShape(), elemTy,
+ /*bWidths=*/{}, /*bHeights=*/{}, packingSize, consumerLayout, numSg,
+ uArch);
+
+ // InstData / Subgroup kinds need block params. Transform / transpose /
+ // upConv are lane-level concerns; treat them as no-op at the propagation
+ // stage (consistent with visitLoadNdOp, which warns on transpose).
auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
elemTy, /*hasTransform=*/false, /*hasTranspose=*/false,
/*upConv=*/false);
if (!blockWHC)
return nullptr;
auto [bWidths, bHeights, bCounts] = blockWHC.value();
- unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
return setupGenericNdAnchorLayout(layoutKind, context, resVecTy.getShape(),
elemTy, bWidths, bHeights, packingSize,
>From 7dba091eaab1f8bb186ba617e469336c5ef6ca45 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 04:06:15 +0000
Subject: [PATCH 10/42] write completed scatter/store-matrix anchor layout back
to the op so the printed IR reflects the resolved lane info, not the user's
incomplete inst_data-only layout.
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 3 +++
1 file changed, 3 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 5a660101879a2..cb9d6eeed9743 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1226,6 +1226,7 @@ void LayoutInfoPropagation::visitLoadGatherOp(
// downstream paths see a fully-populated layout.
requiredAnchorLayoutAttr = xegpu::completeLoadGatherLayoutFromInstData(
anchorLayoutAttr, resVecTy.getElementType(), uArch);
+ load.setLayoutAttr(requiredAnchorLayoutAttr);
} else {
if (!resVecTy) {
load.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1270,6 +1271,7 @@ void LayoutInfoPropagation::visitStoreScatterOp(
// downstream paths see a fully-populated layout.
requiredAnchorLayoutAttr = xegpu::completeStoreScatterLayoutFromInstData(
anchorLayoutAttr, srcVecTy.getElementType(), uArch);
+ storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
} else {
if (!srcVecTy) {
storeScatter.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1340,6 +1342,7 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
// a lane factorization derived from the store-side Lane setup.
auto completed = xegpu::completeStoreScatterLayoutFromInstData(
anchorLayout, srcVecTy.getElementType(), uArch);
+ storeMatrix.setLayoutAttr(completed);
layout = LayoutInfo(completed);
} else {
int chunkSize =
>From 8a816db74e51e666638ecbbfd73b85df12413b53 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 04:28:59 +0000
Subject: [PATCH 11/42] test: update CHECK lines in
propagate-layout-inst-data.mlir to expect lane_layout/lane_data alongside
inst_data on Nd ops. fixes 10 cases. Remaining 7 cases need investigation of
setupGenericLoadAnchorLayout and dpas_mx scale operand layout.
---
.../XeGPU/propagate-layout-inst-data.mlir | 76 +++++++++----------
1 file changed, 38 insertions(+), 38 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 3ea06d52690dc..da4a2e32259de 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -4,13 +4,13 @@
// CHECK-LABEL: func.func @load_store_no_array_len(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<8x32xf32>) {
// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<8x16xf32>
-// CHECK: %[[TDESC_SRC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: %[[TDESC_DST:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.prefetch_nd %[[TDESC_SRC]][0, 0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<inst_data = [8, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: %[[LOADED:.*]] = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}>
-// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>> -> vector<8x32xf32>
-// CHECK: xegpu.store_nd %[[LOADED]], %[[TDESC_DST]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[TDESC_SRC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[TDESC_DST:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.prefetch_nd %[[TDESC_SRC]][0, 0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOADED:.*]] = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
+// CHECK-SAME: !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x32xf32>
+// CHECK: xegpu.store_nd %[[LOADED]], %[[TDESC_DST]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
// Although the uArch allows 8x32 inst data using block count (or array_len),
// it is up to optimization passes to decide on the block count usage.
@@ -29,17 +29,17 @@ func.func @load_store_no_array_len(%arg0: memref<8x32xf32>, %arg1: memref<8x32xf
// CHECK-LABEL: func.func @dpas_f16(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x16xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<8x16xf32>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16]>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16]>> -> vector<8x16xf16>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16]>> -> vector<16x16xf16>
-// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<inst_data = [8, 16]>, layout_b = #xegpu.layout<inst_data = [16, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<8x16xf32>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
+// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
-// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16]>
-// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
func.func @dpas_f16(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
@@ -70,15 +70,15 @@ gpu.module @test_kernel {
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16>
scf.for %k = %c0 to %c1024 step %c32 {
- //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [8, 16]>}> :
- //CHECK-SAME: !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16]>> -> vector<16x32xf16>
+ //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+ //CHECK-SAME: !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x32xf16>
%a = xegpu.load_nd %a_tdesc[0, %k] : !xegpu.tensor_desc<16x32xf16> -> vector<16x32xf16>
%b = xegpu.load_nd %b_tdesc[0, %k] : !xegpu.tensor_desc<16x32xf16> -> vector<16x32xf16>
- //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : vector<16x32xf16>
+ //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x32xf16>
%c = arith.addf %a, %b : vector<16x32xf16>
- //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16]>>
+ //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
xegpu.store_nd %c, %c_tdesc[0, %k] : vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16>
}
gpu.return
@@ -100,15 +100,15 @@ gpu.module @test_kernel {
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf16> -> !xegpu.tensor_desc<12x32xf16>
scf.for %k = %c0 to %c1024 step %c32 {
- //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [4, 16]>}> :
- //CHECK-SAME: !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16]>> -> vector<12x32xf16>
+ //CHECK: xegpu.load_nd {{.*}} <{layout = #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+ //CHECK-SAME: !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<12x32xf16>
%a = xegpu.load_nd %a_tdesc[0, %k] : !xegpu.tensor_desc<12x32xf16> -> vector<12x32xf16>
%b = xegpu.load_nd %b_tdesc[0, %k] : !xegpu.tensor_desc<12x32xf16> -> vector<12x32xf16>
- //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [4, 16]>} : vector<12x32xf16>
+ //CHECK-COUNT: arith.addf {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<12x32xf16>
%c = arith.addf %a, %b : vector<12x32xf16>
- //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<12x32xf16>, !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16]>>
+ //CHECK-COUNT: xegpu.store_nd {{.*}} : vector<12x32xf16>, !xegpu.tensor_desc<12x32xf16, #xegpu.layout<inst_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
xegpu.store_nd %c, %c_tdesc[0, %k] : vector<12x32xf16>, !xegpu.tensor_desc<12x32xf16>
}
gpu.return
@@ -119,7 +119,7 @@ gpu.module @test_kernel {
gpu.module @test {
// CHECK-LABEL: func.func @store_matrix(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: !xegpu.mem_desc<16x64xf16>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16]>} dense<0.000000e+00> : vector<16x16xf16>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x16xf16>
func.func @store_matrix(%arg0: !xegpu.mem_desc<16x64xf16>) {
%cst = arith.constant dense<0.0000> : vector<16x16xf16>
xegpu.store_matrix %cst, %arg0[8, 8]: vector<16x16xf16>, !xegpu.mem_desc<16x64xf16>
@@ -216,7 +216,7 @@ func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
// -----
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]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
+// 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>
@@ -300,8 +300,8 @@ func.func @vector_shape_cast_expand_and_merge(%arg0: memref<256xf16>, %arg1: mem
// -----
gpu.module @test{
// CHECK-LABEL: load_store_matrix
- // CHECK: xegpu.load_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1]>}>
- // CHECK: xegpu.store_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1]>}>
+ // CHECK: xegpu.load_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1], lane_layout = [1, 1], lane_data = [1, 1]>}>
+ // CHECK: xegpu.store_matrix %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 1], lane_layout = [1, 1], lane_data = [1, 1]>}>
func.func @load_store_matrix(%arg0: !xegpu.mem_desc<64x128xf32>, %arg1: i1) {
%c0 = arith.constant 0 : index
scf.if %arg1 {
@@ -315,11 +315,11 @@ gpu.module @test{
// -----
gpu.module @test{
// CHECK-LABEL: broadcast_both_leadingdims_innerdims
- // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16]>} dense<true> : vector<2x2x6x32xi1>
- // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16]>} dense<1.000000e+00> : vector<2x2x6x32xf32>
- // CHECK: vector.step {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1]>, dims = [0, 1]>, dims = [1]>} : vector<6xindex>
- // CHECK: vector.shape_cast {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1]>, dims = [0, 1]>} : vector<6xindex> to vector<6x1xindex>
- // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16]>} : vector<6x1xindex> to vector<2x2x6x32xindex>
+ // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>} dense<true> : vector<2x2x6x32xi1>
+ // CHECK: arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>} dense<1.000000e+00> : vector<2x2x6x32xf32>
+ // CHECK: vector.step {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1]>, dims = [1]>} : vector<6xindex>
+ // CHECK: vector.shape_cast {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1, 1], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1]>} : vector<6xindex> to vector<6x1xindex>
+ // CHECK: vector.broadcast {{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 1, 16], lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>} : vector<6x1xindex> to vector<2x2x6x32xindex>
gpu.func @broadcast_both_leadingdims_innerdims(%arg0: memref<32x2x192xf32>, %arg1: memref<32x2x192xf32>, %arg2: memref<32x2x192xf32>) kernel attributes {known_block_size = array<i32: 768, 1, 1>, known_grid_size = array<i32: 16, 1, 1>} {
%cst = arith.constant dense<true> : vector<2x2x6x32xi1>
%cst_0 = arith.constant dense<1.000000e+00> : vector<2x2x6x32xf32>
@@ -350,7 +350,7 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
%mask = arith.constant dense<true> : vector<32x32xi1>
%data = arith.constant dense<0.0> : vector<32x32xf32>
- // CHECK: xegpu.store {{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [32, 32]>{{.*}}}> :
+ // CHECK: xegpu.store {{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>{{.*}}}> :
xegpu.store %data, %ptr_i64[%1], %mask {
layout = #xegpu.layout<inst_data = [32, 32]>
} : vector<32x32xf32>, i64, vector<32x32xindex>, vector<32x32xi1>
@@ -386,9 +386,9 @@ func.func @bitcast_ui8_to_f4(%arg0: memref<256x16xui8>) {
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui16_to_f4(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x16xui16>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16]>> -> vector<256x16xui16>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x16xui16>
// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 64]>} : vector<256x16xui16> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
>From 9b857eb60cf15df141eea6e16e77b5d2f76def90 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 05:09:00 +0000
Subject: [PATCH 12/42] load_gather setup: take consumer's inst_data as-is and
consumer's lane info when present. fall back to scatter-store default lane
factorization. updates 2 stale test cases.
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 62 +++++++++++++------
.../XeGPU/propagate-layout-inst-data.mlir | 18 +++---
2 files changed, 51 insertions(+), 29 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ad591625754ae..5f50e84e94f6a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1307,15 +1307,24 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
// return {laneLayout, laneData};
// }
+// Forward declaration: defined later in this file.
+static std::pair<SmallVector<int>, SmallVector<int>>
+computeScatterStoreLaneLayoutAndData(ArrayRef<int64_t> srcShape,
+ int subgroupSize, int64_t maxChunkSize);
+
/// Sets up the anchor layout for load gather and load matrix operation.
/// load matrix lowers to load gather and 1d block load. All of them share the
/// same layout setup logic.
+///
/// For Subgroup layout, uses the consumer layout directly.
-/// For InstData layout, the innermost inst_data is taken directly from the
-/// consumer's inst_data, capped by `maxChunkSize * subgroupSize`.
-/// For Lane layout, lane_layout/lane_data are derived via
-/// `computeScatterLoadLaneLayoutAndData` using the consumer's
-/// lane_data[innermost] as the per-lane vector hint.
+///
+/// For InstData layout, takes consumer's inst_data as-is. lane_layout and
+/// lane_data are taken from the consumer when present; otherwise the helper
+/// derives the standard scatter-style default (subgroupSize lanes on the
+/// innermost dim, per-lane vector capped by maxChunkSize).
+///
+/// For Lane layout, lane_layout/lane_data are taken from the consumer when
+/// present; otherwise derived from the same default.
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
@@ -1326,28 +1335,41 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
SmallVector<int64_t> consumerInstData =
consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> consumerLaneData =
- consumerLayout.getEffectiveLaneDataAsInt();
SmallVector<int64_t> consumerLaneLayout =
consumerLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> laneLayout(resShape.size(), 1);
- SmallVector<int64_t> laneData(resShape.size(), 1);
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+
+ // Pick lane_layout / lane_data: prefer consumer's, fall back to the
+ // scatter-store default (subgroupSize lanes on innermost dim, per-lane
+ // vector capped by maxChunkSize).
+ SmallVector<int64_t> laneLayout;
+ SmallVector<int64_t> laneData;
+ if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
+ laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
+ laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
+ } else {
+ auto [defLaneLayout, defLaneData] = computeScatterStoreLaneLayoutAndData(
+ resShape, subgroupSize, maxChunkSize);
+ laneLayout.assign(defLaneLayout.begin(), defLaneLayout.end());
+ laneData.assign(defLaneData.begin(), defLaneData.end());
+ }
if (layoutKind == xegpu::LayoutKind::InstData) {
- SmallVector<int64_t> instData(resShape.size(), 1);
- laneData.back() = std::min(static_cast<int64_t>(consumerLaneData.back()),
- int64_t(maxChunkSize));
- laneLayout.back() = consumerLaneLayout.back();
- instData.back() = laneData.back() * laneLayout.back();
+ // Take consumer's inst_data as-is. If the consumer doesn't have one,
+ // fall back to lane_layout * lane_data per dim.
+ SmallVector<int64_t> instData;
+ if (!consumerInstData.empty()) {
+ instData.assign(consumerInstData.begin(), consumerInstData.end());
+ } else {
+ instData.resize(resShape.size());
+ for (size_t i = 0; i < resShape.size(); ++i)
+ instData[i] = laneLayout[i] * laneData[i];
+ }
return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
- if (layoutKind == xegpu::LayoutKind::Lane) {
-
- laneData.back() = std::min(static_cast<int64_t>(consumerLaneData.back()),
- int64_t(maxChunkSize));
- laneLayout.back() = consumerLaneLayout.back();
+ if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
- }
return nullptr;
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index da4a2e32259de..152543f338083 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -133,11 +133,11 @@ func.func @store_matrix(%arg0: !xegpu.mem_desc<16x64xf16>) {
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 = [16]>} dense<true> : vector<16x32xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16]>} dense<12> : vector<16x32xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16]>}> :
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<true> : vector<16x32xi1>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<12> : vector<16x32xindex>
+// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 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 = [16, 16]>}> :
+// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
func.func @scatter_ops_coalesce_chunksize(%src: memref<512xf32>) {
%1 = arith.constant dense<1>: vector<16x32xi1>
@@ -179,11 +179,11 @@ func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: m
gpu.module @test {
// CHECK-LABEL: func.func @insert_strided_slice_inst_data_no_packing(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf32>) {
-// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<1.000000e+00> : vector<4x16xf32>
-// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<8x32xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>, offsets = [0, 0], strides = [1, 1]} : vector<4x16xf32> into vector<8x32xf32>
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.store_nd %[[INSERT]], %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<1.000000e+00> : vector<4x16xf32>
+// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<8x32xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, offsets = [0, 0], strides = [1, 1]} : vector<4x16xf32> into vector<8x32xf32>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[INSERT]], %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
func.func @insert_strided_slice_inst_data_no_packing(%arg0: memref<8x32xf32>) {
%c0 = arith.constant 0 : index
%cst_small = arith.constant dense<1.0> : vector<4x16xf32>
>From 63f6d96babd30febbc5b0de16d7b7ca258af7d8b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 05:34:36 +0000
Subject: [PATCH 13/42] createScaleLayout: cap scale lane_layout by inst_data
and derive lane_data so the dpas_mx scale operand satisfies Cat A invariant.
setupGenericNdAnchorLayout: fall back to lane_layout * lane_data when
block-WHC lookup yields no divisor (sub-byte float types).
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 41 ++++++++++++++-----
1 file changed, 31 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 5f50e84e94f6a..0bc962bb868d1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1714,19 +1714,26 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
consumerLayout.isForWorkgroup())
return consumerLayout;
- // Compute the default inst_data from hardware block params.
+ // Compute inst_data from hardware block params. For Nd ops, the lane
+ // factorization above (laneLayout / laneData) is rigid; inst_data must be
+ // a multiple of lane_layout * lane_data on each dim (Category A
+ // invariant). If block params are unavailable for this element type
+ // (e.g. sub-byte floats with no uArch entry), fall back to
+ // lane_layout * lane_data (k = 1).
SmallVector<int64_t> instData(rank, 1);
int instWidth = xegpu::getLargestDivisor(
static_cast<int>(dataShape.back()), bWidths);
if (instWidth == -1)
- return nullptr;
- instData.back() = instWidth;
+ instData.back() = laneLayout.back() * laneData.back();
+ else
+ instData.back() = instWidth;
if (rank >= 2) {
int instHeight = xegpu::getLargestDivisor(
static_cast<int>(dataShape[rank - 2]), bHeights);
if (instHeight == -1)
- return nullptr;
- instData[rank - 2] = instHeight;
+ instData[rank - 2] = laneLayout[rank - 2] * laneData[rank - 2];
+ else
+ instData[rank - 2] = instHeight;
}
// Honor the consumer's inst_data if it is uArch-valid.
@@ -2217,15 +2224,29 @@ createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
scaleLaneLayout.assign(laneLayout.begin(), laneLayout.end());
scaleLaneData.assign(laneData.begin(), laneData.end());
bool isRowMajor = uArchInstruction->isLaneLayoutRowMajorOrder();
- if (isBScale ^ isRowMajor) {
+ if (isBScale ^ isRowMajor)
std::swap(scaleLaneLayout[rank - 2], scaleLaneLayout[rank - 1]);
+ // Cap lane_layout by the per-instruction tile (inst_data) on each dim.
+ // Then derive lane_data = inst_data / lane_layout so the Category A
+ // invariant inst_data = lane_layout * lane_data * k (with k = 1) holds
+ // for the scale operand's load_nd consumer.
+ if (!scaleInstData.empty()) {
+ for (int64_t d = rank - 2; d < rank; ++d) {
+ scaleLaneLayout[d] =
+ std::min<int64_t>(scaleInstData[d], scaleLaneLayout[d]);
+ scaleLaneData[d] = std::max<int64_t>(
+ scaleInstData[d] / std::max<int64_t>(scaleLaneLayout[d], 1), 1);
+ }
+ } else {
+ // No inst_data on the matrix layout; fall back to capping by scale
+ // shape and deriving lane_data from it (legacy behavior).
scaleLaneLayout[rank - 2] =
std::min<int64_t>(scaleShape[rank - 2], scaleLaneLayout[rank - 2]);
+ scaleLaneData[rank - 2] = std::max<int64_t>(
+ scaleShape[rank - 2] / scaleLaneLayout[rank - 2], 1);
+ scaleLaneData[rank - 1] = std::max<int64_t>(
+ scaleShape[rank - 1] / scaleLaneLayout[rank - 1], 1);
}
- scaleLaneData[rank - 2] =
- std::max<int64_t>(scaleShape[rank - 2] / scaleLaneLayout[rank - 2], 1);
- scaleLaneData[rank - 1] =
- std::max<int64_t>(scaleShape[rank - 1] / scaleLaneLayout[rank - 1], 1);
}
return xegpu::LayoutAttr::get(
context,
>From 78daa3e7548449023844461267d6947fba02fb69 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 05:58:13 +0000
Subject: [PATCH 14/42] test: update case 19 (dpas_mx_f4e2m1) check lines.
load_nd anchors carry full inst_data + lane info; load_nd's inst_data is
allowed to be larger than dpas_mx's per-instruction tile (multiple-of
relation).
---
.../XeGPU/propagate-layout-inst-data.mlir | 32 +++++++++----------
1 file changed, 16 insertions(+), 16 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 152543f338083..692d61f9797fe 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -452,24 +452,24 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M
// CHECK-LABEL: func.func @dpas_mx_f4e2m1
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x128xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<128x32xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x4xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<4x32xf8E8M0FNU>
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<16x32xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x128xf4E2M1FN> -> !xegpu.tensor_desc<16x128xf4E2M1FN, #xegpu.layout<inst_data = [8, 64]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<128x32xf4E2M1FN> -> !xegpu.tensor_desc<128x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 64]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x128xf4E2M1FN, #xegpu.layout<inst_data = [8, 64]>> -> vector<16x128xf4E2M1FN>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<128x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16]>> -> vector<128x32xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x4xf8E8M0FNU> -> !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [8, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [8, 2]>> -> vector<16x4xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16]>> -> vector<4x32xf8E8M0FNU>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x128xf4E2M1FN> -> !xegpu.tensor_desc<16x128xf4E2M1FN>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<128x32xf4E2M1FN> -> !xegpu.tensor_desc<128x32xf4E2M1FN>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x128xf4E2M1FN> -> vector<16x128xf4E2M1FN>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<128x32xf4E2M1FN> -> vector<128x32xf4E2M1FN>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x4xf8E8M0FNU> -> !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2]>> -> vector<16x4xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [4, 32], lane_layout = [1, 16], lane_data = [2, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [4, 32], lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [4, 32], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<4x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
-// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
func.func @dpas_mx_f4e2m1(%arg0: memref<16x128xf4E2M1FN>, %arg1: memref<128x32xf4E2M1FN>, %arg2: memref<16x32xbf16>,
%arg3: memref<16x4xf8E8M0FNU>, %arg4: memref<4x32xf8E8M0FNU>) {
>From 5225fe7f9bd70a1356174cae4341103f3e67afdb Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 05:59:57 +0000
Subject: [PATCH 15/42] test: update case 6
(load_gather_with_coalesce_chunksize) check lines. mask/offset constants are
2D vector<16x16> matching the source IR (not 1D as previous CHECK assumed);
load_gather anchor includes full inst_data + lane info.
---
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 692d61f9797fe..b0b7af8e801eb 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -155,10 +155,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<lane_layout = [16], lane_data = [1]>}
-// CHECK-SAME: dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
-// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK-SAME: dense<0> : vector<16x16xindex>
+// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], 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, 16], lane_layout = [16, 1], lane_data = [1, 2]>}> : 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>
>From d10c559f68b0d6516313dd047468c1c30b463a73 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 06:08:36 +0000
Subject: [PATCH 16/42] test: update case 12
(vector_shape_cast_expand_and_merge) check lines. layouts now carry full
inst_data + lane_layout + lane_data; lane_data = 2 reflects scatter-load's
per-lane chunk for f16.
---
.../Dialect/XeGPU/propagate-layout-inst-data.mlir | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index b0b7af8e801eb..c8525eefde1e4 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -277,13 +277,13 @@ 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]>} dense<true> : vector<256xi1>
-// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [32]>} : vector<256xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32]>}> : 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]>} : vector<256xf16> to vector<2x4x32xf16>
-// CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<inst_data = [1, 32]>} : vector<2x4x32xf16> to vector<1x256xf16>
-// CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<inst_data = [32]>} : vector<1x256xf16> to vector<256xf16>
-// CHECK: xegpu.store %[[CAST_2]], %arg1[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
+// 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: %[[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>
+// CHECK: xegpu.store %[[CAST_2]], %arg1[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
func.func @vector_shape_cast_expand_and_merge(%arg0: memref<256xf16>, %arg1: memref<256xf16>) {
%cst = arith.constant dense<true> : vector<256xi1>
%0 = vector.step : vector<256xindex>
>From 3159743df131815e4cde47a405b2198a43d12484 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 06:12:33 +0000
Subject: [PATCH 17/42] test: update case 18 (dpas_mx_f8e5m2) check lines.
load_nd anchors carry full inst_data + lane info; load_nd inst_data may
exceed dpas_mx per-instruction tile (multiple-of relation).
---
.../XeGPU/propagate-layout-inst-data.mlir | 32 +++++++++----------
1 file changed, 16 insertions(+), 16 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index c8525eefde1e4..2d4e22f69dd11 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -410,24 +410,24 @@ func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
// CHECK-LABEL: func.func @dpas_mx_f8e5m2
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x64xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<64x32xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x2xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<2x32xf8E8M0FNU>
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0.000000e+00> : vector<16x32xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x64xf8E5M2> -> !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x32xf8E5M2> -> !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 16]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 32]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32]>> -> vector<16x64xf8E5M2>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 16]>> -> vector<64x32xf8E5M2>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x2xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [8, 1]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [8, 1]>> -> vector<16x2xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [1, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16]>> -> vector<2x32xf8E8M0FNU>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x64xf8E5M2> -> !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x32xf8E5M2> -> !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [4, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<16x64xf8E5M2>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [4, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<64x32xf8E5M2>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x2xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<16x2xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 32], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 32], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 32], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
-// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
-// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
+// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M2>, %arg2: memref<16x32xbf16>,
%arg3: memref<16x2xf8E8M0FNU>, %arg4: memref<2x32xf8E8M0FNU>) {
>From a2098b60827a5f678467beb55146495b1a3db97c Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 06:23:47 +0000
Subject: [PATCH 18/42] setupGenericNdAnchorLayout: cap default lane_data and
lane_layout by innermost shape so the lane factorization always fits the
tensor (e.g., ui8 with shape innermost=16 used full packing factor 2 giving
lane_product=32 > 16 and produced an invalid layout). also update case 16
check lines.
---
.../Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 16 +++++++++++++---
.../XeGPU/propagate-layout-inst-data.mlir | 6 +++---
2 files changed, 16 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 0bc962bb868d1..d3244fb677c79 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1667,13 +1667,23 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
int rank = dataShape.size();
assert(rank >= 1 && "Expected at least 1D shape for ND op");
- // Compute the default 2D block IO lane layout / lane data.
+ // Compute the default 2D block IO lane layout / lane data. Cap each by
+ // the innermost shape so the product `lane_layout * lane_data` doesn't
+ // exceed it (e.g. ui8 with shape innermost=16 cannot use the full
+ // packing factor of 2 because subgroupSize * 2 = 32 > 16).
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
SmallVector<int64_t> laneLayout(rank, 1);
SmallVector<int64_t> laneData(rank, 1);
- laneLayout.back() = uArch->getSubgroupSize();
- laneData.back() = packingFactor;
+ int64_t innermostShape = dataShape.back();
+ int64_t lanesOnInnermost =
+ std::min<int64_t>(uArch->getSubgroupSize(), innermostShape);
+ laneLayout.back() = lanesOnInnermost;
+ if (lanesOnInnermost > 0)
+ laneData.back() = std::min<int64_t>(packingFactor,
+ innermostShape / lanesOnInnermost);
+ else
+ laneData.back() = 1;
// Honor the consumer's lane info when valid. The consumer (e.g. dpas) may
// require VNNI-style packing on a non-innermost dim that the load's own
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 2d4e22f69dd11..ae7310bf6fadb 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -363,9 +363,9 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui8_to_f4(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x16xui8>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui8> -> !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16]>> -> vector<256x16xui8>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui8> -> !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x16xui8>
// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32]>} : vector<256x16xui8> to vector<256x32xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
>From 349bd34b3c4c501c6e01f277f8341e9599ccbc82 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 06:31:29 +0000
Subject: [PATCH 19/42] test: fix case 6 (load_gather_with_coalesce_chunksize).
source memref to 1D <256xf16>; offset literal to dense<0> splat (matching 2D
type); CHECK lines updated for 2D mask/offset layouts.
---
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 15 +++++++--------
1 file changed, 7 insertions(+), 8 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index c04e949015368..48891079f85a7 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -134,19 +134,18 @@ 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<16x16xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
-// CHECK-SAME: dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16xindex>
-// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<16x16xf16>, vector<16xindex>, vector<16xi1> -> vector<16x16xf16>
-func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
+// 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>
+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>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %offset = arith.constant dense<[0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]> : vector<16x16xindex>
+ %offset = arith.constant dense<0> : vector<16x16xindex>
%mask = arith.constant dense<true> : vector<16x16xi1>
%3 = xegpu.load %arg1[%offset], %mask
- : memref<16x16xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+ : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
%4 = vector.transpose %3, [1, 0] : vector<16x16xf16> to vector<16x16xf16>
%5 = xegpu.dpas %1, %4 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
%6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
>From e8bbbf36f62a61054fd3f188a2980d401354d6ea Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 10 Jun 2026 06:34:12 +0000
Subject: [PATCH 20/42] test: fix case 7 (scatter_ops_coalesce_chunksize).
source memref to 1D <128xf16>; CHECK lines updated for [1, 8] / [1, 1] lane
factorization that propagation now produces.
---
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 18 +++++++++---------
1 file changed, 9 insertions(+), 9 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 48891079f85a7..e3aad3a3fc965 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -157,19 +157,19 @@ func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: m
// -----
gpu.module @test {
// CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x8xf16>) {
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16x8xi1>
-// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16x8xindex>
-// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}>
-// CHECK-SAME: memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
-// CHECK: xegpu.store %[[LOAD_VEC]], %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8]>}> : vector<16x8xf16>, memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1>
-func.func @scatter_ops_coalesce_chunksize(%src: memref<16x8xf16>) {
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<128xf16>) {
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>} dense<true> : vector<16x8xi1>
+// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>} dense<12> : vector<16x8xindex>
+// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>}>
+// CHECK-SAME: memref<128xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
+// CHECK: xegpu.store %[[LOAD_VEC]], %[[ARG0]][%[[OFFSETS]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>}> : vector<16x8xf16>, memref<128xf16>, vector<16x8xindex>, vector<16x8xi1>
+func.func @scatter_ops_coalesce_chunksize(%src: memref<128xf16>) {
%1 = arith.constant dense<1>: vector<16x8xi1>
%offset = arith.constant dense<12> : vector<16x8xindex>
%3 = xegpu.load %src[%offset], %1
- : memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
+ : memref<128xf16>, vector<16x8xindex>, vector<16x8xi1> -> vector<16x8xf16>
xegpu.store %3, %src[%offset], %1
- : vector<16x8xf16>, memref<16x8xf16>, vector<16x8xindex>, vector<16x8xi1>
+ : vector<16x8xf16>, memref<128xf16>, vector<16x8xindex>, vector<16x8xi1>
return
}
}
>From 7177f3288095529414d0216e78f65ebd87f6eadf Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 11 Jun 2026 03:21:12 +0000
Subject: [PATCH 21/42] remove debug print and polish
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 30 ++--
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 134 ++++++------------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 10 --
.../XeGPU/propagate-layout-inst-data.mlir | 2 +-
4 files changed, 59 insertions(+), 117 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index f3cbf6ae72f5d..22608261512e6 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -208,27 +208,25 @@ DistributeLayoutAttr setupInsertStridedSliceResultLayout(
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a load gather operation.
-DistributeLayoutAttr
-setupLoadGatherAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
- int chunkSize, DistributeLayoutAttr consumerLayout,
- const uArch::uArch *uArch);
+DistributeLayoutAttr setupLoadGatherAnchorLayout(
+ LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize,
+ DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for load matrix operation.
-DistributeLayoutAttr
-setupLoadMatrixAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
- int chunkSize, DistributeLayoutAttr consumerLayout,
- const uArch::uArch *uArch);
+DistributeLayoutAttr setupLoadMatrixAnchorLayout(
+ LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize,
+ DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
/// Sets up the anchor layout for a store scatter operation.
DistributeLayoutAttr setupStoreScatterAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
- int chunkSize,
+ int contigChunkSize,
const uArch::uArch *uArch);
/// Sets up the anchor layout for a store matrix operation.
DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
VectorType vectorTy,
- int chunkSize,
+ int contigChunkSize,
const uArch::uArch *uArch);
/// If the consumer layout has only inst_data (no lane_layout/lane_data),
@@ -245,8 +243,7 @@ completeLoadGatherLayoutFromInstData(DistributeLayoutAttr consumerLayout,
DistributeLayoutAttr
completeStoreScatterLayoutFromInstData(DistributeLayoutAttr consumerLayout,
- Type elemTy,
- const uArch::uArch *uArch);
+ Type elemTy, const uArch::uArch *uArch);
/// Sets up the anchor layout for a store_nd operation. StoreNd does not
/// consider a consumer layout (it is a data sink), and picks its layout from
@@ -269,11 +266,10 @@ DistributeLayoutAttr setupPrefetchNdAnchorLayout(LayoutKind layoutKind,
/// Otherwise defaults derived from uArch block parameters are used.
/// `consumerLayout` may be null. `numSg` is only used for Subgroup-kind
/// layouts when the consumer does not already provide an sg_layout.
-DistributeLayoutAttr setupLoadNdAnchorLayout(LayoutKind layoutKind,
- VectorType vectorTy,
- DistributeLayoutAttr consumerLayout,
- int numSg,
- const uArch::uArch *uArch);
+DistributeLayoutAttr
+setupLoadNdAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
+ DistributeLayoutAttr consumerLayout, int numSg,
+ const uArch::uArch *uArch);
/// Sets up the anchor layouts for a dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout creation.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d3244fb677c79..1421bc0cb9d17 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -718,21 +718,6 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// Use case 2: Dim split from source to result, for multi-stage reduction
SmallVector<SmallVector<int64_t>> splitDimGroups;
if (xegpu::matchSplitDimExpansion(srcShape, resShape, splitDimGroups)) {
- llvm::dbgs() << "[DEBUG inferShapeCastSourceLayout] case2 split. resLayout="
- << resLayout << " srcShape=[";
- for (auto s : srcShape)
- llvm::dbgs() << s << ",";
- llvm::dbgs() << "] resShape=[";
- for (auto s : resShape)
- llvm::dbgs() << s << ",";
- llvm::dbgs() << "] groups=";
- for (auto &g : splitDimGroups) {
- llvm::dbgs() << "{";
- for (auto d : g)
- llvm::dbgs() << d << ",";
- llvm::dbgs() << "} ";
- }
- llvm::dbgs() << "\n";
auto srcLayout = resLayout;
for (const auto &dimGroup : splitDimGroups)
srcLayout = srcLayout.collapseDims(dimGroup);
@@ -811,12 +796,10 @@ static bool leadingDimsAreUnit(ArrayRef<int64_t> shape, int numInnerDims) {
/// result layout can later be distributed without re-deriving the lane
/// layout. `instData`, `laneLayout`, and `laneData` may have different
/// element types; they are normalized to int32 entries.
-static xegpu::LayoutAttr
-buildInstDataLayoutWithLane(mlir::MLIRContext *context,
- ArrayRef<int64_t> instData,
- ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData,
- DenseI32ArrayAttr orderAttr = nullptr) {
+static xegpu::LayoutAttr buildInstDataLayoutWithLane(
+ mlir::MLIRContext *context, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
auto toI32Attr = [&](auto range) {
SmallVector<int32_t> v(range.begin(), range.end());
return DenseI32ArrayAttr::get(context, v);
@@ -1284,33 +1267,26 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
return requiredResLayout;
}
-// /// Computes lane_layout and lane_data for scatter-style load anchor layouts
-// /// (load gather, load matrix). Lanes and the per-lane vector both live on
-// the
-// /// innermost dim; the per-lane vector width is hinted by the consumer's
-// /// lane_data[innermost]:
-// /// - laneData[innermost] = min(consumerLaneData[innermost],
-// maxChunkSize)
-// /// - laneLayout[innermost] = min(subgroupSize,
-// /// resShape[innermost] /
-// laneData[innermost])
-// /// All other entries are 1.
-// static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-// computeScatterLoadLaneLayoutAndData(ArrayRef<int64_t> resShape,
-// ArrayRef<int64_t> consumerLaneData,
-// int subgroupSize, int64_t maxChunkSize) {
-// int rank = resShape.size();
-// SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
-// int innermost = rank - 1;
-// laneData[innermost] = std::min(consumerLaneData.back(), maxChunkSize);
-// laneLayout[innermost] = consumerLaneLayout.back();
-// return {laneLayout, laneData};
-// }
-
-// Forward declaration: defined later in this file.
+/// Computes lane_layout and lane_data for scatter-style store anchor layouts
+/// (store scatter, store matrix). Lanes and the per-lane vector both live on
+/// the innermost dim:
+/// - laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
+/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
+/// maxChunkSize)
+/// All other entries are 1.
static std::pair<SmallVector<int>, SmallVector<int>>
computeScatterStoreLaneLayoutAndData(ArrayRef<int64_t> srcShape,
- int subgroupSize, int64_t maxChunkSize);
+ int subgroupSize, int64_t maxChunkSize) {
+ int rank = srcShape.size();
+ SmallVector<int> laneLayout(rank, 1), laneData(rank, 1);
+ int innermost = rank - 1;
+ laneLayout[innermost] = std::min(static_cast<int>(subgroupSize),
+ static_cast<int>(srcShape[innermost]));
+ laneData[innermost] =
+ std::min(static_cast<int>(srcShape[innermost] / laneLayout[innermost]),
+ static_cast<int>(maxChunkSize));
+ return {laneLayout, laneData};
+}
/// Sets up the anchor layout for load gather and load matrix operation.
/// load matrix lowers to load gather and 1d block load. All of them share the
@@ -1375,7 +1351,7 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
/// Sets up the anchor layout for a load gather operation.
xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
- xegpu::LayoutKind layoutKind, VectorType resVecTy, int chunkSize,
+ xegpu::LayoutKind layoutKind, VectorType resVecTy, int contigChunkSize,
xegpu::DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1386,7 +1362,8 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
maxChunkSize, resShape, subgroupSize);
@@ -1396,7 +1373,7 @@ xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
/// TODO: enhance load matrix to indicate lowering to chunked load or not.
xegpu::DistributeLayoutAttr
xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType resVecTy, int chunkSize,
+ VectorType resVecTy, int contigChunkSize,
xegpu::DistributeLayoutAttr consumerLayout,
const xegpu::uArch::uArch *uArch) {
@@ -1408,33 +1385,12 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize =
- std::min(uArchInstruction->getMaxLaneLoadSize(elemBitWidth), chunkSize);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
maxChunkSize, resShape, subgroupSize);
}
-/// Computes lane_layout and lane_data for scatter-style store anchor layouts
-/// (store scatter, store matrix). Lanes and the per-lane vector both live on
-/// the innermost dim:
-/// - laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
-/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
-/// maxChunkSize)
-/// All other entries are 1.
-static std::pair<SmallVector<int>, SmallVector<int>>
-computeScatterStoreLaneLayoutAndData(ArrayRef<int64_t> srcShape,
- int subgroupSize, int64_t maxChunkSize) {
- int rank = srcShape.size();
- SmallVector<int> laneLayout(rank, 1), laneData(rank, 1);
- int innermost = rank - 1;
- laneLayout[innermost] = std::min(static_cast<int>(subgroupSize),
- static_cast<int>(srcShape[innermost]));
- laneData[innermost] =
- std::min(static_cast<int>(srcShape[innermost] / laneLayout[innermost]),
- static_cast<int>(maxChunkSize));
- return {laneLayout, laneData};
-}
-
/// Sets up the anchor layout for store scatter and store matrix operation.
/// store matrix lowers to store scatter and 1d block store. All of them
/// share the same layout setup logic. For Subgroup layout, not supported
@@ -1477,7 +1433,7 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
/// Sets up the anchor layout for a store scatter operation.
xegpu::DistributeLayoutAttr
xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int chunkSize,
+ VectorType srcVecTy, int contigChunkSize,
const uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1488,8 +1444,8 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize =
- std::min(uArchInstruction->getMaxLaneStoreSize(elemBitWidth), chunkSize);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize);
}
@@ -1497,7 +1453,7 @@ xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
/// Sets up the anchor layout for a store matrix operation.
xegpu::DistributeLayoutAttr
xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int chunkSize,
+ VectorType srcVecTy, int contigChunkSize,
const xegpu::uArch::uArch *uArch) {
const int subgroupSize = uArch->getSubgroupSize();
@@ -1508,8 +1464,8 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize =
- std::min(uArchInstruction->getMaxLaneStoreSize(elemBitWidth), chunkSize);
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
srcShape, subgroupSize);
@@ -1680,8 +1636,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
std::min<int64_t>(uArch->getSubgroupSize(), innermostShape);
laneLayout.back() = lanesOnInnermost;
if (lanesOnInnermost > 0)
- laneData.back() = std::min<int64_t>(packingFactor,
- innermostShape / lanesOnInnermost);
+ laneData.back() =
+ std::min<int64_t>(packingFactor, innermostShape / lanesOnInnermost);
else
laneData.back() = 1;
@@ -1731,8 +1687,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
// (e.g. sub-byte floats with no uArch entry), fall back to
// lane_layout * lane_data (k = 1).
SmallVector<int64_t> instData(rank, 1);
- int instWidth = xegpu::getLargestDivisor(
- static_cast<int>(dataShape.back()), bWidths);
+ int instWidth =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
if (instWidth == -1)
instData.back() = laneLayout.back() * laneData.back();
else
@@ -1780,9 +1736,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
return nullptr;
SmallVector<int> sgLayout = {static_cast<int>(sgLayouts[0].first),
static_cast<int>(sgLayouts[0].second)};
- SmallVector<int> sgData = {
- static_cast<int>(dataShape[0]) / sgLayout[0],
- static_cast<int>(dataShape[1]) / sgLayout[1]};
+ SmallVector<int> sgData = {static_cast<int>(dataShape[0]) / sgLayout[0],
+ static_cast<int>(dataShape[1]) / sgLayout[1]};
return xegpu::LayoutAttr::get(
context, DenseI32ArrayAttr::get(context, sgLayout),
DenseI32ArrayAttr::get(context, sgData),
@@ -1852,10 +1807,11 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
/// 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.
-xegpu::DistributeLayoutAttr xegpu::setupLoadNdAnchorLayout(
- xegpu::LayoutKind layoutKind, VectorType resVecTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
+xegpu::DistributeLayoutAttr
+xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType resVecTy,
+ xegpu::DistributeLayoutAttr consumerLayout,
+ int numSg, const xegpu::uArch::uArch *uArch) {
auto context = resVecTy.getContext();
Type elemTy = resVecTy.getElementType();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index cb9d6eeed9743..1d6297dcb8d6d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -666,18 +666,12 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
auto requiredResLayoutAttr = xegpu::setupMultiReductionResultLayout(
layoutKind, sourceTy, consumerLayoutAttr, reductionDims, numSg, uArch);
- llvm::dbgs() << "[DEBUG visitMultiRed] op=" << *reduction
- << " consumer=" << consumerLayoutAttr
- << " required=" << requiredResLayoutAttr << "\n";
-
xegpu::setTemporaryLayout(reduction->getResult(0), requiredResLayoutAttr);
// derive the source layout from the dominant layout and reduction dims
auto srcLayoutAttr = xegpu::inferMultiReductionSourceLayout(
requiredResLayoutAttr, reductionDims);
- llvm::dbgs() << "[DEBUG visitMultiRed] srcLayout=" << srcLayoutAttr << "\n";
-
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
// Accumulator should have the same layout as the result.
propagateIfChanged(operands[1],
@@ -746,10 +740,6 @@ void LayoutInfoPropagation::visitShapeCastOp(
xegpu::DistributeLayoutAttr srcLayoutAttr =
xegpu::inferShapeCastSourceLayout(resultLayoutAttr, resShape, srcShape);
- llvm::dbgs() << "[DEBUG visitShapeCast] op=" << *shapeCast
- << " consumer=" << resultLayoutAttr
- << " src=" << srcLayoutAttr << "\n";
-
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index ae7310bf6fadb..bb20c63b08bf1 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -280,7 +280,7 @@ gpu.module @test {
// 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: %[[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_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [2, 4, 4], lane_layout = [2, 4, 2], 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>
// CHECK: xegpu.store %[[CAST_2]], %arg1[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
>From 8bd0d2c545c5b0e921272971e57bc1f82e0fd728 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 12 Jun 2026 19:38:37 +0000
Subject: [PATCH 22/42] refactor common function to compute default layout, and
make computeCandidateSgLayouts support nd
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 13 +-
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 15 +-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 457 +++++++++---------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 69 ++-
4 files changed, 274 insertions(+), 280 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 22608261512e6..7b05bfece50cc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -233,17 +233,10 @@ DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
/// completes it by running the corresponding scatter-style Lane-kind setup
/// rule with inst_data as the destination shape. The resulting lane info is
/// merged with the consumer's inst_data so downstream setup* paths see a
-/// fully-populated layout. If the consumer already has lane info (or no
-/// inst_data), returns it unchanged. The "load-side" version uses
-/// setupLoadGatherAnchorLayout under the hood; the "store-side" version uses
-/// setupStoreScatterAnchorLayout.
+/// fully-populated layout.
DistributeLayoutAttr
-completeLoadGatherLayoutFromInstData(DistributeLayoutAttr consumerLayout,
- Type elemTy, const uArch::uArch *uArch);
-
-DistributeLayoutAttr
-completeStoreScatterLayoutFromInstData(DistributeLayoutAttr consumerLayout,
- Type elemTy, const uArch::uArch *uArch);
+completeScatterIOLaneLayoutFromInstData(DistributeLayoutAttr consumerLayout,
+ Type elemTy, const uArch::uArch *uArch);
/// Sets up the anchor layout for a store_nd operation. StoreNd does not
/// consider a consumer layout (it is a data sink), and picks its layout from
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index b780c66594eb0..6dd6c9cb585de 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -751,9 +751,16 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
expSgLayout = spread(origSgLayoutDim, targetShape, /*outerToInner=*/true);
splice(sgLayout, expSgLayout);
}
+ // When sg_data equals the full new-dim extent, the data is replicated across
+ // subgroups along this dim — each subgroup holds the whole extent rather
+ // than a partition. In that case sg_data spreads against the full extent and
+ // the per-sg view is the full extent too; otherwise both are divided by
+ // sg_layout.
+ bool sgDataReplicated =
+ hasSgData && origSgDataDim == computeProduct(targetShape);
if (hasSgData) {
SmallVector<int64_t> dimSizeCap(targetShape.begin(), targetShape.end());
- if (hasSgLayout)
+ if (hasSgLayout && !sgDataReplicated)
for (int64_t i = 0; i < expCount; ++i)
dimSizeCap[i] /= expSgLayout[i];
SmallVector<int64_t> expSgData =
@@ -762,10 +769,10 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
}
// Per-sg view used as the base for lane_layout / lane_data / inst_data:
- // targetShape[i] / sg_layout[i] when sg_layout is present, else
- // targetShape itself.
+ // targetShape[i] / sg_layout[i] when sg_layout is present (and not
+ // replicated), else targetShape itself.
SmallVector<int64_t> perSgShape(targetShape.begin(), targetShape.end());
- if (hasSgLayout)
+ if (hasSgLayout && !sgDataReplicated)
for (int64_t i = 0; i < expCount; ++i)
perSgShape[i] /= expSgLayout[i];
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 1421bc0cb9d17..58a81640af721 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -30,6 +30,7 @@
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
+#include <functional>
#include <numeric>
using namespace mlir;
@@ -1274,17 +1275,15 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
/// maxChunkSize)
/// All other entries are 1.
-static std::pair<SmallVector<int>, SmallVector<int>>
-computeScatterStoreLaneLayoutAndData(ArrayRef<int64_t> srcShape,
- int subgroupSize, int64_t maxChunkSize) {
- int rank = srcShape.size();
- SmallVector<int> laneLayout(rank, 1), laneData(rank, 1);
- int innermost = rank - 1;
- laneLayout[innermost] = std::min(static_cast<int>(subgroupSize),
- static_cast<int>(srcShape[innermost]));
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeScatterIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t maxChunkSize) {
+ int64_t rank = instShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int64_t innermost = rank - 1;
+ laneLayout[innermost] = std::min(subgroupSize, instShape[innermost]);
laneData[innermost] =
- std::min(static_cast<int>(srcShape[innermost] / laneLayout[innermost]),
- static_cast<int>(maxChunkSize));
+ std::min(instShape[innermost] / laneLayout[innermost], maxChunkSize);
return {laneLayout, laneData};
}
@@ -1325,10 +1324,8 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
} else {
- auto [defLaneLayout, defLaneData] = computeScatterStoreLaneLayoutAndData(
- resShape, subgroupSize, maxChunkSize);
- laneLayout.assign(defLaneLayout.begin(), defLaneLayout.end());
- laneData.assign(defLaneData.begin(), defLaneData.end());
+ auto [LaneLayout, LaneData] =
+ computeScatterIOLaneLayoutAndData(resShape, subgroupSize, maxChunkSize);
}
if (layoutKind == xegpu::LayoutKind::InstData) {
@@ -1396,7 +1393,7 @@ xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
/// share the same layout setup logic. For Subgroup layout, not supported
/// yet.
///
-/// Lane layout is derived first via `computeScatterStoreLaneLayoutAndData`;
+/// Lane layout is derived first via `computeScatterIOLaneLayoutAndData`;
/// inst_data is then the element-wise product lane_layout * lane_data.
static xegpu::DistributeLayoutAttr
setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
@@ -1409,23 +1406,17 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
}
- auto [laneLayout, laneData] = computeScatterStoreLaneLayoutAndData(
- srcShape, subgroupSize, maxChunkSize);
+ auto [laneLayout, laneData] =
+ computeScatterIOLaneLayoutAndData(srcShape, subgroupSize, maxChunkSize);
if (layoutKind == xegpu::LayoutKind::InstData) {
- SmallVector<int> instData(srcShape.size());
+ SmallVector<int64_t> instData(srcShape.size());
for (size_t i = 0; i < srcShape.size(); ++i)
- instData[i] = static_cast<int>(laneLayout[i] * laneData[i]);
- return xegpu::LayoutAttr::get(
- context, /*sg_layout=*/nullptr,
- /*sg_data=*/nullptr,
- /*inst_data=*/DenseI32ArrayAttr::get(context, instData),
- /*lane_layout=*/DenseI32ArrayAttr::get(context, laneLayout),
- /*lane_data=*/DenseI32ArrayAttr::get(context, laneData),
- /*order=*/nullptr);
+ instData[i] = laneLayout[i] * laneData[i];
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- return xegpu::LayoutAttr::get(context, laneLayout, laneData);
+ return buildLaneLayout(context, laneLayout, laneData);
}
return nullptr;
}
@@ -1471,13 +1462,22 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
srcShape, subgroupSize);
}
-/// If `consumerLayout` has inst_data set but no lane_layout/lane_data,
-/// derive a lane factorization by re-running the load-side Lane setup with
-/// inst_data as the destination shape, and merge the result back so the
-/// returned LayoutAttr carries inst_data + lane_layout + lane_data. This
-/// guarantees the lane factorization the downstream load setup will see is
-/// the same one its own Lane-kind setup would produce.
-xegpu::DistributeLayoutAttr xegpu::completeLoadGatherLayoutFromInstData(
+/// Completes a scatter IO layout by deriving lane_layout and lane_data from
+/// inst_data when they are missing. If `consumerLayout` already has both
+/// lane_layout and lane_data, or has no inst_data, the layout is returned
+/// unchanged.
+///
+/// When lane info is absent, this function uses inst_data as the effective
+/// shape and computes the standard scatter-style lane factorization:
+/// - laneLayout[innermost] = min(subgroupSize, inst_data[innermost])
+/// - laneData[innermost] = min(inst_data[innermost] /
+/// laneLayout[innermost],
+/// maxChunkSize)
+///
+/// The returned layout carries inst_data + lane_layout + lane_data, ensuring
+/// the lane factorization is consistent with what the downstream load/store
+/// scatter anchor setup would produce.
+xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
const xegpu::uArch::uArch *uArch) {
if (!consumerLayout)
@@ -1500,58 +1500,125 @@ xegpu::DistributeLayoutAttr xegpu::completeLoadGatherLayoutFromInstData(
return consumerLayout;
int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
- auto laneOnly = setupGenericLoadAnchorLayout(
- xegpu::LayoutKind::Lane, context, /*consumerLayout=*/nullptr,
- maxChunkSize, instData, subgroupSize);
- if (!laneOnly)
- return consumerLayout;
+ auto [defLaneLayout, defLaneData] =
+ computeScatterIOLaneLayoutAndData(instData, subgroupSize, maxChunkSize);
- SmallVector<int64_t> laneLayout = laneOnly.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> laneData = laneOnly.getEffectiveLaneDataAsInt();
- return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ return buildInstDataLayoutWithLane(context, instData, defLaneLayout,
+ defLaneData);
}
-/// If `consumerLayout` has inst_data set but no lane_layout/lane_data,
-/// derive a lane factorization by re-running the store-side Lane setup with
-/// inst_data as the destination shape, and merge the result back. Returned
-/// LayoutAttr carries inst_data + lane_layout + lane_data.
-xegpu::DistributeLayoutAttr xegpu::completeStoreScatterLayoutFromInstData(
- xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
- const xegpu::uArch::uArch *uArch) {
- if (!consumerLayout)
- return consumerLayout;
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- if (instData.empty())
- return consumerLayout;
- if (!consumerLayout.getEffectiveLaneLayoutAsInt().empty() &&
- !consumerLayout.getEffectiveLaneDataAsInt().empty())
- return consumerLayout;
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t bitwidth,
+ int64_t packingSize, bool vnni = false) {
+ int64_t rank = instShape.size();
+ assert(leadingDimsAreUnit(instShape, 2));
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int64_t packingDim = vnni ? rank - 2 : rank - 1;
+ laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
+ laneLayout.back() = subgroupSize;
+ // assert that the lane layout and data fit in the inst shape
+ for (int64_t i = 0; i < rank; ++i) {
+ int64_t laneProduct = laneLayout[i] * laneData[i];
+ assert(instShape[i] % laneProduct == 0 &&
+ "lane_layout * lane_data must evenly divide the inst shape");
+ (void)laneProduct;
+ }
+ return {laneLayout, laneData};
+}
- const int subgroupSize = uArch->getSubgroupSize();
- auto *context = consumerLayout.getContext();
- auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- if (!uArchInstruction)
- return consumerLayout;
- int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+// Forward declaration: defined later in the file.
+using LayoutRepresentation = SmallVector<int64_t>;
+
+/// Enumerates all ways to split `total` into `rank` factors whose product
+/// equals `total`. Returns the list of all such factorizations.
+static SmallVector<LayoutRepresentation> enumerateFactorizations(int64_t total,
+ int64_t rank) {
+ SmallVector<LayoutRepresentation> results;
+ SmallVector<int64_t> current(rank, 0);
+
+ // Returns all divisors of `n` in ascending order.
+ auto getDivisors = [](int64_t n) {
+ SmallVector<int64_t> divs;
+ for (int64_t i = 1; i * i <= n; ++i) {
+ if (n % i == 0) {
+ divs.push_back(i);
+ if (i != n / i)
+ divs.push_back(n / i);
+ }
+ }
+ llvm::sort(divs);
+ return divs;
+ };
- auto laneOnly = setupGenericStoreAnchorLayout(
- xegpu::LayoutKind::Lane, context, maxChunkSize, instData, subgroupSize);
- if (!laneOnly)
- return consumerLayout;
+ std::function<void(int64_t, int64_t)> generate = [&](int64_t dim,
+ int64_t remaining) {
+ if (dim == rank - 1) {
+ current[dim] = remaining;
+ results.push_back(LayoutRepresentation(current));
+ return;
+ }
+ for (int64_t factor : getDivisors(remaining)) {
+ current[dim] = factor;
+ generate(dim + 1, remaining / factor);
+ }
+ };
- SmallVector<int64_t> laneLayout = laneOnly.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> laneData = laneOnly.getEffectiveLaneDataAsInt();
- return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ generate(0, total);
+ return results;
}
-// Forward declaration: defined later in the file.
-using LayoutRepresentation = std::pair<int64_t, int64_t>;
+// Computes all valid N-dimensional sg_layout candidates for the given
+// sgCount, whose sgData (= wgShape / sgLayout):
+// 1. Evenly divides wgShape (i.e., wgShape[d] % sgLayout[d] == 0).
+// 2. Is a multiple of instData (i.e., sgData[d] % instData[d] == 0).
+// Results are sorted by balance (smallest max-min spread first), with
+// lexicographic order as a tiebreaker.
+//
+// Example (2D):
+// wgShape = [128, 64], instData = [8, 16], sgCount = 32
+// Returns: [[8,4], [16,2]], corresponding to sgData [16,16] and [8,32].
static SmallVector<LayoutRepresentation>
-getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
- int64_t sgCount);
+computeCandidateSgLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
+ int64_t sgCount) {
+ int64_t rank = wgShape.size();
+ assert(rank > 0 && "wgShape must be non-empty");
+ assert(static_cast<int64_t>(instData.size()) == rank &&
+ "instData rank must match wgShape rank");
+
+ // Step 1: Get all N-D factorizations of sgCount.
+ auto allFactorizations = enumerateFactorizations(sgCount, rank);
+
+ // Step 2: Filter to keep only valid candidates.
+ SmallVector<LayoutRepresentation> candidates;
+ for (const auto &sgLayout : allFactorizations) {
+ bool valid = true;
+ for (int64_t dim = 0; dim < rank; ++dim) {
+ if (wgShape[dim] % sgLayout[dim] != 0) {
+ valid = false;
+ break;
+ }
+ int64_t sgData = wgShape[dim] / sgLayout[dim];
+ if (sgData % instData[dim] != 0) {
+ valid = false;
+ break;
+ }
+ }
+ if (valid)
+ candidates.push_back(sgLayout);
+ }
+
+ // Step 3: Sort by balance (smallest max-min spread), then lexicographic.
+ llvm::sort(candidates, [](const LayoutRepresentation &lhs,
+ const LayoutRepresentation &rhs) {
+ int64_t spreadLhs = *llvm::max_element(lhs) - *llvm::min_element(lhs);
+ int64_t spreadRhs = *llvm::max_element(rhs) - *llvm::min_element(rhs);
+ if (spreadLhs != spreadRhs)
+ return spreadLhs < spreadRhs;
+ return lhs < rhs;
+ });
+ return candidates;
+}
/// Validates whether `instData` is a hardware-viable inst_data for an ND op
/// with the given block params and lane factor. Specifically:
@@ -1596,6 +1663,26 @@ static bool isValidNdInstData(ArrayRef<int64_t> instData,
return true;
}
+static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
+ ArrayRef<int64_t> wgTileShape,
+ ArrayRef<int64_t> sgLayout,
+ int dimK = -1,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ SmallVector<int> sgData(sgLayout.size());
+ SmallVector<int> sgLayoutInt(sgLayout.begin(), sgLayout.end());
+ for (int dim = 0; dim < sgLayout.size(); ++dim) {
+ if (dim == dimK)
+ sgData[dim] = wgTileShape[dim];
+ else
+ sgData[dim] = static_cast<int>(wgTileShape[dim]) / sgLayout[dim];
+ }
+ return xegpu::LayoutAttr::get(context,
+ DenseI32ArrayAttr::get(context, sgLayoutInt),
+ DenseI32ArrayAttr::get(context, sgData),
+ /*inst_data=*/nullptr, /*lane_layout=*/nullptr,
+ /*lane_data=*/nullptr, /*order=*/nullptr);
+}
+
/// Generic anchor-layout setup for ND ops (load_nd, store_nd, prefetch_nd).
///
/// Given hardware-supported block widths/heights, picks the largest divisor
@@ -1613,7 +1700,7 @@ static bool isValidNdInstData(ArrayRef<int64_t> instData,
/// inst_data when it is uArch-valid.
/// For Subgroup kind: if the consumer specifies a workgroup-level layout,
/// reuses it directly; otherwise picks the most balanced sg_layout via
-/// `getValidLayouts` (requires `numSg`).
+/// `computeCandidateSgLayouts` (requires `numSg`).
static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
ArrayRef<int64_t> dataShape, Type elemTy, ArrayRef<int> bWidths,
@@ -1725,24 +1812,15 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
orderAttr);
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- if (rank != 2)
- return nullptr;
// The consumer-fast-path above already returned for the case where the
// consumer carries a workgroup-level layout, so here numSg must be set.
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
- auto sgLayouts = getValidLayouts(dataShape, instData, numSg);
+ auto sgLayouts = computeCandidateSgLayouts(dataShape, instData, numSg);
if (sgLayouts.empty())
return nullptr;
- SmallVector<int> sgLayout = {static_cast<int>(sgLayouts[0].first),
- static_cast<int>(sgLayouts[0].second)};
- SmallVector<int> sgData = {static_cast<int>(dataShape[0]) / sgLayout[0],
- static_cast<int>(dataShape[1]) / sgLayout[1]};
- return xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgData),
- /*inst_data=*/nullptr, /*lane_layout=*/nullptr,
- /*lane_data=*/nullptr, /*order=*/nullptr);
+ return buildSgLayout(context, dataShape, sgLayouts.front(), /*dimK=*/-1,
+ orderAttr);
}
return nullptr;
@@ -1856,84 +1934,6 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout, numSg, uArch);
}
-// Returns the default (lane_layout, lane_data) pair for a given 1D/2D vector
-// type used by 2D block IO ops.
-// - `packingSize` means multiple consecutive elements can be accessed
-// together as a single unit.
-// - `vnni` means data packing is column-wise (i.e., 2x1xf16 with vnni vs.
-// 1x2xf16 w/o vnni).
-template <typename RankedTy>
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-getDefaultLaneLayoutAndData2DBlockIo(
- RankedTy ty, const xegpu::uArch::uArch *uArch,
- std::optional<unsigned> packingSize = std::nullopt, bool vnni = false) {
- // Expecting at least 1D vector. For rank > 2, leading dims are batch dims.
- assert(((ty.getRank() >= 1 && !vnni) || ty.getRank() >= 2) &&
- "Expected at least 1D non-vnni or 2D vector.");
- // Expecting int or float element type.
- assert(ty.getElementType().isIntOrFloat() &&
- "Expected int or float element type.");
-
- auto rank = ty.getRank();
- SmallVector<int64_t> laneLayout(rank, 1);
- SmallVector<int64_t> laneData(rank, 1);
- if (packingSize.has_value()) {
- unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
- int64_t &laneDataPos = vnni ? laneData[rank - 2] : laneData.back();
- laneDataPos = bitwidth < *packingSize ? *packingSize / bitwidth : 1;
- }
- laneLayout.back() = uArch->getSubgroupSize();
- return {laneLayout, laneData};
-}
-
-// Convenience wrapper: returns a LayoutAttr carrying only the default lane
-// layout / lane data for a 2D block IO vector type.
-template <typename RankedTy>
-static xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo(
- RankedTy ty, const xegpu::uArch::uArch *uArch,
- std::optional<unsigned> packingSize = std::nullopt, bool vnni = false) {
- auto [laneLayout, laneData] =
- getDefaultLaneLayoutAndData2DBlockIo(ty, uArch, packingSize, vnni);
- return buildLaneLayout(ty.getContext(), laneLayout, laneData);
-}
-
-// This function returns all layouts for the given sgCount, whose sgData:
-// 1. Evenly divides the wgShape.
-// 2. Is a multiple of instData.
-// Example:
-// wgShape = [128, 64], instData = [8, 16], sgCount = 32
-// Returns layouts:
-// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
-// Definition (forward-declared above).
-static SmallVector<LayoutRepresentation>
-getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
- int64_t sgCount) {
- SmallVector<LayoutRepresentation> candidates;
- for (int sgLayout0 = 1; sgLayout0 <= sgCount; ++sgLayout0) {
- if (sgCount % sgLayout0)
- continue;
- int64_t sgLayout1 = sgCount / sgLayout0;
- int64_t sgData0 = wgShape[0] / sgLayout0;
- int64_t sgData1 = wgShape[1] / sgLayout1;
- if ((wgShape[0] % sgLayout0 || wgShape[1] % sgLayout1) ||
- (sgData0 % instData[0] || sgData1 % instData[1]))
- continue;
- candidates.emplace_back(sgLayout0, sgLayout1);
- }
- // Sort primarily by how balanced they are
- // (i.e., minimize the absolute difference between the two dimensions), and
- // secondarily by the first dimension in ascending order.
- llvm::sort(candidates, [](const LayoutRepresentation &lhs,
- const LayoutRepresentation &rhs) {
- int diffLhs = std::abs(lhs.first - lhs.second);
- int diffRhs = std::abs(rhs.first - rhs.second);
- if (diffLhs != diffRhs)
- return diffLhs < diffRhs;
- return lhs.first < rhs.first;
- });
- return candidates;
-}
-
/// Helper function to compute inst_data vectors for DPAS operands A, B, and
/// C/D.
static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
@@ -2012,28 +2012,26 @@ getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
if (consumerLayout && consumerLayout.isForWorkgroup()) {
- SmallVector<int64_t> sgLayoutD = consumerLayout.getEffectiveSgLayoutAsInt();
- consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
+ consumerSgLayout = consumerLayout.getEffectiveSgLayoutAsInt();
}
// Get all valid layouts for A, B and C/D operands
- auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
- auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
- auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
+ auto layoutsA = computeCandidateSgLayouts(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = computeCandidateSgLayouts(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD =
+ computeCandidateSgLayouts(cdTy.getShape(), instDataCD, numSg);
if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
return std::nullopt;
// Pick the best subgroup layout
- llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
- llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
- layoutsCD.end());
std::optional<LayoutRepresentation> bestPick;
- auto checkAlignedSgDataAB = [&](LayoutRepresentation sgLayout) {
- return aTy.getShape().back() / sgLayout.second ==
- bTy.getShape().front() / sgLayout.first;
+ auto checkAlignedSgDataAB = [&](const LayoutRepresentation &sgLayout) {
+ return aTy.getShape().back() / sgLayout[1] ==
+ bTy.getShape().front() / sgLayout[0];
};
for (auto &sgLayout : layoutsB) {
- if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
+ if (llvm::is_contained(layoutsA, sgLayout) &&
+ llvm::is_contained(layoutsCD, sgLayout)) {
if (!checkAlignedSgDataAB(sgLayout))
continue;
// Is in (A and B and CD) and matches consumer -> best pick
@@ -2052,30 +2050,13 @@ getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
if (!bestPick)
return std::nullopt;
- SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
- static_cast<int>(bestPick->second)};
- SmallVector<int> sgDataA = {static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(aTy.getShape()[1])};
- SmallVector<int> sgDataB = {
- static_cast<int>(bTy.getShape()[0]),
- static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
- SmallVector<int> sgDataCD = {
- static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
-
- auto dpasALayout =
- xegpu::LayoutAttr::get(context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataA), nullptr,
- nullptr, nullptr, nullptr);
- auto dpasBLayout =
- xegpu::LayoutAttr::get(context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataB), nullptr,
- nullptr, nullptr, nullptr);
- auto dpasCDLayout =
- xegpu::LayoutAttr::get(context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataCD), nullptr,
- nullptr, nullptr, nullptr);
+ const auto &picked = *bestPick;
+ auto dpasALayout = buildSgLayout(context, aTy.getShape(), picked,
+ /*dimK=*/aTy.getRank() - 1);
+ auto dpasBLayout = buildSgLayout(context, bTy.getShape(), picked,
+ /*dimK=*/bTy.getRank() - 2);
+ auto dpasCDLayout = buildSgLayout(context, cdTy.getShape(), picked);
return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
}
@@ -2093,12 +2074,23 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- auto [laneLayoutA, laneDataA] = getDefaultLaneLayoutAndData2DBlockIo(
- aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
- auto [laneLayoutB, laneDataB] = getDefaultLaneLayoutAndData2DBlockIo(
- bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
- auto [laneLayoutCD, laneDataCD] =
- getDefaultLaneLayoutAndData2DBlockIo(cdTy, uArch);
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
+
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
+
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
@@ -2243,6 +2235,25 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
xegpu::DistributeLayoutAttr consumerLayout, int numSg,
const xegpu::uArch::uArch *uArch) {
auto context = aTy.getContext();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
+
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
@@ -2270,15 +2281,6 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- auto [laneLayoutA, laneDataA] = getDefaultLaneLayoutAndData2DBlockIo(
- aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
- auto [laneLayoutB, laneDataB] = getDefaultLaneLayoutAndData2DBlockIo(
- bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
- auto [laneLayoutCD, laneDataCD] =
- getDefaultLaneLayoutAndData2DBlockIo(cdTy, uArch);
auto dpasALayout =
buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA);
auto dpasBLayout =
@@ -2286,7 +2288,6 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
auto dpasCDLayout = buildInstDataLayoutWithLane(context, instDataCD,
laneLayoutCD, laneDataCD);
- // Create scale layouts
auto aScaleLayout =
createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
auto bScaleLayout =
@@ -2295,22 +2296,16 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
bScaleLayout);
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- auto aLayout = getDefaultLaneLayout2DBlockIo(
- aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
- auto bLayout = getDefaultLaneLayout2DBlockIo(
- bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
- auto cdLayout = getDefaultLaneLayout2DBlockIo(cdTy, uArch);
+ auto dpasALayout = buildLaneLayout(context, laneLayoutA, laneDataA);
+ auto dpasBLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
+ auto dpasCDLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
- // Create scale layouts
auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, aLayout, false, uArch);
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, bLayout, true, uArch);
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
- return std::make_tuple(aLayout, bLayout, cdLayout, aScaleLayout,
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
bScaleLayout);
}
return std::nullopt;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 1d6297dcb8d6d..352f353dae0d2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -423,18 +423,18 @@ class LayoutInfoPropagation
visitOperation(Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) override;
- void visitBranchOperand(OpOperand &operand) override {};
+ void visitBranchOperand(OpOperand &operand) override{};
- void visitCallOperand(OpOperand &operand) override {};
+ void visitCallOperand(OpOperand &operand) override{};
void
visitNonControlFlowArguments(RegionSuccessor &successor,
- ArrayRef<BlockArgument> arguments) override {};
+ ArrayRef<BlockArgument> arguments) override{};
- void visitExternalCall(CallOpInterface call,
- ArrayRef<LayoutInfoLattice *> operands,
- ArrayRef<const LayoutInfoLattice *> results) override {
- };
+ void
+ visitExternalCall(CallOpInterface call,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) override{};
void setToExitState(LayoutInfoLattice *lattice) override {
(void)lattice->meet(LayoutInfo());
@@ -549,9 +549,9 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
// wgShape = [128, 64], instData = [8, 16], sgCount = 32
// Returns layouts:
// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
-SmallVector<std::pair<int, int>> getValidLayouts(ArrayRef<int64_t> wgShape,
- ArrayRef<int> instData,
- int64_t sgCount) {
+SmallVector<std::pair<int, int>>
+computeCandidateSgLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int> instData,
+ int64_t sgCount) {
SmallVector<std::pair<int, int>> candidates;
for (int sgLayout0 = 1; sgLayout0 <= sgCount; ++sgLayout0) {
if (sgCount % sgLayout0)
@@ -1211,12 +1211,12 @@ void LayoutInfoPropagation::visitLoadGatherOp(
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
- // The user-provided anchor may carry only inst_data; complete it with
- // a lane factorization derived from the load-side Lane setup so
- // downstream paths see a fully-populated layout.
- requiredAnchorLayoutAttr = xegpu::completeLoadGatherLayoutFromInstData(
- anchorLayoutAttr, resVecTy.getElementType(), uArch);
- load.setLayoutAttr(requiredAnchorLayoutAttr);
+ requiredAnchorLayoutAttr = anchorLayoutAttr;
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
+ anchorLayoutAttr, resVecTy.getElementType(), uArch);
+ load.setLayoutAttr(requiredAnchorLayoutAttr);
+ }
} else {
if (!resVecTy) {
load.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1256,12 +1256,12 @@ void LayoutInfoPropagation::visitStoreScatterOp(
int chunkSize = storeScatter.getChunkSize().value_or(1);
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
- // The user-provided anchor may carry only inst_data; complete it with
- // a lane factorization derived from the store-side Lane setup so
- // downstream paths see a fully-populated layout.
- requiredAnchorLayoutAttr = xegpu::completeStoreScatterLayoutFromInstData(
- anchorLayoutAttr, srcVecTy.getElementType(), uArch);
- storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
+ requiredAnchorLayoutAttr = anchorLayoutAttr;
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArch);
+ storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
+ }
} else {
if (!srcVecTy) {
storeScatter.emitWarning("Not propagating, non-vector payload supplied.");
@@ -1320,29 +1320,28 @@ void LayoutInfoPropagation::visitLoadMatrixOp(
void LayoutInfoPropagation::visitStoreMatrixOp(
xegpu::StoreMatrixOp storeMatrix, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
- xegpu::DistributeLayoutAttr anchorLayout = storeMatrix.getLayoutAttr();
+ xegpu::DistributeLayoutAttr requiredAnchorLayoutAttr;
+ xegpu::DistributeLayoutAttr anchorLayoutAttr = storeMatrix.getLayoutAttr();
LayoutInfo layout;
- VectorType srcVecTy =
- llvm::cast<VectorType>(storeMatrix.getData().getType());
+ VectorType srcVecTy = llvm::cast<VectorType>(storeMatrix.getData().getType());
const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
if (!uArch)
return;
- if (hasParamsOfLayoutKind(anchorLayout)) {
- // The user-provided anchor may carry only inst_data; complete it with
- // a lane factorization derived from the store-side Lane setup.
- auto completed = xegpu::completeStoreScatterLayoutFromInstData(
- anchorLayout, srcVecTy.getElementType(), uArch);
- storeMatrix.setLayoutAttr(completed);
- layout = LayoutInfo(completed);
+ if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
+ requiredAnchorLayoutAttr = anchorLayoutAttr;
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArch);
+ storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
+ }
} else {
int chunkSize =
1; // placeHolder for future use when StoreMatrix supports coalescing
- auto requiredAnchorLayoutAttr = xegpu::setupStoreMatrixAnchorLayout(
+ requiredAnchorLayoutAttr = xegpu::setupStoreMatrixAnchorLayout(
layoutKind, srcVecTy, chunkSize, uArch);
storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
- layout = LayoutInfo(requiredAnchorLayoutAttr);
}
-
+ layout = LayoutInfo(requiredAnchorLayoutAttr);
propagateIfChanged(operands[0], operands[0]->meet(layout));
}
>From 1bbcee367de5a8de8fd39a4c653c9982a23f3b0e Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 13 Jun 2026 06:45:24 +0000
Subject: [PATCH 23/42] polish setupLoadNdAnchorLayout
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 326 +++++++-----------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 4 +-
2 files changed, 119 insertions(+), 211 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 58a81640af721..5799bb1bad0c2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1512,7 +1512,6 @@ compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
int64_t subgroupSize, int64_t bitwidth,
int64_t packingSize, bool vnni = false) {
int64_t rank = instShape.size();
- assert(leadingDimsAreUnit(instShape, 2));
SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
int64_t packingDim = vnni ? rank - 2 : rank - 1;
laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
@@ -1579,8 +1578,8 @@ static SmallVector<LayoutRepresentation> enumerateFactorizations(int64_t total,
// wgShape = [128, 64], instData = [8, 16], sgCount = 32
// Returns: [[8,4], [16,2]], corresponding to sgData [16,16] and [8,32].
static SmallVector<LayoutRepresentation>
-computeCandidateSgLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
- int64_t sgCount) {
+getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
+ int64_t sgCount) {
int64_t rank = wgShape.size();
assert(rank > 0 && "wgShape must be non-empty");
assert(static_cast<int64_t>(instData.size()) == rank &&
@@ -1620,47 +1619,36 @@ computeCandidateSgLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
return candidates;
}
-/// Validates whether `instData` is a hardware-viable inst_data for an ND op
-/// with the given block params and lane factor. Specifically:
-/// - leading batch dims must be 1
-/// - innermost dim must be a divisor of `dataShape.back()`, a multiple of
-/// `bWidths.front()` (smallest supported block width), and ≤ a small
-/// multiple of the largest supported block width
-/// - second-to-innermost dim must be in the supported heights for rank >= 2
-/// - each dim must be a multiple of `lane_layout[dim] * lane_data[dim]`
-static bool isValidNdInstData(ArrayRef<int64_t> instData,
- ArrayRef<int64_t> dataShape,
- ArrayRef<int> bWidths, ArrayRef<int> bHeights,
- ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData) {
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D.
+static std::optional<SmallVector<int64_t>>
+get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
+ ArrayRef<int> bHeights, ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
int rank = dataShape.size();
- if (static_cast<int>(instData.size()) != rank)
- return false;
-
- for (int dim = 0; dim < rank - 2; ++dim)
- if (instData[dim] != 1)
- return false;
-
- int64_t inner = instData.back();
- if (inner <= 0 || dataShape.back() % inner != 0)
- return false;
- int minWidth = *llvm::min_element(bWidths);
- int maxWidth = *llvm::max_element(bWidths);
- if (inner % minWidth != 0 || inner > maxWidth * /*maxBlockCount*/ 4)
- return false;
-
- if (rank >= 2) {
- int64_t height = instData[rank - 2];
- if (!llvm::is_contained(bHeights, static_cast<int>(height)))
- return false;
- }
+ // Compute inst_data from hardware block params. For Nd ops, the lane
+ // factorization above (laneLayout / laneData) is rigid; inst_data must be
+ // a multiple of lane_layout * lane_data on each dim (Category A
+ // invariant). If block params are unavailable for this element type
+ // (e.g. sub-byte floats with no uArch entry), fall back to
+ // lane_layout * lane_data (k = 1).
+ SmallVector<int64_t> instData(rank, 1);
+ assert(rank >= 2 && "dataShape must be at least 2D for 2D-block IO");
+ int instWidth =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
+ int instHeight =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape[rank - 2]), bHeights);
+ instData.back() = instWidth;
+ instData[rank - 2] = instHeight;
- for (int dim = 0; dim < rank; ++dim) {
- int64_t laneProduct = laneLayout[dim] * laneData[dim];
- if (laneProduct == 0 || instData[dim] % laneProduct != 0)
- return false;
+ if (instWidth == -1 || instHeight == -1) {
+ instData.back() = laneLayout.back() * laneData.back();
+ instData[rank - 2] = laneLayout[rank - 2] * laneData[rank - 2];
}
- return true;
+ for (int dim = 0; dim < rank; ++dim)
+ assert(instData[dim] % (laneLayout[dim] * laneData[dim]) == 0 &&
+ "inst_data must be a multiple of lane_layout * lane_data for ND op");
+ return instData;
}
static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
@@ -1700,127 +1688,38 @@ static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
/// inst_data when it is uArch-valid.
/// For Subgroup kind: if the consumer specifies a workgroup-level layout,
/// reuses it directly; otherwise picks the most balanced sg_layout via
-/// `computeCandidateSgLayouts` (requires `numSg`).
+/// `getSgLayoutCandidates` (requires `numSg`).
static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
ArrayRef<int64_t> dataShape, Type elemTy, ArrayRef<int> bWidths,
- ArrayRef<int> bHeights, unsigned packingSize,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ ArrayRef<int> bHeights, unsigned packingSize, int numSg,
const xegpu::uArch::uArch *uArch) {
int rank = dataShape.size();
assert(rank >= 1 && "Expected at least 1D shape for ND op");
- // Compute the default 2D block IO lane layout / lane data. Cap each by
- // the innermost shape so the product `lane_layout * lane_data` doesn't
- // exceed it (e.g. ui8 with shape innermost=16 cannot use the full
- // packing factor of 2 because subgroupSize * 2 = 32 > 16).
+ // Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
- int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
- SmallVector<int64_t> laneLayout(rank, 1);
- SmallVector<int64_t> laneData(rank, 1);
- int64_t innermostShape = dataShape.back();
- int64_t lanesOnInnermost =
- std::min<int64_t>(uArch->getSubgroupSize(), innermostShape);
- laneLayout.back() = lanesOnInnermost;
- if (lanesOnInnermost > 0)
- laneData.back() =
- std::min<int64_t>(packingFactor, innermostShape / lanesOnInnermost);
- else
- laneData.back() = 1;
-
- // Honor the consumer's lane info when valid. The consumer (e.g. dpas) may
- // require VNNI-style packing on a non-innermost dim that the load's own
- // packing default doesn't capture.
- auto honorConsumerLane = [&]() {
- if (!consumerLayout)
- return;
- SmallVector<int64_t> consumerLaneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> consumerLaneData =
- consumerLayout.getEffectiveLaneDataAsInt();
- if (consumerLaneLayout.size() != static_cast<size_t>(rank) ||
- consumerLaneData.size() != static_cast<size_t>(rank))
- return;
- // Validate: lane_layout * lane_data must divide each dim of dataShape.
- for (int dim = 0; dim < rank; ++dim) {
- int64_t product = consumerLaneLayout[dim] * consumerLaneData[dim];
- if (product == 0 || dataShape[dim] % product != 0)
- return;
- }
- laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
- laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
- };
- honorConsumerLane();
-
- // If the consumer carries an explicit `order`, propagate it through.
- DenseI32ArrayAttr orderAttr =
- consumerLayout ? consumerLayout.getOrder() : nullptr;
+ auto [laneLayout, laneData] =
+ compute2DBlockIOLaneLayoutAndData(dataShape, uArch->getSubgroupSize(),
+ bitwidth, packingSize, /*vnni=*/false);
if (layoutKind == xegpu::LayoutKind::Lane)
- return buildLaneLayout(context, laneLayout, laneData, orderAttr);
-
- // Subgroup-kind fast path: if the consumer already specifies a
- // workgroup-level layout, reuse it directly. Skip the inst_data
- // computation, which can fail for very small shapes (e.g. dpas_mx scale
- // operands like 128x16 where no supported block width divides 16).
- if (layoutKind == xegpu::LayoutKind::Subgroup && consumerLayout &&
- consumerLayout.isForWorkgroup())
- return consumerLayout;
-
- // Compute inst_data from hardware block params. For Nd ops, the lane
- // factorization above (laneLayout / laneData) is rigid; inst_data must be
- // a multiple of lane_layout * lane_data on each dim (Category A
- // invariant). If block params are unavailable for this element type
- // (e.g. sub-byte floats with no uArch entry), fall back to
- // lane_layout * lane_data (k = 1).
- SmallVector<int64_t> instData(rank, 1);
- int instWidth =
- xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
- if (instWidth == -1)
- instData.back() = laneLayout.back() * laneData.back();
- else
- instData.back() = instWidth;
- if (rank >= 2) {
- int instHeight = xegpu::getLargestDivisor(
- static_cast<int>(dataShape[rank - 2]), bHeights);
- if (instHeight == -1)
- instData[rank - 2] = laneLayout[rank - 2] * laneData[rank - 2];
- else
- instData[rank - 2] = instHeight;
- }
-
- // Honor the consumer's inst_data if it is uArch-valid.
- if (consumerLayout) {
- SmallVector<int64_t> consumerInstData =
- consumerLayout.getEffectiveInstDataAsInt();
- if (!consumerInstData.empty() &&
- isValidNdInstData(consumerInstData, dataShape, bWidths, bHeights,
- laneLayout, laneData))
- instData.assign(consumerInstData.begin(), consumerInstData.end());
- }
+ return buildLaneLayout(context, laneLayout, laneData);
- // Category A invariant: inst_data is a multiple of lane_layout * lane_data.
- for (int dim = 0; dim < rank; ++dim) {
- int64_t laneProduct = laneLayout[dim] * laneData[dim];
- assert(instData[dim] % laneProduct == 0 &&
- "inst_data must be a multiple of lane_layout * lane_data for ND op");
- (void)laneProduct;
- }
+ auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights,
+ laneLayout, laneData);
if (layoutKind == xegpu::LayoutKind::InstData)
- return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData,
- orderAttr);
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ laneData);
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- // The consumer-fast-path above already returned for the case where the
- // consumer carries a workgroup-level layout, so here numSg must be set.
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
- auto sgLayouts = computeCandidateSgLayouts(dataShape, instData, numSg);
+ auto sgLayouts = getSgLayoutCandidates(dataShape, *instData, numSg);
if (sgLayouts.empty())
return nullptr;
- return buildSgLayout(context, dataShape, sgLayouts.front(), /*dimK=*/-1,
- orderAttr);
+ return buildSgLayout(context, dataShape, sgLayouts.front(), /*dimK=*/-1);
}
return nullptr;
@@ -1850,7 +1749,7 @@ xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
return setupGenericNdAnchorLayout(layoutKind, context, srcVecTy.getShape(),
elemTy, bWidths, bHeights, packingSize,
- /*consumerLayout=*/nullptr, numSg, uArch);
+ numSg, uArch);
}
/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
@@ -1877,7 +1776,7 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
return setupGenericNdAnchorLayout(layoutKind, context, tdescTy.getShape(),
elemTy, bWidths, bHeights, packingSize,
- /*consumerLayout=*/nullptr, numSg, uArch);
+ numSg, uArch);
}
/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
@@ -1890,17 +1789,31 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
VectorType resVecTy,
xegpu::DistributeLayoutAttr consumerLayout,
int numSg, const xegpu::uArch::uArch *uArch) {
- auto context = resVecTy.getContext();
- Type elemTy = resVecTy.getElementType();
- // Subgroup-kind fast path: if the consumer already specifies a complete
- // workgroup-level layout, reuse it directly. We don't need the uArch block
- // params at all (which may be unavailable for unusual element types like
- // sub-byte floats used in dpas_mx scale operands).
- if (layoutKind == xegpu::LayoutKind::Subgroup && consumerLayout &&
- consumerLayout.isForWorkgroup())
+ assert(consumerLayout && "Expected a valid consumer layout");
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(consumerLayout.isForWorkgroup() &&
+ "Expected consumer layout to be a complete workgroup-level layout");
return consumerLayout;
+ }
+ int rank = resVecTy.getRank();
+ SmallVector<int64_t> consumerInstData =
+ consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> consumerOrder = consumerLayout.getEffectiveOrderAsInt();
+ bool hasTransform = false;
+ bool hasTranspose = false;
+ if (!consumerLaneLayout.empty()) {
+ hasTransform = consumerLaneData[rank - 2] != 1;
+ hasTranspose = consumerOrder[0] != rank;
+ }
+ auto context = resVecTy.getContext();
+ Type elemTy = resVecTy.getElementType();
+ auto dataShape = resVecTy.getShape();
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
uArch->getInstruction(
@@ -1908,50 +1821,48 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
if (!uArchInstruction)
return nullptr;
unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
-
- // Lane kind only needs subgroupSize + packingSize. Skip the block-WHC
- // lookup, which can fail for element types without a uArch entry (e.g.
- // sub-byte floats like f4E2M1FN), and let the generic helper produce a
- // default lane layout from packingSize alone.
- if (layoutKind == xegpu::LayoutKind::Lane)
- return setupGenericNdAnchorLayout(
- layoutKind, context, resVecTy.getShape(), elemTy,
- /*bWidths=*/{}, /*bHeights=*/{}, packingSize, consumerLayout, numSg,
- uArch);
-
- // InstData / Subgroup kinds need block params. Transform / transpose /
- // upConv are lane-level concerns; treat them as no-op at the propagation
- // stage (consistent with visitLoadNdOp, which warns on transpose).
auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
- elemTy, /*hasTransform=*/false, /*hasTranspose=*/false,
+ elemTy, hasTransform, hasTranspose,
/*upConv=*/false);
if (!blockWHC)
return nullptr;
auto [bWidths, bHeights, bCounts] = blockWHC.value();
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ int64_t height = consumerInstData[rank - 2];
+ int64_t width = consumerInstData[rank - 1];
+ auto maxBlockCount = *llvm::max_element(bCounts);
+ auto maxWidth = *llvm::max_element(bWidths);
+ if (llvm::is_contained(bWidths, static_cast<int>(width)) ||
+ (width % maxWidth == 0 && width / maxWidth < maxBlockCount)) {
+ if (!llvm::is_contained(bHeights, static_cast<int>(height)))
+ return consumerLayout;
+ }
+ }
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ bool validLaneLayout = true;
+ for (int dim = 0; dim < rank; ++dim) {
+ int64_t laneProduct = consumerLaneLayout[dim] * consumerLaneData[dim];
+ if (dataShape[dim] % laneProduct != 0)
+ validLaneLayout = false;
+ }
+ if (validLaneLayout)
+ return consumerLayout;
+ }
+
return setupGenericNdAnchorLayout(layoutKind, context, resVecTy.getShape(),
elemTy, bWidths, bHeights, packingSize,
- consumerLayout, numSg, uArch);
+ numSg, uArch);
}
/// Helper function to compute inst_data vectors for DPAS operands A, B, and
/// C/D.
static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
SmallVector<int64_t>>>
-getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
- const xegpu::uArch::uArch *uArch,
- bool isDpasMx = false) {
- const int subgroupSize = uArch->getSubgroupSize();
-
- const xegpu::uArch::MMAInstructionInterface *uArchInstruction;
- if (isDpasMx)
- uArchInstruction = dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
- else
- uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+getDpasInstDataLayouts(
+ VectorType aTy, VectorType bTy, VectorType cdTy,
+ const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
+ const int subgroupSize, bool isDpasMx = false) {
// M dimension is the second-to-last dim of A (handles batch dims).
const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
@@ -1998,17 +1909,12 @@ getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
static std::optional<
std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
xegpu::DistributeLayoutAttr>>
-getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
- VectorType bTy, VectorType cdTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
- if (!instDataVecs)
- return std::nullopt;
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- assert(instDataA.size() == 2 && instDataB.size() == 2 &&
- instDataCD.size() == 2 &&
- "Sg layout creation expects valid 2D inst data");
+getDpasSubgroupLayouts(
+ mlir::MLIRContext *context, VectorType aTy, VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, SmallVector<int64_t>>
+ instDataVecs) {
+ auto [instDataA, instDataB, instDataCD] = instDataVecs;
std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
if (consumerLayout && consumerLayout.isForWorkgroup()) {
@@ -2016,10 +1922,9 @@ getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
}
// Get all valid layouts for A, B and C/D operands
- auto layoutsA = computeCandidateSgLayouts(aTy.getShape(), instDataA, numSg);
- auto layoutsB = computeCandidateSgLayouts(bTy.getShape(), instDataB, numSg);
- auto layoutsCD =
- computeCandidateSgLayouts(cdTy.getShape(), instDataCD, numSg);
+ auto layoutsA = getSgLayoutCandidates(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = getSgLayoutCandidates(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD = getSgLayoutCandidates(cdTy.getShape(), instDataCD, numSg);
if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
return std::nullopt;
@@ -2091,15 +1996,17 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
cdTy.getElementType().getIntOrFloatBitWidth(),
cdTy.getElementType().getIntOrFloatBitWidth());
+ auto instDataVecs =
+ getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction, subgroupSize);
+ if (!instDataVecs)
+ return std::nullopt;
+
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
- return getupDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
- numSg, uArch);
+ return getDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
+ numSg, *instDataVecs);
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
- if (!instDataVecs)
- return std::nullopt;
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
return std::make_tuple(
buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
@@ -2254,12 +2161,16 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
cdTy.getShape(), subgroupSize,
cdTy.getElementType().getIntOrFloatBitWidth(),
cdTy.getElementType().getIntOrFloatBitWidth());
+ auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction,
+ subgroupSize, /*isDpasMx=*/true);
+ if (!instDataVecs)
+ return std::nullopt;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
- auto dpasLayouts = getupDpasSubgroupLayouts(context, aTy, bTy, cdTy,
- consumerLayout, numSg, uArch);
+ auto dpasLayouts = getDpasSubgroupLayouts(
+ context, aTy, bTy, cdTy, consumerLayout, numSg, *instDataVecs);
if (!dpasLayouts)
return std::nullopt;
@@ -2275,10 +2186,7 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
bScaleLayout);
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto instDataVecs =
- getDpasInstDataVectors(aTy, bTy, cdTy, uArch, /*isDpasMx=*/true);
- if (!instDataVecs)
- return std::nullopt;
+
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
auto dpasALayout =
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 352f353dae0d2..fdb767f66ca18 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -550,8 +550,8 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
// Returns layouts:
// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
SmallVector<std::pair<int, int>>
-computeCandidateSgLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int> instData,
- int64_t sgCount) {
+getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int> instData,
+ int64_t sgCount) {
SmallVector<std::pair<int, int>> candidates;
for (int sgLayout0 = 1; sgLayout0 <= sgCount; ++sgLayout0) {
if (sgCount % sgLayout0)
>From 5d1bfb97d39cb994ab540f3721bf52ef47b078ed Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 13 Jun 2026 18:35:52 +0000
Subject: [PATCH 24/42] fix dpas_mx_f8e5m2 inst_data
---
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 8 +-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 151 +++++++++++++++---
.../XeGPU/propagate-layout-inst-data.mlir | 114 +++++--------
3 files changed, 167 insertions(+), 106 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index eeb1100cc8eab..60d404c0385be 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -97,6 +97,7 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
static const int kWidth32[] = {32};
static const int kWidth16[] = {16};
+ static const int kWidthAtLeast16[] = {16, 32};
static const int kWidth8[] = {8};
static const int32_t kCount1[] = {1};
@@ -109,7 +110,7 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
using Value = std::tuple<llvm::ArrayRef<int32_t>, llvm::ArrayRef<int32_t>,
llvm::ArrayRef<int32_t>>;
static const llvm::DenseMap<Key, Value> kMap = {
- {{1, false, false, false}, {kWidth32, kHeightAtLeast1, kCount2}},
+ {{1, false, false, false}, {kWidthAtLeast16, kHeightAtLeast1, kCount2}},
{{1, false, false, true}, {kWidth16, kHeightAtLeast8, kCount4Only}},
{{2, false, false, false}, {kWidth16, kHeightAtLeast1, kCount2}},
{{4, false, false, false}, {kWidth16, kHeightAtLeast1, kCount1}},
@@ -117,8 +118,9 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
{{1, true, false, false}, {kWidth16, kHeightAtLeast32, kCount4}},
{{2, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
// Block Loads with Transpose:
- {{4, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}},
- };
+ {{1, false, true, false}, {kWidth32, kHeightAtLeast16, kCount1}},
+ {{2, false, true, false}, {kWidth16, kHeightAtLeast16, kCount1}},
+ {{4, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}}};
const int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
auto it = kMap.find({elemByteSize, hasTransform, hasTranspose, upConv});
if (it != kMap.end())
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 5799bb1bad0c2..fac2621fcbe23 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1510,15 +1510,50 @@ xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
int64_t subgroupSize, int64_t bitwidth,
- int64_t packingSize, bool vnni = false) {
+ int64_t packingSize, bool vnni = false,
+ bool transpose = false) {
int64_t rank = instShape.size();
SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
int64_t packingDim = vnni ? rank - 2 : rank - 1;
laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
- laneLayout.back() = subgroupSize;
+ assert(
+ !(vnni && transpose) &&
+ "transpose and VNNI cannot be enabled at the same time for 2D block IO");
+ if (transpose)
+ laneLayout[rank - 2] = subgroupSize;
+ else
+ laneLayout.back() = subgroupSize;
+ llvm::dbgs() << "[DEBUG compute2DBlockIOLaneLayoutAndData] instShape = [";
+ for (int64_t i = 0; i < rank; ++i) {
+ if (i > 0)
+ llvm::dbgs() << ", ";
+ llvm::dbgs() << instShape[i];
+ }
+ llvm::dbgs() << "], subgroupSize = " << subgroupSize
+ << ", bitwidth = " << bitwidth
+ << ", packingSize = " << packingSize << ", vnni = " << vnni
+ << "\n";
+ llvm::dbgs() << "[DEBUG compute2DBlockIOLaneLayoutAndData] laneLayout = [";
+ for (int64_t i = 0; i < rank; ++i) {
+ if (i > 0)
+ llvm::dbgs() << ", ";
+ llvm::dbgs() << laneLayout[i];
+ }
+ llvm::dbgs() << "], laneData = [";
+ for (int64_t i = 0; i < rank; ++i) {
+ if (i > 0)
+ llvm::dbgs() << ", ";
+ llvm::dbgs() << laneData[i];
+ }
+ llvm::dbgs() << "]\n";
// assert that the lane layout and data fit in the inst shape
for (int64_t i = 0; i < rank; ++i) {
int64_t laneProduct = laneLayout[i] * laneData[i];
+ llvm::dbgs() << "[DEBUG compute2DBlockIOLaneLayoutAndData] dim " << i
+ << ": instShape[i] = " << instShape[i]
+ << ", laneProduct = " << laneProduct
+ << ", divisible = " << (instShape[i] % laneProduct == 0)
+ << "\n";
assert(instShape[i] % laneProduct == 0 &&
"lane_layout * lane_data must evenly divide the inst shape");
(void)laneProduct;
@@ -1658,7 +1693,7 @@ static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
DenseI32ArrayAttr orderAttr = nullptr) {
SmallVector<int> sgData(sgLayout.size());
SmallVector<int> sgLayoutInt(sgLayout.begin(), sgLayout.end());
- for (int dim = 0; dim < sgLayout.size(); ++dim) {
+ for (int dim = 0; dim < (int)sgLayout.size(); ++dim) {
if (dim == dimK)
sgData[dim] = wgTileShape[dim];
else
@@ -1697,6 +1732,17 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
int rank = dataShape.size();
assert(rank >= 1 && "Expected at least 1D shape for ND op");
+ llvm::dbgs() << "[DEBUG setupGenericNdAnchorLayout] ENTRY, layoutKind = "
+ << static_cast<int>(layoutKind) << "\n";
+ llvm::dbgs() << "[DEBUG setupGenericNdAnchorLayout] dataShape = [";
+ for (int i = 0; i < rank; ++i) {
+ if (i > 0)
+ llvm::dbgs() << ", ";
+ llvm::dbgs() << dataShape[i];
+ }
+ llvm::dbgs() << "], elemTy = " << elemTy << ", packingSize = " << packingSize
+ << "\n";
+
// Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
auto [laneLayout, laneData] =
@@ -1791,11 +1837,29 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
int numSg, const xegpu::uArch::uArch *uArch) {
assert(consumerLayout && "Expected a valid consumer layout");
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] ENTRY, layoutKind = "
+ << static_cast<int>(layoutKind) << ", resVecTy = " << resVecTy
+ << ", consumerLayout = " << consumerLayout << "\n";
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(consumerLayout.isForWorkgroup() &&
"Expected consumer layout to be a complete workgroup-level layout");
return consumerLayout;
}
+
+ auto context = resVecTy.getContext();
+ Type elemTy = resVecTy.getElementType();
+ auto subgroupSize = uArch->getSubgroupSize();
+ auto dataShape = resVecTy.getShape();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
+ if (!uArchInstruction) {
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] uArchInstruction is null, "
+ "returning nullptr\n";
+ return nullptr;
+ }
+
int rank = resVecTy.getRank();
SmallVector<int64_t> consumerInstData =
consumerLayout.getEffectiveInstDataAsInt();
@@ -1805,27 +1869,29 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout.getEffectiveLaneDataAsInt();
SmallVector<int64_t> consumerOrder = consumerLayout.getEffectiveOrderAsInt();
- bool hasTransform = false;
- bool hasTranspose = false;
- if (!consumerLaneLayout.empty()) {
- hasTransform = consumerLaneData[rank - 2] != 1;
- hasTranspose = consumerOrder[0] != rank;
+ if (consumerLaneLayout.empty() || consumerLaneData.empty()) {
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(),
+ elemTy.getIntOrFloatBitWidth(), false, false);
}
- auto context = resVecTy.getContext();
- Type elemTy = resVecTy.getElementType();
- auto dataShape = resVecTy.getShape();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
- if (!uArchInstruction)
- return nullptr;
- unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
+ bool hasTransform = consumerLaneData[rank - 2] != 1;
+ bool hasTranspose = consumerLaneLayout[rank - 2] != 1;
+ unsigned packingFactor =
+ hasTransform ? consumerLaneData[rank - 2] : consumerLaneData[rank - 1];
+ unsigned packingSize = packingFactor * elemTy.getIntOrFloatBitWidth();
+
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] calling "
+ "getBlockWidthHeightCount with elemTy = "
+ << elemTy << ", hasTransform = " << hasTransform
+ << ", hasTranspose = " << hasTranspose << "\n";
auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
elemTy, hasTransform, hasTranspose,
/*upConv=*/false);
- if (!blockWHC)
+ if (!blockWHC) {
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] blockWHC is nullopt, "
+ "returning nullptr\n";
return nullptr;
+ }
auto [bWidths, bHeights, bCounts] = blockWHC.value();
if (layoutKind == xegpu::LayoutKind::InstData) {
@@ -1833,11 +1899,44 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
int64_t width = consumerInstData[rank - 1];
auto maxBlockCount = *llvm::max_element(bCounts);
auto maxWidth = *llvm::max_element(bWidths);
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] InstData check: height = "
+ << height << ", width = " << width
+ << ", maxBlockCount = " << maxBlockCount
+ << ", maxWidth = " << maxWidth << "\n";
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] bWidths = [";
+ for (size_t i = 0; i < bWidths.size(); ++i) {
+ if (i > 0)
+ llvm::dbgs() << ", ";
+ llvm::dbgs() << bWidths[i];
+ }
+ llvm::dbgs() << "], bHeights = [";
+ for (size_t i = 0; i < bHeights.size(); ++i) {
+ if (i > 0)
+ llvm::dbgs() << ", ";
+ llvm::dbgs() << bHeights[i];
+ }
+ llvm::dbgs() << "]\n";
if (llvm::is_contained(bWidths, static_cast<int>(width)) ||
(width % maxWidth == 0 && width / maxWidth < maxBlockCount)) {
- if (!llvm::is_contained(bHeights, static_cast<int>(height)))
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] width check PASSED\n";
+ if (llvm::is_contained(bHeights, static_cast<int>(height))) {
+ llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] height check PASSED, "
+ "honoring consumer layout\n";
return consumerLayout;
+ }
}
+ llvm::dbgs()
+ << "[DEBUG setupLoadNdAnchorLayout] height and width check FAILED, "
+ "falling through\n";
+
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
+ hasTransform, hasTranspose);
+ auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights,
+ laneLayout, laneData);
+
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
bool validLaneLayout = true;
@@ -1846,13 +1945,15 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
if (dataShape[dim] % laneProduct != 0)
validLaneLayout = false;
}
- if (validLaneLayout)
+ if (validLaneLayout) {
return consumerLayout;
+ } else {
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
+ hasTransform, hasTranspose);
+ return buildLaneLayout(context, laneLayout, laneData);
+ }
}
-
- return setupGenericNdAnchorLayout(layoutKind, context, resVecTy.getShape(),
- elemTy, bWidths, bHeights, packingSize,
- numSg, uArch);
}
/// Helper function to compute inst_data vectors for DPAS operands A, B, and
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index bb20c63b08bf1..4121c5af5c780 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -362,22 +362,22 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
// -----
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui8_to_f4(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x16xui8>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui8> -> !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x16xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x16xui8>
-// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32]>} : vector<256x16xui8> to vector<256x32xf4E2M1FN>
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x32xui8>) {
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<256x32xui8>
+// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32]>} : vector<256x32xui8> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
-// CHECK-SAME: : vector<256x32xf4E2M1FN>
-func.func @bitcast_ui8_to_f4(%arg0: memref<256x16xui8>) {
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256x16xui8> -> !xegpu.tensor_desc<256x16xui8>
- %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x16xui8> -> vector<256x16xui8>
- %2 = vector.bitcast %1 : vector<256x16xui8> to vector<256x32xf4E2M1FN>
+// CHECK-SAME: : vector<256x64xf4E2M1FN>
+func.func @bitcast_ui8_to_f4(%arg0: memref<256x32xui8>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xui8> -> vector<256x32xui8>
+ %2 = vector.bitcast %1 : vector<256x32xui8> to vector<256x64xf4E2M1FN>
%3 = xegpu.convert_layout %2
<{input_layout = #xegpu.layout<inst_data = [32, 32]>,
target_layout = #xegpu.layout<inst_data = [32, 32]>}>
- : vector<256x32xf4E2M1FN>
+ : vector<256x64xf4E2M1FN>
return
}
}
@@ -408,82 +408,40 @@ func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
// -----
// CHECK-LABEL: func.func @dpas_mx_f8e5m2
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x64xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<64x32xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
-// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x2xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<2x32xf8E8M0FNU>
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x1024xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<1024x32xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
+// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x32xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<32x32xf8E8M0FNU>
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x64xf8E5M2> -> !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x32xf8E5M2> -> !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [4, 1]>>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x1024xf8E5M2> -> !xegpu.tensor_desc<16x1024xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<1024x32xf8E5M2> -> !xegpu.tensor_desc<1024x32xf8E5M2, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>>
// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x64xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<16x64xf8E5M2>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [4, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<64x32xf8E5M2, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<64x32xf8E5M2>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x2xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x2xf8E8M0FNU, #xegpu.layout<inst_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<16x2xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 32], lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 32], lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 32], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x32xf8E8M0FNU>
+// CHECK-SAME: !xegpu.tensor_desc<16x1024xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<16x1024xf8E5M2>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<1024x32xf8E5M2, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<1024x32xf8E5M2>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
-// CHECK-SAME: (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
-// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-gpu.module @test {
-func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M2>, %arg2: memref<16x32xbf16>,
- %arg3: memref<16x2xf8E8M0FNU>, %arg4: memref<2x32xf8E8M0FNU>) {
- %c0 = arith.constant 0 : index
- %cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<16x64xf8E5M2> -> !xegpu.tensor_desc<16x64xf8E5M2>
- %1 = xegpu.create_nd_tdesc %arg1 : memref<64x32xf8E5M2> -> !xegpu.tensor_desc<64x32xf8E5M2>
- %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x64xf8E5M2> -> vector<16x64xf8E5M2>
- %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<64x32xf8E5M2> -> vector<64x32xf8E5M2>
- %4 = xegpu.create_nd_tdesc %arg3 : memref<16x2xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU>
- %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x2xf8E8M0FNU> -> vector<16x2xf8E8M0FNU>
- %6 = xegpu.create_nd_tdesc %arg4 : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU>
- %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<2x32xf8E8M0FNU> -> vector<2x32xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
- %9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
- xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
- return
-}
-}
-
-// -----
-// CHECK-LABEL: func.func @dpas_mx_f4e2m1
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x128xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<128x32xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
-// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x4xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<4x32xf8E8M0FNU>
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x128xf4E2M1FN> -> !xegpu.tensor_desc<16x128xf4E2M1FN>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<128x32xf4E2M1FN> -> !xegpu.tensor_desc<128x32xf4E2M1FN>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x128xf4E2M1FN> -> vector<16x128xf4E2M1FN>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<128x32xf4E2M1FN> -> vector<128x32xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x4xf8E8M0FNU> -> !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x4xf8E8M0FNU, #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2]>> -> vector<16x4xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [4, 32], lane_layout = [1, 16], lane_data = [2, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [4, 32], lane_layout = [1, 16], lane_data = [2, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [4, 32], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<4x32xf8E8M0FNU>
-// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
-// CHECK-SAME: (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
+// CHECK-SAME: (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
gpu.module @test {
-func.func @dpas_mx_f4e2m1(%arg0: memref<16x128xf4E2M1FN>, %arg1: memref<128x32xf4E2M1FN>, %arg2: memref<16x32xbf16>,
- %arg3: memref<16x4xf8E8M0FNU>, %arg4: memref<4x32xf8E8M0FNU>) {
+func.func @dpas_mx_f8e5m2(%arg0: memref<16x1024xf8E5M2>, %arg1: memref<1024x32xf8E5M2>, %arg2: memref<16x32xbf16>,
+ %arg3: memref<16x32xf8E8M0FNU>, %arg4: memref<32x32xf8E8M0FNU>) {
%c0 = arith.constant 0 : index
%cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xf4E2M1FN> -> !xegpu.tensor_desc<16x128xf4E2M1FN>
- %1 = xegpu.create_nd_tdesc %arg1 : memref<128x32xf4E2M1FN> -> !xegpu.tensor_desc<128x32xf4E2M1FN>
- %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xf4E2M1FN> -> vector<16x128xf4E2M1FN>
- %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<128x32xf4E2M1FN> -> vector<128x32xf4E2M1FN>
- %4 = xegpu.create_nd_tdesc %arg3 : memref<16x4xf8E8M0FNU> -> !xegpu.tensor_desc<16x4xf8E8M0FNU>
- %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x4xf8E8M0FNU> -> vector<16x4xf8E8M0FNU>
- %6 = xegpu.create_nd_tdesc %arg4 : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU>
- %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<4x32xf8E8M0FNU> -> vector<4x32xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x1024xf8E5M2> -> !xegpu.tensor_desc<16x1024xf8E5M2>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<1024x32xf8E5M2> -> !xegpu.tensor_desc<1024x32xf8E5M2>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x1024xf8E5M2> -> vector<16x1024xf8E5M2>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<1024x32xf8E5M2> -> vector<1024x32xf8E5M2>
+ %4 = xegpu.create_nd_tdesc %arg3 : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU>
+ %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x32xf8E8M0FNU> -> vector<16x32xf8E8M0FNU>
+ %6 = xegpu.create_nd_tdesc %arg4 : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU>
+ %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<32x32xf8E8M0FNU> -> vector<32x32xf8E8M0FNU>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
return
>From 8fed92485869cfe2eb415aa1b035aa5e649fce7f Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 14 Jun 2026 18:04:32 +0000
Subject: [PATCH 25/42] passing tests, adjust uArch for subByte
---
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 27 ++++-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 13 ++-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 102 ++++++++---------
.../XeGPU/propagate-layout-inst-data.mlir | 64 +++++++++--
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 108 ++++++++++--------
5 files changed, 192 insertions(+), 122 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index 60d404c0385be..c0cd8be341adf 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -121,10 +121,33 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
{{1, false, true, false}, {kWidth32, kHeightAtLeast16, kCount1}},
{{2, false, true, false}, {kWidth16, kHeightAtLeast16, kCount1}},
{{4, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}}};
- const int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
+ int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
+ // handle sub-byte elements by treating them as 1 byte elements
+ if (elemByteSize == 0)
+ elemByteSize = 1;
auto it = kMap.find({elemByteSize, hasTransform, hasTranspose, upConv});
- if (it != kMap.end())
+ if (it != kMap.end()) {
+ // for sub-byte elements, need to double width retrieved from map since
+ // the map is based on byte-sized elements
+ if (elemTy.getIntOrFloatBitWidth() < 8) {
+ int subByteElemCount = 8 / elemTy.getIntOrFloatBitWidth();
+ auto [widths, heights, counts] = it->second;
+ if (hasTransform) {
+ llvm::SmallVector<int, 8> newHeights;
+ for (int h : heights)
+ newHeights.push_back(h * subByteElemCount);
+ return std::make_tuple(widths, llvm::ArrayRef<int>(newHeights),
+ counts);
+ } else {
+ llvm::SmallVector<int, 8> newWidths;
+ for (int w : widths)
+ newWidths.push_back(w * subByteElemCount);
+ return std::make_tuple(llvm::ArrayRef<int>(newWidths), heights,
+ counts);
+ }
+ }
return it->second;
+ }
return std::nullopt;
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index fac2621fcbe23..eeff926e3d8a3 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1869,11 +1869,9 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout.getEffectiveLaneDataAsInt();
SmallVector<int64_t> consumerOrder = consumerLayout.getEffectiveOrderAsInt();
- if (consumerLaneLayout.empty() || consumerLaneData.empty()) {
- auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
- dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(),
- elemTy.getIntOrFloatBitWidth(), false, false);
- }
+ assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
+ "Expected consumer layout to have lane_layout and lane_data");
+
bool hasTransform = consumerLaneData[rank - 2] != 1;
bool hasTranspose = consumerLaneLayout[rank - 2] != 1;
unsigned packingFactor =
@@ -1922,7 +1920,9 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
if (llvm::is_contained(bHeights, static_cast<int>(height))) {
llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] height check PASSED, "
"honoring consumer layout\n";
- return consumerLayout;
+ return buildInstDataLayoutWithLane(context, consumerInstData,
+ consumerLaneLayout, consumerLaneData,
+ consumerLayout.getOrder());
}
}
llvm::dbgs()
@@ -1954,6 +1954,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
return buildLaneLayout(context, laneLayout, laneData);
}
}
+ return nullptr;
}
/// Helper function to compute inst_data vectors for DPAS operands A, B, and
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index fdb767f66ca18..923e6966234ec 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -263,52 +263,54 @@ struct LayoutInfoLattice : public Lattice<LayoutInfo> {
/// is assigned to a value when the layout is not fixed by some anchor operation
/// (like DPAS).
-/// Helper Function to get the default layout for uniform values like constants.
-/// For 1D vector, lane_layout is [subgroupSize] and lane_data is [1].
-/// For 2D vector, lane_layout is [1, subgroupSize] and lane_data is [1, 1].
-/// For ND vector (N>2), leading dims get unit lane_layout and lane_data.
-static LayoutInfo getDefaultSIMTLayoutInfo(mlir::MLIRContext *ctx,
- unsigned rank,
- const xegpu::uArch::uArch *uArch) {
- assert(rank >= 1 && "Expected at least 1D vector.");
- if (rank == 1) {
- return LayoutInfo(
- xegpu::LayoutAttr::get(ctx, {uArch->getSubgroupSize()}, {1}));
- }
- // For rank >= 2, lane_layout is [1, ..., 1, subgroupSize] and
- // lane_data is [1, ..., 1, 1].
- SmallVector<int32_t> laneLayout(rank, 1);
- SmallVector<int32_t> laneData(rank, 1);
- laneLayout[rank - 1] = uArch->getSubgroupSize();
- return LayoutInfo(xegpu::LayoutAttr::get(ctx, laneLayout, laneData));
-}
-
-/// Helper to get the default layout for 2D block operations.
-/// For ND (N>2) types, leading dimensions get unit layout/data values.
-template <typename Ty>
-static LayoutInfo getSIMTLayoutInfoBlockIO(Ty ty,
- const xegpu::uArch::uArch *uArch,
- unsigned packingSize) {
- // Expecting at least 1D.
- assert(ty.getRank() >= 1 && "Expected at least 1D vector.");
- // Expecting int or float element type.
- assert(ty.getElementType().isIntOrFloat() &&
- "Expected int or float element type.");
- // If the rank is 1, then return default layout for 1D vector.
- if (ty.getRank() == 1)
- return getDefaultSIMTLayoutInfo(ty.getContext(), 1, uArch);
- // Packing factor is determined by the element type bitwidth.
- unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
- int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
- // For rank >= 2, distribute along the last dimension with leading units.
- unsigned rank = ty.getRank();
- SmallVector<int32_t> laneLayout(rank, 1);
- SmallVector<int32_t> laneData(rank, 1);
- laneLayout[rank - 1] = uArch->getSubgroupSize();
- laneData[rank - 1] = packingFactor;
- return LayoutInfo(
- xegpu::LayoutAttr::get(ty.getContext(), laneLayout, laneData));
-}
+// /// Helper Function to get the default layout for uniform values like
+// constants.
+// /// For 1D vector, lane_layout is [subgroupSize] and lane_data is [1].
+// /// For 2D vector, lane_layout is [1, subgroupSize] and lane_data is [1, 1].
+// /// For ND vector (N>2), leading dims get unit lane_layout and lane_data.
+// static LayoutInfo getDefaultSIMTLayoutInfo(mlir::MLIRContext *ctx,
+// unsigned rank,
+// const xegpu::uArch::uArch *uArch)
+// {
+// assert(rank >= 1 && "Expected at least 1D vector.");
+// if (rank == 1) {
+// return LayoutInfo(
+// xegpu::LayoutAttr::get(ctx, {uArch->getSubgroupSize()}, {1}));
+// }
+// // For rank >= 2, lane_layout is [1, ..., 1, subgroupSize] and
+// // lane_data is [1, ..., 1, 1].
+// SmallVector<int32_t> laneLayout(rank, 1);
+// SmallVector<int32_t> laneData(rank, 1);
+// laneLayout[rank - 1] = uArch->getSubgroupSize();
+// return LayoutInfo(xegpu::LayoutAttr::get(ctx, laneLayout, laneData));
+// }
+
+// /// Helper to get the default layout for 2D block operations.
+// /// For ND (N>2) types, leading dimensions get unit layout/data values.
+// template <typename Ty>
+// static LayoutInfo getSIMTLayoutInfoBlockIO(Ty ty,
+// const xegpu::uArch::uArch *uArch,
+// unsigned packingSize) {
+// // Expecting at least 1D.
+// assert(ty.getRank() >= 1 && "Expected at least 1D vector.");
+// // Expecting int or float element type.
+// assert(ty.getElementType().isIntOrFloat() &&
+// "Expected int or float element type.");
+// // If the rank is 1, then return default layout for 1D vector.
+// if (ty.getRank() == 1)
+// return getDefaultSIMTLayoutInfo(ty.getContext(), 1, uArch);
+// // Packing factor is determined by the element type bitwidth.
+// unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
+// int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
+// // For rank >= 2, distribute along the last dimension with leading units.
+// unsigned rank = ty.getRank();
+// SmallVector<int32_t> laneLayout(rank, 1);
+// SmallVector<int32_t> laneData(rank, 1);
+// laneLayout[rank - 1] = uArch->getSubgroupSize();
+// laneData[rank - 1] = packingFactor;
+// return LayoutInfo(
+// xegpu::LayoutAttr::get(ty.getContext(), laneLayout, laneData));
+// }
//===----------------------------------------------------------------------===//
// LayoutInfoPropagation
@@ -1011,13 +1013,7 @@ void LayoutInfoPropagation::visitLoadNdOp(
return;
int numSg = 0;
- // numSg is only needed when the helper has to derive a fresh sg_layout
- // for Subgroup kind. If the consumer already provides a complete
- // sg_layout, the helper will reuse it without consulting numSg.
- bool consumerHasSgLayout =
- consumerLayoutAttr &&
- !consumerLayoutAttr.getEffectiveSgLayoutAsInt().empty();
- if (layoutKind == xegpu::LayoutKind::Subgroup && !consumerHasSgLayout) {
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
auto numSgOrErr = getNumSg(load, uArch->getSubgroupSize());
if (failed(numSgOrErr)) {
load.emitWarning(
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 4121c5af5c780..33d779cf1b975 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -363,20 +363,20 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui8_to_f4(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x32xui8>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<256x32xui8>
-// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32]>} : vector<256x32xui8> to vector<256x64xf4E2M1FN>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x32xui8>
+// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>} : vector<256x32xui8> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>, target_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}>
// CHECK-SAME: : vector<256x64xf4E2M1FN>
func.func @bitcast_ui8_to_f4(%arg0: memref<256x32xui8>) {
%0 = xegpu.create_nd_tdesc %arg0 : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xui8> -> vector<256x32xui8>
%2 = vector.bitcast %1 : vector<256x32xui8> to vector<256x64xf4E2M1FN>
%3 = xegpu.convert_layout %2
- <{input_layout = #xegpu.layout<inst_data = [32, 32]>,
- target_layout = #xegpu.layout<inst_data = [32, 32]>}>
+ <{input_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>,
+ target_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}>
: vector<256x64xf4E2M1FN>
return
}
@@ -389,17 +389,17 @@ gpu.module @test {
// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<256x16xui16, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x16xui16>
-// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 64]>} : vector<256x16xui16> to vector<256x64xf4E2M1FN>
+// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>} : vector<256x16xui16> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 32]>, target_layout = #xegpu.layout<inst_data = [32, 32]>}>
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>, target_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
// CHECK-SAME: : vector<256x64xf4E2M1FN>
func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
%0 = xegpu.create_nd_tdesc %arg0 : memref<256x16xui16> -> !xegpu.tensor_desc<256x16xui16>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x16xui16> -> vector<256x16xui16>
%2 = vector.bitcast %1 : vector<256x16xui16> to vector<256x64xf4E2M1FN>
%3 = xegpu.convert_layout %2
- <{input_layout = #xegpu.layout<inst_data = [32, 32]>,
- target_layout = #xegpu.layout<inst_data = [32, 32]>}>
+ <{input_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>,
+ target_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
: vector<256x64xf4E2M1FN>
return
}
@@ -448,6 +448,48 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x1024xf8E5M2>, %arg1: memref<1024x32xf
}
}
+// -----
+// CHECK-LABEL: func.func @dpas_mx_f4e2m1
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<16x1024xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<1024x32xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<16x32xbf16>
+// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<16x32xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<32x32xf8E8M0FNU>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<16x32xbf16>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<16x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x1024xf4E2M1FN, #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<1024x32xf4E2M1FN> -> !xegpu.tensor_desc<1024x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x1024xf4E2M1FN, #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<16x1024xf4E2M1FN>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<1024x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<1024x32xf4E2M1FN>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 2]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 2]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<32x32xf8E8M0FNU>
+// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
+// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+gpu.module @test {
+func.func @dpas_mx_f4e2m1(%arg0: memref<16x1024xf4E2M1FN>, %arg1: memref<1024x32xf4E2M1FN>, %arg2: memref<16x32xbf16>,
+ %arg3: memref<16x32xf8E8M0FNU>, %arg4: memref<32x32xf8E8M0FNU>) {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x1024xf4E2M1FN>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<1024x32xf4E2M1FN> -> !xegpu.tensor_desc<1024x32xf4E2M1FN>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x1024xf4E2M1FN> -> vector<16x1024xf4E2M1FN>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<1024x32xf4E2M1FN> -> vector<1024x32xf4E2M1FN>
+ %4 = xegpu.create_nd_tdesc %arg3 : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU>
+ %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x32xf8E8M0FNU> -> vector<16x32xf8E8M0FNU>
+ %6 = xegpu.create_nd_tdesc %arg4 : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU>
+ %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<32x32xf8E8M0FNU> -> vector<32x32xf8E8M0FNU>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
+ %9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
+ xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
+ return
+}
+}
+
// -----
// shape_cast that collapses all src dims into a single innermost dst dim.
// Consumer carries inst_data + lane_layout=[..,subgroupSize] + lane_data=[..,1]
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index e3aad3a3fc965..37ec5769450d6 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -455,37 +455,45 @@ 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]+]]: !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>) {
+// 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>
-func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: !xegpu.tensor_desc<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>
+ %offset = vector.step : vector<16xindex>
%0 = vector.multi_reduction <add>, %arg0, %cst [0] : vector<16x16xf32> to vector<16xf32>
- xegpu.store_nd %0, %arg1[0] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %0, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
return
}
}
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_inner_reduction(
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>) {
+// 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 = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
-func.func @vector_inner_reduction(%arg0: vector<16x16xf32>, %arg1: !xegpu.tensor_desc<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>
+ %offset = vector.step : vector<16xindex>
%0 = vector.multi_reduction <add>, %arg0, %cst [1] : vector<16x16xf32> to vector<16xf32>
- xegpu.store_nd %0, %arg1[0] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %0, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
return
}
}
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @store_nd_with_offset(
+// CHECK-LABEL: func.func @store_with_offset(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf32>) {
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256xf32> -> !xegpu.tensor_desc<16xf32, #xegpu.layout<lane_layout = [16], lane_data = [1]>>
-func.func @store_nd_with_offset(%arg0: memref<256xf32>){
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
+// CHECK: %[[OFFSET:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16xindex>
+// CHECK: %[[VAL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1.000000e+00> : vector<16xf32>
+// CHECK: xegpu.store %[[VAL]], %[[ARG0]][%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+func.func @store_with_offset(%arg0: memref<256xf32>){
%c32 = arith.constant 32 : index
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
%1 = arith.constant dense<1.000000e+00> : vector<16xf32>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256xf32> -> !xegpu.tensor_desc<16xf32>
- xegpu.store_nd %1, %0[%c32] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %1, %arg0[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
return
}
}
@@ -519,12 +527,10 @@ func.func @prefetch_2d(%arg0: memref<256x256xf16>){
gpu.module @test {
// CHECK-LABEL: func.func @prefetch_1d(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256xf16> -> !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>
-// CHECK-NEXT: xegpu.prefetch_nd %[[T0]][0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : !xegpu.tensor_desc<16xf16, #xegpu.layout<lane_layout = [16], lane_data = [1]>>
+// CHECK: xegpu.prefetch %[[ARG0]][%{{.*}}] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<16xindex>
func.func @prefetch_1d(%arg0: memref<256xf16>){
- %c0 = arith.constant 0 : index
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256xf16> -> !xegpu.tensor_desc<16xf16>
- xegpu.prefetch_nd %0[0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>: !xegpu.tensor_desc<16xf16>
+ %offset = vector.step : vector<16xindex>
+ xegpu.prefetch %arg0[%offset] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<16xindex>
return
}
}
@@ -544,9 +550,9 @@ func.func @scf_while_and_condition(%arg0: memref<256xf32>, %arg1: memref<256xf32
%c16 = arith.constant 16 : i32
%c16_idx = arith.constant 16 : index
%c256 = arith.constant 256 : i32
- %0 = xegpu.create_nd_tdesc %arg0 : memref<256xf32> -> !xegpu.tensor_desc<16xf32>
- %1 = xegpu.load_nd %0[0] : !xegpu.tensor_desc<16xf32> -> vector<16xf32>
- %2 = xegpu.create_nd_tdesc %arg1 : memref<256xf32> -> !xegpu.tensor_desc<16xf32>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
+ %1 = xegpu.load %arg0[%offset], %mask : memref<256xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
%3:2 = scf.while (%arg2 = %1, %arg3 = %c0) : (vector<16xf32>, i32)
-> (vector<16xf32>, i32) {
@@ -554,12 +560,14 @@ func.func @scf_while_and_condition(%arg0: memref<256xf32>, %arg1: memref<256xf32
scf.condition(%4) %arg2, %arg3 : vector<16xf32>, i32
} do {
^bb0(%arg2: vector<16xf32>, %arg3: i32):
- xegpu.store_nd %arg2, %2[0] : vector<16xf32>, !xegpu.tensor_desc<16xf32>
+ xegpu.store %arg2, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
%4 = arith.addi %arg3, %c16 : i32
- %offset = arith.index_cast %4 : i32 to index
- %6 = xegpu.load_nd %0[%offset] : !xegpu.tensor_desc<16xf32> -> vector<16xf32>
+ %offset2 = arith.index_cast %4 : i32 to index
+ %offset2_v = vector.broadcast %offset2 : index to vector<16xindex>
+ %6 = xegpu.load %arg0[%offset], %mask : memref<256xf32>, vector<16xindex>, vector<16xi1> -> vector<16xf32>
scf.yield %6, %4 : vector<16xf32>, i32
}
+
return
}
}
@@ -965,40 +973,40 @@ gpu.module @test{
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @dpas_mx_f8e5m2
-// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x32xf8E5M2>, %[[ARG1:[0-9a-zA-Z]+]]: memref<32x16xf8E5M2>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xbf16>
-// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<8x1xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<1x16xf8E8M0FNU>
+// CHECK-LABEL: func.func @dpas_mx_fp4
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x64xf4E2M1FN>, %[[ARG1:[0-9a-zA-Z]+]]: memref<64x16xf4E2M1FN>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xbf16>
+// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: memref<8x2xf8E8M0FNU>, %[[ARG4:[0-9a-zA-Z]+]]: memref<2x16xf8E8M0FNU>
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.000000e+00> : vector<8x16xbf16>
-// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x32xf8E5M2> -> !xegpu.tensor_desc<8x32xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>>
-// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<32x16xf8E5M2> -> !xegpu.tensor_desc<32x16xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>>
-// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x32xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<8x32xf8E5M2>
-// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<32x16xf8E5M2, #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<32x16xf8E5M2>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x1xf8E8M0FNU> -> !xegpu.tensor_desc<8x1xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x1xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<8x1xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<1x16xf8E8M0FNU> -> !xegpu.tensor_desc<1x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<1x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<1x16xf8E8M0FNU>
+// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x64xf4E2M1FN> -> !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>>
+// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x16xf4E2M1FN> -> !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
+// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<8x64xf4E2M1FN>
+// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<64x16xf4E2M1FN>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>> -> vector<8x2xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
-// CHECK-SAME: (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
+// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
-func.func @dpas_mx_f8e5m2(%arg0: memref<8x32xf8E5M2>, %arg1: memref<32x16xf8E5M2>, %arg2: memref<8x16xbf16>,
- %arg3: memref<8x1xf8E8M0FNU>, %arg4: memref<1x16xf8E8M0FNU>) {
+func.func @dpas_mx_fp4(%arg0: memref<8x64xf4E2M1FN>, %arg1: memref<64x16xf4E2M1FN>, %arg2: memref<8x16xbf16>,
+ %arg3: memref<8x2xf8E8M0FNU>, %arg4: memref<2x16xf8E8M0FNU>) {
%c0 = arith.constant 0 : index
%cst = arith.constant dense<0.000000e+00> : vector<8x16xbf16>
- %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf8E5M2> -> !xegpu.tensor_desc<8x32xf8E5M2>
- %1 = xegpu.create_nd_tdesc %arg1 : memref<32x16xf8E5M2> -> !xegpu.tensor_desc<32x16xf8E5M2>
- %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x32xf8E5M2> -> vector<8x32xf8E5M2>
- %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<32x16xf8E5M2> -> vector<32x16xf8E5M2>
- %4 = xegpu.create_nd_tdesc %arg3 : memref<8x1xf8E8M0FNU> -> !xegpu.tensor_desc<8x1xf8E8M0FNU>
- %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<8x1xf8E8M0FNU> -> vector<8x1xf8E8M0FNU>
- %6 = xegpu.create_nd_tdesc %arg4 : memref<1x16xf8E8M0FNU> -> !xegpu.tensor_desc<1x16xf8E8M0FNU>
- %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<1x16xf8E8M0FNU> -> vector<1x16xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x64xf4E2M1FN> -> !xegpu.tensor_desc<8x64xf4E2M1FN>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<64x16xf4E2M1FN> -> !xegpu.tensor_desc<64x16xf4E2M1FN>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x64xf4E2M1FN> -> vector<8x64xf4E2M1FN>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<64x16xf4E2M1FN> -> vector<64x16xf4E2M1FN>
+ %4 = xegpu.create_nd_tdesc %arg3 : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU>
+ %5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<8x2xf8E8M0FNU> -> vector<8x2xf8E8M0FNU>
+ %6 = xegpu.create_nd_tdesc %arg4 : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU>
+ %7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<2x16xf8E8M0FNU> -> vector<2x16xf8E8M0FNU>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16>
return
>From 7d12d0a6dd2dd5e6cf23dcbcd06c9725636d635e Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 14 Jun 2026 18:35:18 +0000
Subject: [PATCH 26/42] passing all, refactor getNumSg
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 141 +++++++++---------
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 4 +
2 files changed, 74 insertions(+), 71 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 923e6966234ec..65cb48f19af68 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -580,7 +580,15 @@ getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int> instData,
return candidates;
}
-FailureOr<int64_t> getNumSg(Operation *op, const int sgSize) {
+FailureOr<int64_t>
+getNumSg(Operation *op, const int sgSize,
+ xegpu::DistributeLayoutAttr consumerLayout = nullptr) {
+ // first look for the number of subgroups required by the consumer layout
+ if (consumerLayout) {
+ auto sgLayout = consumerLayout.getEffectiveSgLayoutAsInt();
+ if (!sgLayout.empty())
+ return llvm::product_of(sgLayout);
+ }
// Oblivious to workitem layout, the total count matters.
auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>();
if (!gpuFunc)
@@ -606,19 +614,15 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
if (!uArch)
return;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- auto numSgOrErr = getNumSg(prefetch, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- prefetch.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ auto numSgOrErr = getNumSg(prefetch, uArch->getSubgroupSize());
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ prefetch.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
- auto layoutAttr =
- xegpu::setupPrefetchNdAnchorLayout(layoutKind, tdescTy, numSg, uArch);
+ auto layoutAttr = xegpu::setupPrefetchNdAnchorLayout(
+ layoutKind, tdescTy, numSgOrErr.value_or(0), uArch);
if (!layoutAttr) {
prefetch.emitWarning(
"Failed to determine required layout for prefetch_nd.");
@@ -653,11 +657,13 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
const uArch *uArch = getUArch(xegpu::getChipStr(reduction).value_or(""));
if (!uArch)
return;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- auto numSgOrErr = getNumSg(reduction, uArch->getSubgroupSize());
- if (succeeded(numSgOrErr))
- numSg = numSgOrErr.value();
+
+ auto numSgOrErr =
+ getNumSg(reduction, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ reduction.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
// The result layout represents the layout requirements of the operation.
@@ -666,7 +672,8 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
// propagated from consumer op, the conflict is resolved in later phase by
// converting the required result layout to the consumer layout
auto requiredResLayoutAttr = xegpu::setupMultiReductionResultLayout(
- layoutKind, sourceTy, consumerLayoutAttr, reductionDims, numSg, uArch);
+ layoutKind, sourceTy, consumerLayoutAttr, reductionDims,
+ numSgOrErr.value_or(0), uArch);
xegpu::setTemporaryLayout(reduction->getResult(0), requiredResLayoutAttr);
@@ -776,23 +783,23 @@ void LayoutInfoPropagation::visitDpasOp(
xegpu::DistributeLayoutAttr requiredCDLayoutAttr, requiredALayout,
requiredBLayout;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- LayoutInfo consumerLayout = results[0]->getValue();
- if (!consumerLayout.isAssigned())
- return;
- consumerLayoutAttr =
- dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
- auto numSgOrErr = getNumSg(dpas, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- dpas.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ LayoutInfo consumerLayout = results[0]->getValue();
+ if (!consumerLayout.isAssigned())
+ return;
+ consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
+
+ auto numSgOrErr =
+ getNumSg(dpas, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ dpas.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
- auto layouts = xegpu::setupDpasLayout(layoutKind, aTy, bTy, cdTy,
- consumerLayoutAttr, numSg, uArch);
+
+ auto layouts =
+ xegpu::setupDpasLayout(layoutKind, aTy, bTy, cdTy, consumerLayoutAttr,
+ numSgOrErr.value_or(0), uArch);
if (!layouts.has_value()) {
dpas.emitWarning(
"Failed to determine required layouts for DPAS operands.");
@@ -872,25 +879,23 @@ void LayoutInfoPropagation::visitDpasMxOp(
xegpu::DistributeLayoutAttr requiredCDLayoutAttr, requiredALayout,
requiredBLayout, requiredAScaleLayout, requiredBScaleLayout;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- LayoutInfo consumerLayout = results[0]->getValue();
- if (!consumerLayout.isAssigned())
- return;
- consumerLayoutAttr =
- dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
- auto numSgOrErr = getNumSg(dpasMx, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- dpasMx.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ LayoutInfo consumerLayout = results[0]->getValue();
+ if (!consumerLayout.isAssigned())
+ return;
+ consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(consumerLayout.get());
+
+ auto numSgOrErr =
+ getNumSg(dpasMx, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ dpasMx.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
- auto layouts =
- xegpu::setupDpasMxLayout(layoutKind, aTy, bTy, cdTy, aScaleTy, bScaleTy,
- consumerLayoutAttr, numSg, uArch);
+ auto layouts = xegpu::setupDpasMxLayout(
+ layoutKind, aTy, bTy, cdTy, aScaleTy, bScaleTy, consumerLayoutAttr,
+ numSgOrErr.value_or(0), uArch);
if (!layouts.has_value()) {
dpasMx.emitWarning(
"Failed to determine required layouts for DPAS_MX operands.");
@@ -955,19 +960,15 @@ void LayoutInfoPropagation::visitStoreNdOp(
if (!uArch)
return;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- auto numSgOrErr = getNumSg(store, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- store.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ auto numSgOrErr = getNumSg(store, uArch->getSubgroupSize());
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ store.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
auto layoutAttr = xegpu::setupStoreNdAnchorLayout(
- layoutKind, store.getValueType(), numSg, uArch);
+ layoutKind, store.getValueType(), numSgOrErr.value_or(0), uArch);
if (!layoutAttr) {
store.emitWarning("Failed to determine required layout for store_nd.");
return;
@@ -1012,19 +1013,17 @@ void LayoutInfoPropagation::visitLoadNdOp(
if (!uArch)
return;
- int numSg = 0;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- auto numSgOrErr = getNumSg(load, uArch->getSubgroupSize());
- if (failed(numSgOrErr)) {
- load.emitWarning(
- "Unable to determine the number of subgroups for the operation.");
- return;
- }
- numSg = numSgOrErr.value();
+ auto numSgOrErr =
+ getNumSg(load, uArch->getSubgroupSize(), consumerLayoutAttr);
+ if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
+ load.emitWarning(
+ "Unable to determine the number of subgroups for the operation.");
+ return;
}
auto layoutAttr = xegpu::setupLoadNdAnchorLayout(
- layoutKind, load.getType(), consumerLayoutAttr, numSg, uArch);
+ layoutKind, load.getType(), consumerLayoutAttr, numSgOrErr.value_or(0),
+ uArch);
if (!layoutAttr) {
load.emitWarning("Failed to determine required layout for load_nd.");
return;
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 37ec5769450d6..285f051094aad 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -127,6 +127,10 @@ func.func @extf_truncf(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor
%2 = arith.extf %1 : vector<16x16xf16> to vector<16x16xf32>
%3 = arith.truncf %2 : vector<16x16xf32> to vector<16x16xf16>
%4 = xegpu.dpas %0, %3 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
+ %5 = xegpu.convert_layout %4
+ <{input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
+ : vector<8x16xf32>
return %4 : vector<8x16xf32>
}
}
>From 4093160ac471c7d2ac8c0ad559b2e7b72bb8a5e4 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 15 Jun 2026 04:59:15 +0000
Subject: [PATCH 27/42] refactor computeReductionLaneLayoutAndData, passing all
tests
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 99 +++++++++----------
.../XeGPU/propagate-layout-inst-data.mlir | 14 +--
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 30 +++---
3 files changed, 68 insertions(+), 75 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index eeff926e3d8a3..6f2f11c497956 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -806,11 +806,9 @@ static xegpu::LayoutAttr buildInstDataLayoutWithLane(
return DenseI32ArrayAttr::get(context, v);
};
return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
- /*sg_data=*/nullptr,
- /*inst_data=*/toI32Attr(instData),
- /*lane_layout=*/toI32Attr(laneLayout),
- /*lane_data=*/toI32Attr(laneData),
- /*order=*/orderAttr);
+ /*sg_data=*/nullptr, toI32Attr(instData),
+ toI32Attr(laneLayout), toI32Attr(laneData),
+ orderAttr);
}
static xegpu::LayoutAttr
@@ -823,10 +821,25 @@ buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
};
return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
/*sg_data=*/nullptr,
- /*inst_data=*/nullptr,
- /*lane_layout=*/toI32Attr(laneLayout),
- /*lane_data=*/toI32Attr(laneData),
- /*order=*/orderAttr);
+ /*inst_data=*/nullptr, toI32Attr(laneLayout),
+ toI32Attr(laneData), orderAttr);
+}
+
+static xegpu::LayoutAttr
+buildLayout(mlir::MLIRContext *context, ArrayRef<int64_t> sgLayout,
+ ArrayRef<int64_t> sgData, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(
+ context, sgLayout.empty() ? nullptr : toI32Attr(sgLayout),
+ sgData.empty() ? nullptr : toI32Attr(sgData),
+ instData.empty() ? nullptr : toI32Attr(instData),
+ laneLayout.empty() ? nullptr : toI32Attr(laneLayout),
+ laneData.empty() ? nullptr : toI32Attr(laneData), orderAttr);
}
/// Computes the lane_layout and lane_data for a multi-reduction's source
@@ -852,6 +865,7 @@ buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
ArrayRef<int64_t> reductionDims,
+ ArrayRef<int64_t> consumerReductionDims,
int subgroupSize,
int64_t maxReduceVectorSize) {
int srcRank = srcShape.size();
@@ -863,6 +877,10 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
// `laneDim` carries the subgroupSize lanes; `vectorDim` packs the reduced
// elements into lane_data. Default: lanes on the innermost dim, reduced
// vector on the second-to-innermost dim.
+ if (secondInnermost >= 0 && llvm::is_contained(reductionDims, innermost) &&
+ reductionDims.size() == 1 && consumerReductionDims.empty()) {
+ std::swap(innermost, secondInnermost);
+ }
int laneDim = innermost;
int vectorDim = secondInnermost; // negative for rank 1
@@ -916,7 +934,7 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
/// * Consumer Layout:
/// #xegpu.layout<sgLayout=[32], sgData=[1]>
/// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[32,1], sgData=[1, 64]>, dims =
+/// #xegpu.slice<#xegpu.layout<sgLayout=[32,1], sgData=[1, 128]>, dims =
/// [1]>}
/// * Consumer Layout:
/// #xegpu.slice<#xegpu.layout<sgLayout=[8, 2, 4], sgData=[4, 64, 32]>,
@@ -1026,12 +1044,12 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
} else if (layoutKind == xegpu::LayoutKind::InstData) {
xegpu::SliceAttr consumerSliceLayout =
dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- auto reductionDimsOverrideConsumer =
+ auto consumerReductionDims =
consumerSliceLayout
? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
- : reductionDims;
+ : SmallVector<int64_t>({});
auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDimsOverrideConsumer, subgroupSize,
+ srcShape, reductionDims, consumerReductionDims, subgroupSize,
maxReduceVectorSize);
// inst_data is the per-instruction data, i.e. the element-wise product of
// lane_layout and lane_data.
@@ -1048,12 +1066,12 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
"dimensions are unit dimensions");
xegpu::SliceAttr consumerSliceLayout =
dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- auto reductionDimsOverrideConsumer =
+ auto consumerReductionDims =
consumerSliceLayout
? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
- : reductionDims;
+ : SmallVector<int64_t>({});
auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDimsOverrideConsumer, subgroupSize,
+ srcShape, reductionDims, consumerReductionDims, subgroupSize,
maxReduceVectorSize);
srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(laneLayout),
toInt32Attr(laneData));
@@ -2158,8 +2176,8 @@ createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
SmallVector<int64_t> laneData = matrixLayout.getEffectiveLaneDataAsInt();
auto order = matrixLayout.getOrder();
- SmallVector<int> scaleSgLayout;
- SmallVector<int> scaleSgData;
+ SmallVector<int64_t> scaleSgLayout;
+ SmallVector<int64_t> scaleSgData;
if (!sgLayout.empty() && !sgData.empty()) {
scaleSgLayout.assign(sgLayout.begin(), sgLayout.end());
scaleSgData.assign(sgData.begin(), sgData.end());
@@ -2172,7 +2190,7 @@ createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
// For DPAS_MX scales: if matrix has inst_data, scale needs adjusted
// inst_data. Scale inst_data is derived from matrix inst_data divided by
// scale factor.
- SmallVector<int> scaleInstData;
+ SmallVector<int64_t> scaleInstData;
if (!instData.empty()) {
scaleInstData.assign(instData.begin(), instData.end());
if (isBScale)
@@ -2185,11 +2203,12 @@ createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
1);
}
- SmallVector<int> scaleLaneLayout;
- SmallVector<int> scaleLaneData;
+ SmallVector<int64_t> scaleLaneLayout;
+ SmallVector<int64_t> scaleLaneData;
if (!laneLayout.empty() && !laneData.empty()) {
scaleLaneLayout.assign(laneLayout.begin(), laneLayout.end());
- scaleLaneData.assign(laneData.begin(), laneData.end());
+ scaleLaneData.assign(laneData.size(), 1);
+
bool isRowMajor = uArchInstruction->isLaneLayoutRowMajorOrder();
if (isBScale ^ isRowMajor)
std::swap(scaleLaneLayout[rank - 2], scaleLaneLayout[rank - 1]);
@@ -2197,38 +2216,12 @@ createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
// Then derive lane_data = inst_data / lane_layout so the Category A
// invariant inst_data = lane_layout * lane_data * k (with k = 1) holds
// for the scale operand's load_nd consumer.
- if (!scaleInstData.empty()) {
- for (int64_t d = rank - 2; d < rank; ++d) {
- scaleLaneLayout[d] =
- std::min<int64_t>(scaleInstData[d], scaleLaneLayout[d]);
- scaleLaneData[d] = std::max<int64_t>(
- scaleInstData[d] / std::max<int64_t>(scaleLaneLayout[d], 1), 1);
- }
- } else {
- // No inst_data on the matrix layout; fall back to capping by scale
- // shape and deriving lane_data from it (legacy behavior).
- scaleLaneLayout[rank - 2] =
- std::min<int64_t>(scaleShape[rank - 2], scaleLaneLayout[rank - 2]);
- scaleLaneData[rank - 2] = std::max<int64_t>(
- scaleShape[rank - 2] / scaleLaneLayout[rank - 2], 1);
- scaleLaneData[rank - 1] = std::max<int64_t>(
- scaleShape[rank - 1] / scaleLaneLayout[rank - 1], 1);
- }
+ auto layoutCap = scaleInstData.empty() ? scaleShape : scaleInstData;
+ for (int64_t d = rank - 2; d < rank; ++d)
+ scaleLaneLayout[d] = std::min<int64_t>(layoutCap[d], scaleLaneLayout[d]);
}
- return xegpu::LayoutAttr::get(
- context,
- scaleSgLayout.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleSgLayout),
- scaleSgData.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleSgData),
- scaleInstData.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleInstData),
- scaleLaneLayout.empty()
- ? nullptr
- : DenseI32ArrayAttr::get(context, scaleLaneLayout),
- scaleLaneData.empty() ? nullptr
- : DenseI32ArrayAttr::get(context, scaleLaneData),
- order);
+ return buildLayout(context, scaleSgLayout, scaleSgData, scaleInstData,
+ scaleLaneLayout, scaleLaneData, order);
}
/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 33d779cf1b975..ba169a04bcb4e 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -459,14 +459,14 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x1024xf8E5M2>, %arg1: memref<1024x32xf
// CHECK-SAME: !xegpu.tensor_desc<16x1024xf4E2M1FN, #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<16x1024xf4E2M1FN>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<1024x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<1024x32xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 2]>> -> vector<16x32xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<32x32xf8E8M0FNU>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 285f051094aad..25d713ccc8a0f 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -474,7 +474,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 = [1, 16], 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, 1]>, 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>
@@ -987,14 +987,14 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<8x64xf4E2M1FN>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<64x16xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>> -> vector<8x2xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<2x16xf8E8M0FNU>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<8x2xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -1029,14 +1029,14 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<8x64xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<8x64xf4E2M1FN>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<64x16xf4E2M1FN, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<64x16xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>> -> vector<8x2xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<2x16xf8E8M0FNU>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<8x2xf8E8M0FNU> -> !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<8x2xf8E8M0FNU, #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>> -> vector<8x2xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
>From b07274f937804917a886922c8fd4a01b2849f939 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 15 Jun 2026 05:20:00 +0000
Subject: [PATCH 28/42] polish
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 129 ++++--------------
1 file changed, 25 insertions(+), 104 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 6f2f11c497956..726dc1dc2fd13 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -30,7 +30,6 @@
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
-#include <functional>
#include <numeric>
using namespace mlir;
@@ -842,26 +841,19 @@ buildLayout(mlir::MLIRContext *context, ArrayRef<int64_t> sgLayout,
laneData.empty() ? nullptr : toI32Attr(laneData), orderAttr);
}
-/// Computes the lane_layout and lane_data for a multi-reduction's source
-/// layout. Only the innermost two dimensions are distributed; all leading
-/// dimensions are assumed to be unit (the caller verifies this via
-/// `leadingDimsAreUnit`).
+/// 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.
///
-/// The layout is chosen to minimize cross-lane reduction: whenever possible a
-/// reduction dimension is reduced *within* a lane (lane_layout == 1, with up to
-/// `maxReduceVectorSize` elements packed into lane_data), and the subgroup's
-/// lanes are spread across a non-reduction dimension instead.
-///
-/// - Exactly one of the innermost two dims is a reduction dim: place
-/// `subgroupSize` lanes on the non-reduction dim and keep lane_layout == 1
-/// on the reduction dim, packing up to `maxReduceVectorSize` reduced
-/// elements into lane_data along that reduction dim.
-/// - Both innermost dims are reduction dims (or the source is rank 1): fall
-/// back to the default of `subgroupSize` lanes on the innermost dim,
-/// packing `maxReduceVectorSize` elements on the second-to-innermost dim.
-///
-/// Returns the (lane_layout, lane_data) pair. The corresponding inst_data is
-/// simply the element-wise product lane_layout * lane_data.
+/// 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,
@@ -947,15 +939,19 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
/// first and then distribute remaining subgroups on the reduction
/// dimension.
///
-/// 2. InstData layout - Column reduction:
+/// 3. Lane layout - Default (lanes on innermost dim):
/// srcShape=[32, 64], reductionDims=[0], subgroupSize=16
-/// Result: instData=[1, 16] (maxReduceVectorSize=1, subgroupSize on
-/// innermost)
+/// Result: laneLayout=[1, 16], laneData=[1, 1]. The innermost dim is not
+/// reduced, so lanes stay on it.
///
-/// 3. Lane layout - Multi-dimensional reduction:
-/// srcShape=[16, 32, 64], reductionDims=[1], subgroupSize=16
-/// Result: laneLayout=[1, 1, 16], laneData=[1, 1, 1]
-/// (subgroupSize on innermost dim, max vector size on reduction dim)
+/// 4. Lane layout - Switch (lanes moved off the reduction dim):
+/// srcShape=[32, 64], reductionDims=[1], subgroupSize=16
+/// Result: laneLayout=[16, 1], laneData=[1, 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.
xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
xegpu::LayoutKind layoutKind, VectorType srcVecTy,
@@ -1541,37 +1537,9 @@ compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
laneLayout[rank - 2] = subgroupSize;
else
laneLayout.back() = subgroupSize;
- llvm::dbgs() << "[DEBUG compute2DBlockIOLaneLayoutAndData] instShape = [";
- for (int64_t i = 0; i < rank; ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << instShape[i];
- }
- llvm::dbgs() << "], subgroupSize = " << subgroupSize
- << ", bitwidth = " << bitwidth
- << ", packingSize = " << packingSize << ", vnni = " << vnni
- << "\n";
- llvm::dbgs() << "[DEBUG compute2DBlockIOLaneLayoutAndData] laneLayout = [";
- for (int64_t i = 0; i < rank; ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << laneLayout[i];
- }
- llvm::dbgs() << "], laneData = [";
- for (int64_t i = 0; i < rank; ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << laneData[i];
- }
- llvm::dbgs() << "]\n";
// assert that the lane layout and data fit in the inst shape
for (int64_t i = 0; i < rank; ++i) {
int64_t laneProduct = laneLayout[i] * laneData[i];
- llvm::dbgs() << "[DEBUG compute2DBlockIOLaneLayoutAndData] dim " << i
- << ": instShape[i] = " << instShape[i]
- << ", laneProduct = " << laneProduct
- << ", divisible = " << (instShape[i] % laneProduct == 0)
- << "\n";
assert(instShape[i] % laneProduct == 0 &&
"lane_layout * lane_data must evenly divide the inst shape");
(void)laneProduct;
@@ -1750,17 +1718,6 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
int rank = dataShape.size();
assert(rank >= 1 && "Expected at least 1D shape for ND op");
- llvm::dbgs() << "[DEBUG setupGenericNdAnchorLayout] ENTRY, layoutKind = "
- << static_cast<int>(layoutKind) << "\n";
- llvm::dbgs() << "[DEBUG setupGenericNdAnchorLayout] dataShape = [";
- for (int i = 0; i < rank; ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << dataShape[i];
- }
- llvm::dbgs() << "], elemTy = " << elemTy << ", packingSize = " << packingSize
- << "\n";
-
// Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
auto [laneLayout, laneData] =
@@ -1855,9 +1812,6 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
int numSg, const xegpu::uArch::uArch *uArch) {
assert(consumerLayout && "Expected a valid consumer layout");
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] ENTRY, layoutKind = "
- << static_cast<int>(layoutKind) << ", resVecTy = " << resVecTy
- << ", consumerLayout = " << consumerLayout << "\n";
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(consumerLayout.isForWorkgroup() &&
"Expected consumer layout to be a complete workgroup-level layout");
@@ -1872,11 +1826,8 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
uArch->getInstruction(
xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
- if (!uArchInstruction) {
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] uArchInstruction is null, "
- "returning nullptr\n";
+ if (!uArchInstruction)
return nullptr;
- }
int rank = resVecTy.getRank();
SmallVector<int64_t> consumerInstData =
@@ -1896,18 +1847,11 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
hasTransform ? consumerLaneData[rank - 2] : consumerLaneData[rank - 1];
unsigned packingSize = packingFactor * elemTy.getIntOrFloatBitWidth();
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] calling "
- "getBlockWidthHeightCount with elemTy = "
- << elemTy << ", hasTransform = " << hasTransform
- << ", hasTranspose = " << hasTranspose << "\n";
auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
elemTy, hasTransform, hasTranspose,
/*upConv=*/false);
- if (!blockWHC) {
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] blockWHC is nullopt, "
- "returning nullptr\n";
+ if (!blockWHC)
return nullptr;
- }
auto [bWidths, bHeights, bCounts] = blockWHC.value();
if (layoutKind == xegpu::LayoutKind::InstData) {
@@ -1915,37 +1859,14 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
int64_t width = consumerInstData[rank - 1];
auto maxBlockCount = *llvm::max_element(bCounts);
auto maxWidth = *llvm::max_element(bWidths);
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] InstData check: height = "
- << height << ", width = " << width
- << ", maxBlockCount = " << maxBlockCount
- << ", maxWidth = " << maxWidth << "\n";
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] bWidths = [";
- for (size_t i = 0; i < bWidths.size(); ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << bWidths[i];
- }
- llvm::dbgs() << "], bHeights = [";
- for (size_t i = 0; i < bHeights.size(); ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << bHeights[i];
- }
- llvm::dbgs() << "]\n";
if (llvm::is_contained(bWidths, static_cast<int>(width)) ||
(width % maxWidth == 0 && width / maxWidth < maxBlockCount)) {
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] width check PASSED\n";
if (llvm::is_contained(bHeights, static_cast<int>(height))) {
- llvm::dbgs() << "[DEBUG setupLoadNdAnchorLayout] height check PASSED, "
- "honoring consumer layout\n";
return buildInstDataLayoutWithLane(context, consumerInstData,
consumerLaneLayout, consumerLaneData,
consumerLayout.getOrder());
}
}
- llvm::dbgs()
- << "[DEBUG setupLoadNdAnchorLayout] height and width check FAILED, "
- "falling through\n";
auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
>From ea2c3a079eb1dca7c96a76896baaf9161f8abca4 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 15 Jun 2026 20:51:56 +0000
Subject: [PATCH 29/42] fix multireduction 4d to 3d case
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 193 ++++++++----------
1 file changed, 81 insertions(+), 112 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 726dc1dc2fd13..7a89170dabd6c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -370,6 +370,61 @@ void xegpu::removeTemporaryLayoutAttrs(Operation *op) {
});
}
+/// Returns true if every dimension of `shape` except the innermost
+/// `numInnerDims` is a unit (size-1) dimension.
+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,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr, toI32Attr(instData),
+ toI32Attr(laneLayout), toI32Attr(laneData),
+ orderAttr);
+}
+
+static xegpu::LayoutAttr
+buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr,
+ /*inst_data=*/nullptr, toI32Attr(laneLayout),
+ toI32Attr(laneData), orderAttr);
+}
+
+static xegpu::LayoutAttr
+buildLayout(mlir::MLIRContext *context, ArrayRef<int64_t> sgLayout,
+ ArrayRef<int64_t> sgData, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ auto toI32Attr = [&](auto range) {
+ SmallVector<int32_t> v(range.begin(), range.end());
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ return xegpu::LayoutAttr::get(
+ context, sgLayout.empty() ? nullptr : toI32Attr(sgLayout),
+ sgData.empty() ? nullptr : toI32Attr(sgData),
+ instData.empty() ? nullptr : toI32Attr(instData),
+ laneLayout.empty() ? nullptr : toI32Attr(laneLayout),
+ laneData.empty() ? nullptr : toI32Attr(laneData), orderAttr);
+}
+
/// Infers the source layout attribute for a broadcast operation given the
/// result layout attribute, result shape, source shape.
xegpu::DistributeLayoutAttr
@@ -666,21 +721,13 @@ xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
order.push_back(dimDiff - 1 - i);
}
- DenseI32ArrayAttr orderAttr = resLayout ? resLayout.getOrder() : nullptr;
- auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
- if (v.empty())
- return DenseI32ArrayAttr();
- SmallVector<int32_t> v32(v.begin(), v.end());
- return DenseI32ArrayAttr::get(context, v32);
- };
- auto srcLayout = xegpu::LayoutAttr::get(
- context, sgLayout.empty() ? nullptr : toAttr(sgLayout),
- sgData.empty() ? nullptr : toAttr(sgData),
- instData.empty() ? nullptr : toAttr(instData),
- laneLayout.empty() ? nullptr : toAttr(laneLayout),
- laneData.empty() ? nullptr : toAttr(laneData),
- (!orderAttr || orderAttr.empty()) ? nullptr : toAttr(order));
- return srcLayout;
+ DenseI32ArrayAttr orderAttr = DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
+ if (!resLayout.getOrder())
+ orderAttr = nullptr;
+
+ return buildLayout(context, sgLayout, sgData, instData, laneLayout,
+ laneData, orderAttr);
}
return resLayout;
}
@@ -727,18 +774,6 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// Use case 3: General dim collapse, for cross-sg reduction to SLM and other
// shape casts where consecutive src dims fold into a single dst dim.
- //
- // Mirrors use case 2's elegant shape: walk the dst-side groups and call
- // a single layout-attribute primitive per group. Here the primitive is
- // `expandDim(dim, targetShape)`, the inverse of `collapseDims`. It applies
- // the per-field distribution policy required for a no-data-movement collapse
- // (sg_layout/lane_layout spread outer-to-inner; sg_data/lane_data/inst_data
- // fill innermost-first; inst_data is seeded from lane_layout * lane_data).
- // See LayoutAttr::expandDim for the full policy.
- //
- // Iteration goes innermost-first (reverse dst order) so that each
- // expandDim/dropDims call only mutates dst positions whose indices are
- // unaffected by earlier calls.
SmallVector<SmallVector<int64_t>> collapseDims;
if (xegpu::matchDimCollapse(srcShape, resShape, collapseDims)) {
auto srcLayout = resLayout;
@@ -746,12 +781,10 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
dstIdx >= 0; --dstIdx) {
ArrayRef<int64_t> srcDims = collapseDims[dstIdx];
if (srcDims.empty()) {
- // Unit dst dim with no backing src dim: drop it.
srcLayout = srcLayout.dropDims({dstIdx});
continue;
}
if (srcDims.size() == 1)
- // 1:1 mapping, nothing to do for this dim.
continue;
SmallVector<int64_t> targetShape;
targetShape.reserve(srcDims.size());
@@ -761,7 +794,6 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
}
return srcLayout;
}
- llvm_unreachable("running into unsupported shape cast scenarios");
return nullptr;
}
@@ -776,71 +808,6 @@ xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
return payloadLayout;
}
-/// Returns true if every dimension of `shape` except the innermost
-/// `numInnerDims` is a unit (size-1) dimension.
-///
-/// Several reduction layout-setup paths (InstData, Lane) only distribute the
-/// innermost one or two dimensions and rely on all the leading dimensions
-/// being degenerate. This helper makes that assumption explicit and checkable
-/// instead of silently leaving leading dimensions undistributed.
-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; });
-}
-
-/// Builds a LayoutAttr carrying inst_data, lane_layout, and lane_data (no
-/// sg_layout / sg_data / order). Used by InstData-kind setup paths so the
-/// result layout can later be distributed without re-deriving the lane
-/// layout. `instData`, `laneLayout`, and `laneData` may have different
-/// element types; they are normalized to int32 entries.
-static xegpu::LayoutAttr buildInstDataLayoutWithLane(
- mlir::MLIRContext *context, ArrayRef<int64_t> instData,
- ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
- DenseI32ArrayAttr orderAttr = nullptr) {
- auto toI32Attr = [&](auto range) {
- SmallVector<int32_t> v(range.begin(), range.end());
- return DenseI32ArrayAttr::get(context, v);
- };
- return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
- /*sg_data=*/nullptr, toI32Attr(instData),
- toI32Attr(laneLayout), toI32Attr(laneData),
- orderAttr);
-}
-
-static xegpu::LayoutAttr
-buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData,
- DenseI32ArrayAttr orderAttr = nullptr) {
- auto toI32Attr = [&](auto range) {
- SmallVector<int32_t> v(range.begin(), range.end());
- return DenseI32ArrayAttr::get(context, v);
- };
- return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
- /*sg_data=*/nullptr,
- /*inst_data=*/nullptr, toI32Attr(laneLayout),
- toI32Attr(laneData), orderAttr);
-}
-
-static xegpu::LayoutAttr
-buildLayout(mlir::MLIRContext *context, ArrayRef<int64_t> sgLayout,
- ArrayRef<int64_t> sgData, ArrayRef<int64_t> instData,
- ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
- DenseI32ArrayAttr orderAttr = nullptr) {
- auto toI32Attr = [&](auto range) {
- SmallVector<int32_t> v(range.begin(), range.end());
- return DenseI32ArrayAttr::get(context, v);
- };
- return xegpu::LayoutAttr::get(
- context, sgLayout.empty() ? nullptr : toI32Attr(sgLayout),
- sgData.empty() ? nullptr : toI32Attr(sgData),
- instData.empty() ? nullptr : toI32Attr(instData),
- laneLayout.empty() ? nullptr : toI32Attr(laneLayout),
- laneData.empty() ? nullptr : toI32Attr(laneData), orderAttr);
-}
-
/// 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
@@ -962,12 +929,6 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
int srcRank = srcShape.size();
auto context = srcVecTy.getContext();
- // Helper lambda to convert int64 vectors to int32 DenseArrayAttr
- auto toInt32Attr = [&](ArrayRef<int64_t> vec) {
- SmallVector<int32_t> vec32(vec.begin(), vec.end());
- return DenseI32ArrayAttr::get(context, vec32);
- };
-
const int subgroupSize = uArch->getSubgroupSize();
int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
xegpu::DistributeLayoutAttr srcLayout;
@@ -1029,13 +990,14 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
order[i] = remainOrder++;
}
}
-
+ DenseI32ArrayAttr resOrderAttr = DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
+ if (!orderAttr || orderAttr.empty())
+ resOrderAttr = nullptr;
assert(remainingSgCount == 1 && "not all subgroups distributed");
- srcLayout = xegpu::LayoutAttr::get(
- context, toInt32Attr(sgLayout), toInt32Attr(sgData),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/
- (!orderAttr || orderAttr.empty()) ? nullptr : toInt32Attr(order));
+ srcLayout = buildLayout(context, sgLayout, sgData,
+ /*instData=*/{}, /*laneLayout=*/{},
+ /*laneData=*/{}, resOrderAttr);
}
} else if (layoutKind == xegpu::LayoutKind::InstData) {
xegpu::SliceAttr consumerSliceLayout =
@@ -1066,11 +1028,18 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
consumerSliceLayout
? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
: SmallVector<int64_t>({});
- auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDims, consumerReductionDims, subgroupSize,
- maxReduceVectorSize);
- srcLayout = xegpu::LayoutAttr::get(context, toInt32Attr(laneLayout),
- toInt32Attr(laneData));
+ if (consumerSliceLayout &&
+ consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
+ // at the lane level, the consumerSliceLayout can be directly reused
+ // since the inst_data propagation already insert convert_layout if
+ // the layout is not consistent
+ srcLayout = consumerSliceLayout.getParent();
+ } else {
+ auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ srcShape, reductionDims, consumerReductionDims, subgroupSize,
+ maxReduceVectorSize);
+ srcLayout = buildLaneLayout(context, laneLayout, laneData);
+ }
}
return xegpu::SliceAttr::get(context, srcLayout,
>From 423a33fc6c0852887a21273e9bbd3330d7af74dd Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 15 Jun 2026 21:20:13 +0000
Subject: [PATCH 30/42] polish
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 107 +++++++++---------
1 file changed, 51 insertions(+), 56 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 7a89170dabd6c..5ccad0f8335b8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1031,7 +1031,7 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
if (consumerSliceLayout &&
consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
// at the lane level, the consumerSliceLayout can be directly reused
- // since the inst_data propagation already insert convert_layout if
+ // since the inst_data propagation already insert convert_layout if
// the layout is not consistent
srcLayout = consumerSliceLayout.getParent();
} else {
@@ -1065,12 +1065,10 @@ xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
assert(true && "instData layout assignment not supported for reduction (op "
"is not expected at this level).");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- SmallVector<int32_t> laneLayout(1), laneData(1);
- laneLayout[0] = std::min(subgroupSize, static_cast<int32_t>(srcShape[0]));
+ SmallVector<int64_t> laneLayout(1), laneData(1);
+ laneLayout[0] = std::min(static_cast<int64_t>(subgroupSize), srcShape[0]);
laneData[0] = 1;
- srcLayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, laneLayout),
- DenseI32ArrayAttr::get(context, laneData));
+ srcLayout = buildLaneLayout(context, laneLayout, laneData);
}
auto result = xegpu::SliceAttr::get(context, srcLayout,
@@ -1490,33 +1488,6 @@ xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
defLaneData);
}
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
- int64_t subgroupSize, int64_t bitwidth,
- int64_t packingSize, bool vnni = false,
- bool transpose = false) {
- int64_t rank = instShape.size();
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- int64_t packingDim = vnni ? rank - 2 : rank - 1;
- laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
- assert(
- !(vnni && transpose) &&
- "transpose and VNNI cannot be enabled at the same time for 2D block IO");
- if (transpose)
- laneLayout[rank - 2] = subgroupSize;
- else
- laneLayout.back() = subgroupSize;
- // assert that the lane layout and data fit in the inst shape
- for (int64_t i = 0; i < rank; ++i) {
- int64_t laneProduct = laneLayout[i] * laneData[i];
- assert(instShape[i] % laneProduct == 0 &&
- "lane_layout * lane_data must evenly divide the inst shape");
- (void)laneProduct;
- }
- return {laneLayout, laneData};
-}
-
-// Forward declaration: defined later in the file.
using LayoutRepresentation = SmallVector<int64_t>;
/// Enumerates all ways to split `total` into `rank` factors whose product
@@ -1609,6 +1580,52 @@ getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
return candidates;
}
+static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
+ ArrayRef<int64_t> wgTileShape,
+ ArrayRef<int64_t> sgLayout,
+ int dimK = -1,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ SmallVector<int64_t> sgData(sgLayout.size());
+ for (int dim = 0; dim < (int)sgLayout.size(); ++dim) {
+ if (dim == dimK)
+ sgData[dim] = wgTileShape[dim];
+ else
+ sgData[dim] = wgTileShape[dim] / sgLayout[dim];
+ }
+ return buildLayout(context, sgLayout, sgData,
+ /*inst_data=*/{}, /*lane_layout=*/{},
+ /*lane_data=*/{}, /*order=*/nullptr);
+}
+
+// Computes the per-lane layout and data for a 2D block load/store/prefetch:
+// lanes are spread across the subgroup along the last dim (or rank-2 if
+// transposed), and laneData packs sub-bitwidth elements along the packing dim.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t bitwidth,
+ int64_t packingSize, bool vnni = false,
+ bool transpose = false) {
+ int64_t rank = instShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int64_t packingDim = vnni ? rank - 2 : rank - 1;
+ laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
+ assert(
+ !(vnni && transpose) &&
+ "transpose and VNNI cannot be enabled at the same time for 2D block IO");
+ if (transpose)
+ laneLayout[rank - 2] = subgroupSize;
+ else
+ laneLayout.back() = subgroupSize;
+ // assert that the lane layout and data fit in the inst shape
+ for (int64_t i = 0; i < rank; ++i) {
+ int64_t laneProduct = laneLayout[i] * laneData[i];
+ assert(instShape[i] % laneProduct == 0 &&
+ "lane_layout * lane_data must evenly divide the inst shape");
+ (void)laneProduct;
+ }
+ return {laneLayout, laneData};
+}
+
/// Helper function to compute inst_data vectors for DPAS operands A, B, and
/// C/D.
static std::optional<SmallVector<int64_t>>
@@ -1619,9 +1636,7 @@ get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
// Compute inst_data from hardware block params. For Nd ops, the lane
// factorization above (laneLayout / laneData) is rigid; inst_data must be
// a multiple of lane_layout * lane_data on each dim (Category A
- // invariant). If block params are unavailable for this element type
- // (e.g. sub-byte floats with no uArch entry), fall back to
- // lane_layout * lane_data (k = 1).
+ // invariant).
SmallVector<int64_t> instData(rank, 1);
assert(rank >= 2 && "dataShape must be at least 2D for 2D-block IO");
int instWidth =
@@ -1641,26 +1656,6 @@ get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
return instData;
}
-static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
- ArrayRef<int64_t> wgTileShape,
- ArrayRef<int64_t> sgLayout,
- int dimK = -1,
- DenseI32ArrayAttr orderAttr = nullptr) {
- SmallVector<int> sgData(sgLayout.size());
- SmallVector<int> sgLayoutInt(sgLayout.begin(), sgLayout.end());
- for (int dim = 0; dim < (int)sgLayout.size(); ++dim) {
- if (dim == dimK)
- sgData[dim] = wgTileShape[dim];
- else
- sgData[dim] = static_cast<int>(wgTileShape[dim]) / sgLayout[dim];
- }
- return xegpu::LayoutAttr::get(context,
- DenseI32ArrayAttr::get(context, sgLayoutInt),
- DenseI32ArrayAttr::get(context, sgData),
- /*inst_data=*/nullptr, /*lane_layout=*/nullptr,
- /*lane_data=*/nullptr, /*order=*/nullptr);
-}
-
/// Generic anchor-layout setup for ND ops (load_nd, store_nd, prefetch_nd).
///
/// Given hardware-supported block widths/heights, picks the largest divisor
>From 5ed1bbebe6683cefd54c46c2de43502d7edb369d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 16 Jun 2026 05:52:49 +0000
Subject: [PATCH 31/42] rearrange the code location, put utitlies/helpers
togehter
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 408 +++++++++---------
1 file changed, 209 insertions(+), 199 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 5ccad0f8335b8..d701b1887e822 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -425,6 +425,23 @@ buildLayout(mlir::MLIRContext *context, ArrayRef<int64_t> sgLayout,
laneData.empty() ? nullptr : toI32Attr(laneData), orderAttr);
}
+static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
+ ArrayRef<int64_t> wgTileShape,
+ ArrayRef<int64_t> sgLayout,
+ int dimK = -1,
+ DenseI32ArrayAttr orderAttr = nullptr) {
+ SmallVector<int64_t> sgData(sgLayout.size());
+ for (int dim = 0; dim < (int)sgLayout.size(); ++dim) {
+ if (dim == dimK)
+ sgData[dim] = wgTileShape[dim];
+ else
+ sgData[dim] = wgTileShape[dim] / sgLayout[dim];
+ }
+ return buildLayout(context, sgLayout, sgData,
+ /*inst_data=*/{}, /*lane_layout=*/{},
+ /*lane_data=*/{}, /*order=*/nullptr);
+}
+
/// Infers the source layout attribute for a broadcast operation given the
/// result layout attribute, result shape, source shape.
xegpu::DistributeLayoutAttr
@@ -808,6 +825,183 @@ xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
return payloadLayout;
}
+//===----------------------------------------------------------------------===//
+// Layout derivation helpers: factorize sgCount into
+// sg_layout candidates, then
+// compute per-subgroup (sgData) and per-lane
+// (lane_layout/lane_data/inst_data).
+//===----------------------------------------------------------------------===//
+
+using LayoutRepresentation = SmallVector<int64_t>;
+
+/// Enumerates all ways to split `total` into `rank` factors whose product
+/// equals `total`. Returns the list of all such factorizations.
+static SmallVector<LayoutRepresentation> enumerateFactorizations(int64_t total,
+ int64_t rank) {
+ SmallVector<LayoutRepresentation> results;
+ SmallVector<int64_t> current(rank, 0);
+
+ // Returns all divisors of `n` in ascending order.
+ auto getDivisors = [](int64_t n) {
+ SmallVector<int64_t> divs;
+ for (int64_t i = 1; i * i <= n; ++i) {
+ if (n % i == 0) {
+ divs.push_back(i);
+ if (i != n / i)
+ divs.push_back(n / i);
+ }
+ }
+ llvm::sort(divs);
+ return divs;
+ };
+
+ std::function<void(int64_t, int64_t)> generate = [&](int64_t dim,
+ int64_t remaining) {
+ if (dim == rank - 1) {
+ current[dim] = remaining;
+ results.push_back(LayoutRepresentation(current));
+ return;
+ }
+ for (int64_t factor : getDivisors(remaining)) {
+ current[dim] = factor;
+ generate(dim + 1, remaining / factor);
+ }
+ };
+
+ generate(0, total);
+ return results;
+}
+
+// Computes all valid N-dimensional sg_layout candidates for the given
+// sgCount, whose sgData (= wgShape / sgLayout):
+// 1. Evenly divides wgShape (i.e., wgShape[d] % sgLayout[d] == 0).
+// 2. Is a multiple of instData (i.e., sgData[d] % instData[d] == 0).
+// Results are sorted by balance (smallest max-min spread first), with
+// lexicographic order as a tiebreaker.
+//
+// Example (2D):
+// wgShape = [128, 64], instData = [8, 16], sgCount = 32
+// Returns: [[8,4], [16,2]], corresponding to sgData [16,16] and [8,32].
+static SmallVector<LayoutRepresentation>
+getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
+ int64_t sgCount) {
+ int64_t rank = wgShape.size();
+ assert(rank > 0 && "wgShape must be non-empty");
+ assert(static_cast<int64_t>(instData.size()) == rank &&
+ "instData rank must match wgShape rank");
+
+ // Step 1: Get all N-D factorizations of sgCount.
+ auto allFactorizations = enumerateFactorizations(sgCount, rank);
+
+ // Step 2: Filter to keep only valid candidates.
+ SmallVector<LayoutRepresentation> candidates;
+ for (const auto &sgLayout : allFactorizations) {
+ bool valid = true;
+ for (int64_t dim = 0; dim < rank; ++dim) {
+ if (wgShape[dim] % sgLayout[dim] != 0) {
+ valid = false;
+ break;
+ }
+ int64_t sgData = wgShape[dim] / sgLayout[dim];
+ if (sgData % instData[dim] != 0) {
+ valid = false;
+ break;
+ }
+ }
+ if (valid)
+ candidates.push_back(sgLayout);
+ }
+
+ // Step 3: Sort by balance (smallest max-min spread), then lexicographic.
+ llvm::sort(candidates, [](const LayoutRepresentation &lhs,
+ const LayoutRepresentation &rhs) {
+ int64_t spreadLhs = *llvm::max_element(lhs) - *llvm::min_element(lhs);
+ int64_t spreadRhs = *llvm::max_element(rhs) - *llvm::min_element(rhs);
+ if (spreadLhs != spreadRhs)
+ return spreadLhs < spreadRhs;
+ return lhs < rhs;
+ });
+ return candidates;
+}
+
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D.
+static std::optional<SmallVector<int64_t>>
+get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
+ ArrayRef<int> bHeights, ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
+ int rank = dataShape.size();
+ // Compute inst_data from hardware block params. For Nd ops, the lane
+ // factorization above (laneLayout / laneData) is rigid; inst_data must be
+ // a multiple of lane_layout * lane_data on each dim (Category A
+ // invariant).
+ SmallVector<int64_t> instData(rank, 1);
+ assert(rank >= 2 && "dataShape must be at least 2D for 2D-block IO");
+ int instWidth =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
+ int instHeight =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape[rank - 2]), bHeights);
+ instData.back() = instWidth;
+ instData[rank - 2] = instHeight;
+
+ if (instWidth == -1 || instHeight == -1) {
+ instData.back() = laneLayout.back() * laneData.back();
+ instData[rank - 2] = laneLayout[rank - 2] * laneData[rank - 2];
+ }
+ for (int dim = 0; dim < rank; ++dim)
+ assert(instData[dim] % (laneLayout[dim] * laneData[dim]) == 0 &&
+ "inst_data must be a multiple of lane_layout * lane_data for ND op");
+ return instData;
+}
+
+/// Computes lane_layout and lane_data for scatter-style store anchor layouts
+/// (store scatter, store matrix). Lanes and the per-lane vector both live on
+/// the innermost dim:
+/// - laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
+/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
+/// maxChunkSize)
+/// All other entries are 1.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+computeScatterIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t maxChunkSize) {
+ int64_t rank = instShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int64_t innermost = rank - 1;
+ laneLayout[innermost] = std::min(subgroupSize, instShape[innermost]);
+ laneData[innermost] =
+ std::min(instShape[innermost] / laneLayout[innermost], maxChunkSize);
+ return {laneLayout, laneData};
+}
+
+// Computes the per-lane layout and data for a 2D block load/store/prefetch:
+// lanes are spread across the subgroup along the last dim (or rank-2 if
+// transposed), and laneData packs sub-bitwidth elements along the packing dim.
+static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
+compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
+ int64_t subgroupSize, int64_t bitwidth,
+ int64_t packingSize, bool vnni = false,
+ bool transpose = false) {
+ int64_t rank = instShape.size();
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
+ int64_t packingDim = vnni ? rank - 2 : rank - 1;
+ laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
+ assert(
+ !(vnni && transpose) &&
+ "transpose and VNNI cannot be enabled at the same time for 2D block IO");
+ if (transpose)
+ laneLayout[rank - 2] = subgroupSize;
+ else
+ laneLayout.back() = subgroupSize;
+ // assert that the lane layout and data fit in the inst shape
+ for (int64_t i = 0; i < rank; ++i) {
+ int64_t laneProduct = laneLayout[i] * laneData[i];
+ assert(instShape[i] % laneProduct == 0 &&
+ "lane_layout * lane_data must evenly divide the inst shape");
+ (void)laneProduct;
+ }
+ return {laneLayout, laneData};
+}
+
/// 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
@@ -824,20 +1018,15 @@ xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
ArrayRef<int64_t> reductionDims,
- ArrayRef<int64_t> consumerReductionDims,
- int subgroupSize,
- int64_t maxReduceVectorSize) {
+ int subgroupSize, int64_t maxReduceVectorSize,
+ bool verticalLaneLayout = false) {
int srcRank = srcShape.size();
SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
int innermost = srcRank - 1;
int secondInnermost = srcRank - 2;
- // `laneDim` carries the subgroupSize lanes; `vectorDim` packs the reduced
- // elements into lane_data. Default: lanes on the innermost dim, reduced
- // vector on the second-to-innermost dim.
- if (secondInnermost >= 0 && llvm::is_contained(reductionDims, innermost) &&
- reductionDims.size() == 1 && consumerReductionDims.empty()) {
+ if (verticalLaneLayout && secondInnermost >= 0) {
std::swap(innermost, secondInnermost);
}
int laneDim = innermost;
@@ -1006,9 +1195,14 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
consumerSliceLayout
? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
: SmallVector<int64_t>({});
+ // A[i] reduced from A[i, j] is stored out directly, use veritical Lane
+ // layout like [16, 1]
+ bool verticalLaneLayout = consumerReductionDims.empty() &&
+ reductionDims.size() == 1 &&
+ reductionDims[0] == (srcRank - 1);
auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDims, consumerReductionDims, subgroupSize,
- maxReduceVectorSize);
+ 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);
@@ -1035,9 +1229,12 @@ 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(
- srcShape, reductionDims, consumerReductionDims, subgroupSize,
- maxReduceVectorSize);
+ srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
+ verticalLaneLayout);
srcLayout = buildLaneLayout(context, laneLayout, laneData);
}
}
@@ -1249,25 +1446,6 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
return requiredResLayout;
}
-/// Computes lane_layout and lane_data for scatter-style store anchor layouts
-/// (store scatter, store matrix). Lanes and the per-lane vector both live on
-/// the innermost dim:
-/// - laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
-/// - laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost],
-/// maxChunkSize)
-/// All other entries are 1.
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-computeScatterIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
- int64_t subgroupSize, int64_t maxChunkSize) {
- int64_t rank = instShape.size();
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- int64_t innermost = rank - 1;
- laneLayout[innermost] = std::min(subgroupSize, instShape[innermost]);
- laneData[innermost] =
- std::min(instShape[innermost] / laneLayout[innermost], maxChunkSize);
- return {laneLayout, laneData};
-}
-
/// Sets up the anchor layout for load gather and load matrix operation.
/// load matrix lowers to load gather and 1d block load. All of them share the
/// same layout setup logic.
@@ -1488,174 +1666,6 @@ xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
defLaneData);
}
-using LayoutRepresentation = SmallVector<int64_t>;
-
-/// Enumerates all ways to split `total` into `rank` factors whose product
-/// equals `total`. Returns the list of all such factorizations.
-static SmallVector<LayoutRepresentation> enumerateFactorizations(int64_t total,
- int64_t rank) {
- SmallVector<LayoutRepresentation> results;
- SmallVector<int64_t> current(rank, 0);
-
- // Returns all divisors of `n` in ascending order.
- auto getDivisors = [](int64_t n) {
- SmallVector<int64_t> divs;
- for (int64_t i = 1; i * i <= n; ++i) {
- if (n % i == 0) {
- divs.push_back(i);
- if (i != n / i)
- divs.push_back(n / i);
- }
- }
- llvm::sort(divs);
- return divs;
- };
-
- std::function<void(int64_t, int64_t)> generate = [&](int64_t dim,
- int64_t remaining) {
- if (dim == rank - 1) {
- current[dim] = remaining;
- results.push_back(LayoutRepresentation(current));
- return;
- }
- for (int64_t factor : getDivisors(remaining)) {
- current[dim] = factor;
- generate(dim + 1, remaining / factor);
- }
- };
-
- generate(0, total);
- return results;
-}
-
-// Computes all valid N-dimensional sg_layout candidates for the given
-// sgCount, whose sgData (= wgShape / sgLayout):
-// 1. Evenly divides wgShape (i.e., wgShape[d] % sgLayout[d] == 0).
-// 2. Is a multiple of instData (i.e., sgData[d] % instData[d] == 0).
-// Results are sorted by balance (smallest max-min spread first), with
-// lexicographic order as a tiebreaker.
-//
-// Example (2D):
-// wgShape = [128, 64], instData = [8, 16], sgCount = 32
-// Returns: [[8,4], [16,2]], corresponding to sgData [16,16] and [8,32].
-static SmallVector<LayoutRepresentation>
-getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
- int64_t sgCount) {
- int64_t rank = wgShape.size();
- assert(rank > 0 && "wgShape must be non-empty");
- assert(static_cast<int64_t>(instData.size()) == rank &&
- "instData rank must match wgShape rank");
-
- // Step 1: Get all N-D factorizations of sgCount.
- auto allFactorizations = enumerateFactorizations(sgCount, rank);
-
- // Step 2: Filter to keep only valid candidates.
- SmallVector<LayoutRepresentation> candidates;
- for (const auto &sgLayout : allFactorizations) {
- bool valid = true;
- for (int64_t dim = 0; dim < rank; ++dim) {
- if (wgShape[dim] % sgLayout[dim] != 0) {
- valid = false;
- break;
- }
- int64_t sgData = wgShape[dim] / sgLayout[dim];
- if (sgData % instData[dim] != 0) {
- valid = false;
- break;
- }
- }
- if (valid)
- candidates.push_back(sgLayout);
- }
-
- // Step 3: Sort by balance (smallest max-min spread), then lexicographic.
- llvm::sort(candidates, [](const LayoutRepresentation &lhs,
- const LayoutRepresentation &rhs) {
- int64_t spreadLhs = *llvm::max_element(lhs) - *llvm::min_element(lhs);
- int64_t spreadRhs = *llvm::max_element(rhs) - *llvm::min_element(rhs);
- if (spreadLhs != spreadRhs)
- return spreadLhs < spreadRhs;
- return lhs < rhs;
- });
- return candidates;
-}
-
-static xegpu::LayoutAttr buildSgLayout(mlir::MLIRContext *context,
- ArrayRef<int64_t> wgTileShape,
- ArrayRef<int64_t> sgLayout,
- int dimK = -1,
- DenseI32ArrayAttr orderAttr = nullptr) {
- SmallVector<int64_t> sgData(sgLayout.size());
- for (int dim = 0; dim < (int)sgLayout.size(); ++dim) {
- if (dim == dimK)
- sgData[dim] = wgTileShape[dim];
- else
- sgData[dim] = wgTileShape[dim] / sgLayout[dim];
- }
- return buildLayout(context, sgLayout, sgData,
- /*inst_data=*/{}, /*lane_layout=*/{},
- /*lane_data=*/{}, /*order=*/nullptr);
-}
-
-// Computes the per-lane layout and data for a 2D block load/store/prefetch:
-// lanes are spread across the subgroup along the last dim (or rank-2 if
-// transposed), and laneData packs sub-bitwidth elements along the packing dim.
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
- int64_t subgroupSize, int64_t bitwidth,
- int64_t packingSize, bool vnni = false,
- bool transpose = false) {
- int64_t rank = instShape.size();
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- int64_t packingDim = vnni ? rank - 2 : rank - 1;
- laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
- assert(
- !(vnni && transpose) &&
- "transpose and VNNI cannot be enabled at the same time for 2D block IO");
- if (transpose)
- laneLayout[rank - 2] = subgroupSize;
- else
- laneLayout.back() = subgroupSize;
- // assert that the lane layout and data fit in the inst shape
- for (int64_t i = 0; i < rank; ++i) {
- int64_t laneProduct = laneLayout[i] * laneData[i];
- assert(instShape[i] % laneProduct == 0 &&
- "lane_layout * lane_data must evenly divide the inst shape");
- (void)laneProduct;
- }
- return {laneLayout, laneData};
-}
-
-/// Helper function to compute inst_data vectors for DPAS operands A, B, and
-/// C/D.
-static std::optional<SmallVector<int64_t>>
-get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
- ArrayRef<int> bHeights, ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData) {
- int rank = dataShape.size();
- // Compute inst_data from hardware block params. For Nd ops, the lane
- // factorization above (laneLayout / laneData) is rigid; inst_data must be
- // a multiple of lane_layout * lane_data on each dim (Category A
- // invariant).
- SmallVector<int64_t> instData(rank, 1);
- assert(rank >= 2 && "dataShape must be at least 2D for 2D-block IO");
- int instWidth =
- xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
- int instHeight =
- xegpu::getLargestDivisor(static_cast<int>(dataShape[rank - 2]), bHeights);
- instData.back() = instWidth;
- instData[rank - 2] = instHeight;
-
- if (instWidth == -1 || instHeight == -1) {
- instData.back() = laneLayout.back() * laneData.back();
- instData[rank - 2] = laneLayout[rank - 2] * laneData[rank - 2];
- }
- for (int dim = 0; dim < rank; ++dim)
- assert(instData[dim] % (laneLayout[dim] * laneData[dim]) == 0 &&
- "inst_data must be a multiple of lane_layout * lane_data for ND op");
- return instData;
-}
-
/// Generic anchor-layout setup for ND ops (load_nd, store_nd, prefetch_nd).
///
/// Given hardware-supported block widths/heights, picks the largest divisor
>From 736c9b23c39d1703abbac1a9c91e06cc38daf78d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 16 Jun 2026 06:30:34 +0000
Subject: [PATCH 32/42] add comments and reorder functions
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 1850 +++++++++--------
1 file changed, 958 insertions(+), 892 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d701b1887e822..bb8e96fec5f6f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1040,630 +1040,400 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
return {laneLayout, laneData};
}
-/// Sets up layout for reduction operations by creating a SliceAttr for the
-/// result.
-///
-/// Algorithm Overview:
-/// This function attempts to construct a source layout that, when sliced along
-/// reduction dimensions, produces a result layout compatible with the
-/// consumer layout.
-///
-/// For subgroup layouts, it first tries to align the source layout's subgroup
-/// layout and data with the consumer's layout on non-reduction dimensions.
-/// Then, it distributes remaining subgroups across reduction dimensions. This
-/// avoids subgroup data redistribution overhead between the reduced result and
-/// its consumer. When the consumer layout is a slice layout, it attempts to
-/// reuse the slice layout's parent layout for the source to further minimize
-/// potential data redistribution.
-///
-/// 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.
-///
-/// Examples:
-/// 1. Subgroup layout - Row reduction on 2D tensor:
-/// srcShape=[32, 128], reductionDims=[1], resShape=[32], subgroupSize=16,
-/// NumSg=32
-/// * Consumer Layout:
-/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 8]>, dims =
-/// [1]>}
-//// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8],sg_data=[8, 16]>, dims =
-/// [1]>}
-/// Note that the sg_layout is reused but sg_data needs to be adjusted to
-/// evenly distribute the source tensor tile among the reduction dim.
-///
-/// 2. Subgroup layout - Same example above but consumer doesn't have a
-/// reusable slice layout.
-/// * Consumer Layout:
-/// #xegpu.layout<sgLayout=[32], sgData=[1]>
-/// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[32,1], sgData=[1, 128]>, dims =
-/// [1]>}
-/// * Consumer Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[8, 2, 4], sgData=[4, 64, 32]>,
-/// dims = [1, 2]>}
-/// * Result Layout:
-/// #xegpu.slice<#xegpu.layout<sgLayout=[8,4], sgData=[4, 32]>, dims =
-/// [1]>}
-/// Note that the consumer's layout can't be directly reused as is.
-/// So the algorithm distributes all subgroups on non reduction dimensions
-/// first and then distribute remaining subgroups on the reduction
-/// dimension.
-///
-/// 3. Lane layout - Default (lanes on innermost dim):
-/// srcShape=[32, 64], reductionDims=[0], subgroupSize=16
-/// Result: laneLayout=[1, 16], laneData=[1, 1]. The innermost dim is not
-/// reduced, so lanes stay on it.
-///
-/// 4. Lane layout - Switch (lanes moved off the reduction dim):
-/// srcShape=[32, 64], reductionDims=[1], subgroupSize=16
-/// Result: laneLayout=[16, 1], laneData=[1, 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.
-
-xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy,
- DistributeLayoutAttr consumerLayout, SmallVector<int64_t> reductionDims,
- int numSg, const xegpu::uArch::uArch *uArch) {
-
- auto srcShape = srcVecTy.getShape();
- int srcRank = srcShape.size();
- auto context = srcVecTy.getContext();
-
- const int subgroupSize = uArch->getSubgroupSize();
- int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
- xegpu::DistributeLayoutAttr srcLayout;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- xegpu::SliceAttr consumerSliceLayout =
- dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- if (consumerSliceLayout &&
- consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
- srcLayout = consumerSliceLayout.getParent();
- SmallVector<int64_t> sgLayoutFromConsumer =
- srcLayout.getEffectiveSgLayoutAsInt();
- auto srcSgData = computeShapeRatio(srcShape, sgLayoutFromConsumer);
- if (srcSgData)
- for (int dim = 0; dim < srcRank; dim++) {
- if (llvm::is_contained(reductionDims, dim))
- srcLayout =
- srcLayout.setDimData(dim, srcSgData.value()[dim], -1, -1);
- }
- } else {
- SmallVector<int64_t> consumerSgLayout =
- consumerLayout ? consumerLayout.getEffectiveSgLayoutAsInt()
- : SmallVector<int64_t>();
- SmallVector<int64_t> consumerSgData =
- consumerLayout ? consumerLayout.getEffectiveSgDataAsInt()
- : SmallVector<int64_t>();
- SmallVector<int64_t> consumerOrder =
- consumerLayout ? consumerLayout.getEffectiveOrderAsInt()
- : SmallVector<int64_t>();
- DenseI32ArrayAttr orderAttr =
- consumerLayout ? consumerLayout.getOrder() : nullptr;
- SmallVector<int64_t> sgLayout(srcRank), sgData(srcRank), order(srcRank);
- int remainingSgCount =
- consumerLayout ? consumerLayout.getNumSubgroups() : numSg;
- int consumerIdx = 0;
-
- // First pass: Match consumer's layout on non-reduction dimensions
- for (int i = 0; i < srcRank; i++) {
- if (!llvm::is_contained(reductionDims, i) &&
- consumerIdx < static_cast<int>(consumerSgLayout.size())) {
- sgLayout[i] = consumerSgLayout[consumerIdx];
- sgData[i] = consumerSgData[consumerIdx];
- remainingSgCount /= sgLayout[i];
- order[i] = consumerOrder[consumerIdx];
- consumerIdx++;
- }
- }
-
- // Second pass: Distribute remaining subgroups across reduction dimensions
- // the reduction to scalar case is handled only by this loop
- int64_t remainOrder = consumerSgLayout.size();
- for (int i = 0; i < srcRank; i++) {
- if (llvm::is_contained(reductionDims, i)) {
- sgLayout[i] =
- std::min(srcShape[i], static_cast<int64_t>(remainingSgCount));
- assert((srcShape[i] % sgLayout[i] == 0) &&
- "source shape not divisible by sg_layout");
- sgData[i] = srcShape[i] / sgLayout[i];
- remainingSgCount /= sgLayout[i];
- order[i] = remainOrder++;
- }
- }
- DenseI32ArrayAttr resOrderAttr = DenseI32ArrayAttr::get(
- context, SmallVector<int32_t>(order.begin(), order.end()));
- if (!orderAttr || orderAttr.empty())
- resOrderAttr = nullptr;
- assert(remainingSgCount == 1 && "not all subgroups distributed");
- srcLayout = buildLayout(context, sgLayout, sgData,
- /*instData=*/{}, /*laneLayout=*/{},
- /*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 veritical 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];
- 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
- // since the inst_data propagation already insert convert_layout if
- // the layout is not consistent
- srcLayout = consumerSliceLayout.getParent();
- } else {
- bool verticalLaneLayout = consumerReductionDims.empty() &&
- reductionDims.size() == 1 &&
- reductionDims[0] == (srcRank - 1);
- auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
- verticalLaneLayout);
- srcLayout = buildLaneLayout(context, laneLayout, laneData);
- }
- }
-
- return xegpu::SliceAttr::get(context, srcLayout,
- DenseI64ArrayAttr::get(context, reductionDims));
-}
-
-/// Sets up layout for Reduction operations by creating a SliceAttr for the
-/// result.
-xegpu::SliceAttr
-xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy,
- const xegpu::uArch::uArch *uArch) {
-
- auto srcShape = srcVecTy.getShape();
- auto context = srcVecTy.getContext();
- auto subgroupSize = uArch->getSubgroupSize();
- xegpu::LayoutAttr srcLayout;
+//===----------------------------------------------------------------------===//
+// Result/anchor-layout setup. Each op category derives lane_layout/lane_data
+// (and inst_data / sgData) differently. Two things vary across ops:
+//
+// * Consumer dependence: consumer-driven ops prefer the layout requested by
+// their downstream uses and fall back to uArch defaults only when it is
+// absent/invalid; sinks (StoreNd, PrefetchNd) have no consumer and always
+// pick their own layout from uArch.
+//
+// * Derivation direction between inst_data and lane_layout/lane_data. Both
+// obey the Category-A invariant inst_data = k * lane_layout * lane_data
+// (k >= 1, per dim), but ops solve it from opposite ends:
+// - Rigid-lane ops (Nd block IO, DPAS): hardware fixes lane_layout /
+// lane_data first, then inst_data is built as a multiple of their
+// product (get2DBlockIOInstDataLayout / getDpasInstDataLayouts).
+// - inst_data-first ops (scatter load): take inst_data from the consumer
+// and derive lane_layout/lane_data underneath it.
+//
+// - DPAS (+DPAS_MX) : rigid lanes — inst_data from HW block dims; A/B/C/D
+// lanes/data follow each operand's matmul role; DPAS_MX
+// additionally lays out the scale operand.
+// - LoadNd : consumer-driven, rigid lanes — honors the consumer's
+// inst_data / lane / sg_layout (incl. transpose & VNNI
+// packing) when it satisfies uArch block constraints,
+// else falls back to the default 2D-block scheme (lanes
+// on the last dim, rank-2 if transposed). The fallback
+// picks the LARGEST uArch block that divides the data
+// shape, so the resulting inst_data block can be bigger
+// than what the consumer asked for (fewer, wider
+// loads).
+// - StoreNd/PrefetchNd: data sinks, no consumer, rigid lanes — pick the
+// 2D-block layout directly from uArch (no VNNI
+// packing).
+// - Load (scatter) : load_gather / load_matrix, consumer-driven,
+// inst_data-first — reuse the consumer's inst_data and
+// derive lane_layout/lane_data, else default to lanes +
+// per-lane chunk on the innermost dim (chunk capped by
+// maxChunkSize).
+// - Store (scatter) : store_scatter / store_matrix — same scatter scheme,
+// but always self-derived from the scatter default.
+// - Reduction : (multi_)reduction, consumer-driven — distribute the
+// inner two dims. Default: lanes stay on the innermost
+// dim, even when it is reduced. They only switch to a
+// non-reduction dim (to keep the reduction within a
+// lane) when the innermost is the sole reduction dim
+// AND the consumer has nowhere to broadcast the result
+// back (i.e. not a 2D reduction-to-scalar and the
+// consumer is not a slice over this reduction dim).
+// Reuses the consumer's slice layout when it slices
+// exactly the reduction dims, otherwise re-derives.
+// - BitCast/Interleave: scale the innermost data field by the bitwidth /
+// interleave ratio so the source layout divides back
+// out.
+// - InsertStridedSlice: clamp lane_data per dim to fit the inserted slice
+// (Lane kind only; sg/inst layouts unsupported).
+//===----------------------------------------------------------------------===//
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true && "subgroup layout assignment not supported for reduction (op "
- "is not expected at this level).");
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- assert(true && "instData layout assignment not supported for reduction (op "
- "is not expected at this level).");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- SmallVector<int64_t> laneLayout(1), laneData(1);
- laneLayout[0] = std::min(static_cast<int64_t>(subgroupSize), srcShape[0]);
- laneData[0] = 1;
- srcLayout = buildLaneLayout(context, laneLayout, laneData);
- }
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D.
+static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>>
+getDpasInstDataLayouts(
+ VectorType aTy, VectorType bTy, VectorType cdTy,
+ const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
+ const int subgroupSize, bool isDpasMx = false) {
- auto result = xegpu::SliceAttr::get(context, srcLayout,
- DenseI64ArrayAttr::get(context, 0));
- return result;
-}
+ // M dimension is the second-to-last dim of A (handles batch dims).
+ const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
+ auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
+ const int maxALen =
+ xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
-/// Adjusts `consumerLayout`'s innermost-dim data field selected by
-/// `layoutKind` so that the source layout can be safely inferred by dividing
-/// that value by `ratio`. Doubles the value until the divisibility constraint
-/// is met, bounded above by `bound` like result-shape.
-///
-/// Used by ops whose source relates to the result by a fixed factor along the
-/// innermost dim (e.g., bitcast: bitwidth ratio; interleave: 2x).
-///
-/// Divisibility constraints per LayoutKind:
-/// - Subgroup: sgData[innermost] % ratio == 0
-/// - InstData: instData[innermost] % (laneLayout[innermost] * ratio) == 0
-/// (laneLayout falls back to subgroupSize if absent)
-/// - Lane: laneData[innermost] % ratio == 0
-static xegpu::DistributeLayoutAttr
-adjustInnermostDimForDivisibility(xegpu::DistributeLayoutAttr consumerLayout,
- xegpu::LayoutKind layoutKind,
- size_t innerMostDim, int ratio, int64_t bound,
- const xegpu::uArch::uArch *uArch) {
- SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> laneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
+ // N dimension is the last dim of B.
+ const unsigned dataBLen = bTy.getShape().back();
+ auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
+ const int maxBLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
- int64_t sgDataValue = -1;
- int64_t instDataValue = -1;
- int64_t laneDataValue = -1;
+ auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
+ const int maxCLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
+ if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+ return std::nullopt;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[innerMostDim];
- while ((sgDataValue <= bound) && (sgDataValue % ratio) != 0)
- sgDataValue *= 2;
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[innerMostDim];
- const int innermostDimLaneLayout = laneLayout.empty()
- ? uArch->getSubgroupSize()
- : laneLayout[innerMostDim];
- while ((instDataValue <= bound) &&
- (instDataValue % (innermostDimLaneLayout * ratio) != 0))
- instDataValue *= 2;
- assert((bound % instDataValue) == 0 &&
- "bound, instData, and laneLayout for innermost must be 2^n!");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[innerMostDim];
- while ((laneDataValue <= bound) && (laneDataValue % ratio) != 0)
- laneDataValue *= 2;
+ // For DPAS_MX, use getSupportedK to get the scaled K dimension.
+ // assume single element in the returned vector.
+ int kDimSize = subgroupSize;
+ if (isDpasMx) {
+ auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
+ if (supportedKLen.empty())
+ return std::nullopt;
+ kDimSize = supportedKLen[0];
}
- return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
- laneDataValue);
-}
-
-/// Sets up the result layout for a bitcast operation.
-/// When casting to a smaller bitwidth, adjusts the layout dimensions (sgData,
-/// instData, or laneData) by multiplying by the bitwidth ratio to ensure the
-/// result layout can be correctly divided back to the source layout during
-/// inference.
-///
-/// Examples:
-/// 1. Casting f32 -> f16 (32-bit to 16-bit, bitWidthRatio = 2):
-/// Consumer layout: instData=[1, 16], subgroupSize=16
-/// Source shape: [8, 32]
-/// Result layout: instData=[1, 32] (16 * 2)
-/// The innermost dimension is multiplied by 2 to maintain consistency.
-///
-/// 2. Casting f32 -> i8 (32-bit to 8-bit, bitWidthRatio = 4):
-/// Consumer instData=[1, 16], subgroupSize=16
-/// Source shape: [4, 128]
-/// adjust the instData from [1, 16] to [1, 16 * 4 = 64]
-///
-/// 3. Casting i8 -> i32 (8-bit to 32-bit, bitWidthRatio = 1/4):
-/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4]
-/// No adjustment needed - returns consumer layout directly.
-///
-xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
- DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
-
- int srcElemTyBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
- int resElemTyBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
-
- ArrayRef<int64_t> srcShape = srcVecTy.getShape();
- ArrayRef<int64_t> resShape = resVecTy.getShape();
-
- assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
- "laneData must be available for all dimensions");
-
- // Casting to same/larger element type: result has fewer (or equal) elements
- // along the innermost dim, no adjustment needed.
- if (srcElemTyBitWidth <= resElemTyBitWidth)
- return consumerLayout;
-
- // Casting to smaller element type: result has more elements along innermost
- // dim. Adjust the innermost data field upward so the source layout can be
- // recovered by dividing by bitWidthRatio.
- size_t innerMostDim = srcShape.size() - 1;
- int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
- return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
- innerMostDim, bitWidthRatio,
- resShape[innerMostDim], uArch);
-}
-
-/// Sets up the result layout for an interleave operation to ensure the source
-/// layout can be safely derived. Interleave doubles the innermost dimension,
-/// so the result layout must ensure that laneData is a multiple
-/// of 2, and instData must be divisible by innermostDimLaneLayout * 2.
-///
-/// Example:
-/// Interleave: vector<128x256xf4> -> vector<128x512xf4>
-/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4], instData=[1, 64]
-/// Result layout adjustment to ensure source can be safely inferred:
-/// - laneData must be >= 2 and multiple of 2 (so source = laneData/2 is
-/// valid)
-/// - instData must be divisible by (16 * 2 = 32) (so source = instData/2 is
-/// valid)
-/// - Adjusted instData: ensure (instData % 32 == 0)
-///
-xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
- DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
-
- ArrayRef<int64_t> resShape = resVecTy.getShape();
- assert(consumerLayout.getRank() == static_cast<int64_t>(resShape.size()) &&
- "consumer layout rank must match source shape rank");
-
- // Interleave doubles the innermost dimension (ratio = 2). Adjust the
- // innermost data field so the source layout can be recovered by dividing
- // by 2.
- const size_t innerMostDim = resShape.size() - 1;
- constexpr int ratio = 2;
- return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
- innerMostDim, ratio,
- resShape[innerMostDim], uArch);
+ SmallVector<int64_t> instDataA(aTy.getRank(), 1);
+ instDataA[aTy.getRank() - 2] = maxALen;
+ instDataA[aTy.getRank() - 1] = kDimSize;
+ SmallVector<int64_t> instDataB(bTy.getRank(), 1);
+ instDataB[bTy.getRank() - 2] = kDimSize;
+ instDataB[bTy.getRank() - 1] = maxBLen;
+ SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
+ instDataCD[cdTy.getRank() - 2] = maxALen;
+ instDataCD[cdTy.getRank() - 1] = maxCLen;
+ return std::make_tuple(instDataA, instDataB, instDataCD);
}
-/// Sets up the result layout for an insert strided slice operation.
-/// Creates a result layout based on the specified layout kind (InstData or
-/// Lane).
-xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVectorTy,
- VectorType resVectorTy, xegpu::DistributeLayoutAttr consumerLayout,
- const xegpu::uArch::uArch *uArch) {
-
- xegpu::DistributeLayoutAttr requiredResLayout;
- SmallVector<int64_t> consumerInstData =
- consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> consumerLaneData =
- consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> consumerLaneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
- ArrayRef<int64_t> srcShape = srcVectorTy.getShape();
- int64_t laneDataValue = -1;
-
- requiredResLayout = consumerLayout;
- int srcRank = srcShape.size();
+/// Helper function to set up subgroup layouts for DPAS operands A, B, and
+/// C/D. Returns the three layouts if successful, nullopt otherwise.
+static std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+getDpasSubgroupLayouts(
+ mlir::MLIRContext *context, VectorType aTy, VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, SmallVector<int64_t>>
+ instDataVecs) {
+ auto [instDataA, instDataB, instDataCD] = instDataVecs;
- if (layoutKind == xegpu::LayoutKind::Subgroup ||
- layoutKind == xegpu::LayoutKind::InstData) {
- assert(true &&
- "subgroup layout assignment not supported for insertStridedSlice.");
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- for (int dim = 0; dim < srcRank; dim++) {
- assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
- "srcShape must be divisible by laneLayout for all dimensions");
- laneDataValue = std::min(srcShape[dim] / consumerLaneLayout[dim],
- consumerLaneData[dim]);
- requiredResLayout =
- requiredResLayout.setDimData(dim, -1, -1, laneDataValue);
- }
+ std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
+ if (consumerLayout && consumerLayout.isForWorkgroup()) {
+ consumerSgLayout = consumerLayout.getEffectiveSgLayoutAsInt();
}
- return requiredResLayout;
-}
-
-/// Sets up the anchor layout for load gather and load matrix operation.
-/// load matrix lowers to load gather and 1d block load. All of them share the
-/// same layout setup logic.
-///
-/// For Subgroup layout, uses the consumer layout directly.
-///
-/// For InstData layout, takes consumer's inst_data as-is. lane_layout and
-/// lane_data are taken from the consumer when present; otherwise the helper
-/// derives the standard scatter-style default (subgroupSize lanes on the
-/// innermost dim, per-lane vector capped by maxChunkSize).
-///
-/// For Lane layout, lane_layout/lane_data are taken from the consumer when
-/// present; otherwise derived from the same default.
-static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
- xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
- xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
- ArrayRef<int64_t> resShape, int subgroupSize) {
-
- if (layoutKind == xegpu::LayoutKind::Subgroup)
- return consumerLayout;
- SmallVector<int64_t> consumerInstData =
- consumerLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> consumerLaneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> consumerLaneData =
- consumerLayout.getEffectiveLaneDataAsInt();
-
- // Pick lane_layout / lane_data: prefer consumer's, fall back to the
- // scatter-store default (subgroupSize lanes on innermost dim, per-lane
- // vector capped by maxChunkSize).
- SmallVector<int64_t> laneLayout;
- SmallVector<int64_t> laneData;
- if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
- laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
- laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
- } else {
- auto [LaneLayout, LaneData] =
- computeScatterIOLaneLayoutAndData(resShape, subgroupSize, maxChunkSize);
- }
+ // Get all valid layouts for A, B and C/D operands
+ auto layoutsA = getSgLayoutCandidates(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = getSgLayoutCandidates(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD = getSgLayoutCandidates(cdTy.getShape(), instDataCD, numSg);
+ if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
+ return std::nullopt;
- if (layoutKind == xegpu::LayoutKind::InstData) {
- // Take consumer's inst_data as-is. If the consumer doesn't have one,
- // fall back to lane_layout * lane_data per dim.
- SmallVector<int64_t> instData;
- if (!consumerInstData.empty()) {
- instData.assign(consumerInstData.begin(), consumerInstData.end());
- } else {
- instData.resize(resShape.size());
- for (size_t i = 0; i < resShape.size(); ++i)
- instData[i] = laneLayout[i] * laneData[i];
+ // Pick the best subgroup layout
+ std::optional<LayoutRepresentation> bestPick;
+ auto checkAlignedSgDataAB = [&](const LayoutRepresentation &sgLayout) {
+ return aTy.getShape().back() / sgLayout[1] ==
+ bTy.getShape().front() / sgLayout[0];
+ };
+ for (auto &sgLayout : layoutsB) {
+ if (llvm::is_contained(layoutsA, sgLayout) &&
+ llvm::is_contained(layoutsCD, sgLayout)) {
+ if (!checkAlignedSgDataAB(sgLayout))
+ continue;
+ // Is in (A and B and CD) and matches consumer -> best pick
+ if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
+ bestPick = sgLayout;
+ break;
+ }
+ // Is in (A and B and CD) layoutsB is ordered from most
+ // balanced to least. So the first one we see is the most balanced one,
+ // remember it and later only update if there is one that matches the
+ // consumer.
+ if (!bestPick)
+ bestPick = sgLayout;
}
- return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
- if (layoutKind == xegpu::LayoutKind::Lane)
- return buildLaneLayout(context, laneLayout, laneData);
- return nullptr;
-}
+ if (!bestPick)
+ return std::nullopt;
-/// Sets up the anchor layout for a load gather operation.
-xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
- xegpu::LayoutKind layoutKind, VectorType resVecTy, int contigChunkSize,
- xegpu::DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch) {
+ const auto &picked = *bestPick;
- const int subgroupSize = uArch->getSubgroupSize();
- ArrayRef<int64_t> resShape = resVecTy.getShape();
- auto context = resVecTy.getContext();
- auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
+ auto dpasALayout = buildSgLayout(context, aTy.getShape(), picked,
+ /*dimK=*/aTy.getRank() - 1);
+ auto dpasBLayout = buildSgLayout(context, bTy.getShape(), picked,
+ /*dimK=*/bTy.getRank() - 2);
+ auto dpasCDLayout = buildSgLayout(context, cdTy.getShape(), picked);
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
+}
+/// Sets up the anchor layouts for dpas operands (A, B, and C/D).
+/// The numSg and consumerLayout (optional) are only used by sg layout
+/// creation.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
+ VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
-
- return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
- maxChunkSize, resShape, subgroupSize);
-}
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
-/// Sets up the anchor layout for load matrix operation.
-/// TODO: enhance load matrix to indicate lowering to chunked load or not.
-xegpu::DistributeLayoutAttr
-xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType resVecTy, int contigChunkSize,
- xegpu::DistributeLayoutAttr consumerLayout,
- const xegpu::uArch::uArch *uArch) {
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
- const int subgroupSize = uArch->getSubgroupSize();
- ArrayRef<int64_t> resShape = resVecTy.getShape();
- auto context = resVecTy.getContext();
- auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
+ auto instDataVecs =
+ getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction, subgroupSize);
+ if (!instDataVecs)
+ return std::nullopt;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
- return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
- maxChunkSize, resShape, subgroupSize);
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ return getDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
+ numSg, *instDataVecs);
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+ return std::make_tuple(
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
+ buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
+ laneDataCD));
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ auto aLayout = buildLaneLayout(context, laneLayoutA, laneDataA);
+ auto bLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
+ auto cdLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
+ return std::make_tuple(aLayout, bLayout, cdLayout);
+ }
+ return std::nullopt;
}
-/// Sets up the anchor layout for store scatter and store matrix operation.
-/// store matrix lowers to store scatter and 1d block store. All of them
-/// share the same layout setup logic. For Subgroup layout, not supported
-/// yet.
-///
-/// Lane layout is derived first via `computeScatterIOLaneLayoutAndData`;
-/// inst_data is then the element-wise product lane_layout * lane_data.
+/// Helper to create a scale layout derived from a matrix operand layout.
+/// The scale layout is computed by mapping each dimension of the matrix
+/// layout to the corresponding scale tensor dimension using the ratio
+/// between the matrix and scale shapes.
static xegpu::DistributeLayoutAttr
-setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
- mlir::MLIRContext *context, int maxChunkSize,
- ArrayRef<int64_t> srcShape, int subgroupSize) {
+createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
+ VectorType scaleTy, xegpu::DistributeLayoutAttr matrixLayout,
+ bool isBScale, const xegpu::uArch::uArch *uArch) {
+ if (!scaleTy || !matrixLayout)
+ return nullptr;
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true &&
- "subgroup layout assignment not supported for storeScatter.");
+ // Calculate scaling factor by dividing matrix shape by scale shape
+ ArrayRef<int64_t> matrixShape = matrixTy.getShape();
+ ArrayRef<int64_t> scaleShape = scaleTy.getShape();
+
+ // Scale shapes can be 1D or 2D, handle both cases
+ if (scaleShape.empty())
return nullptr;
- }
- auto [laneLayout, laneData] =
- computeScatterIOLaneLayoutAndData(srcShape, subgroupSize, maxChunkSize);
+ auto uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
- if (layoutKind == xegpu::LayoutKind::InstData) {
- SmallVector<int64_t> instData(srcShape.size());
- for (size_t i = 0; i < srcShape.size(); ++i)
- instData[i] = laneLayout[i] * laneData[i];
- return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ int64_t rank = matrixLayout.getRank();
+ assert(rank >= 2 && "dpas layouts must be at least two dimensions");
+
+ SmallVector<int64_t> sgLayout = matrixLayout.getEffectiveSgLayoutAsInt();
+ SmallVector<int64_t> sgData = matrixLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = matrixLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneLayout = matrixLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = matrixLayout.getEffectiveLaneDataAsInt();
+ auto order = matrixLayout.getOrder();
+
+ SmallVector<int64_t> scaleSgLayout;
+ SmallVector<int64_t> scaleSgData;
+ if (!sgLayout.empty() && !sgData.empty()) {
+ scaleSgLayout.assign(sgLayout.begin(), sgLayout.end());
+ scaleSgData.assign(sgData.begin(), sgData.end());
+ scaleSgData[rank - 2] = std::max<int64_t>(
+ scaleShape[rank - 2] / (matrixShape[rank - 2] / sgData[rank - 2]), 1);
+ scaleSgData[rank - 1] = std::max<int64_t>(
+ scaleShape[rank - 1] / (matrixShape[rank - 1] / sgData[rank - 1]), 1);
}
- if (layoutKind == xegpu::LayoutKind::Lane) {
- return buildLaneLayout(context, laneLayout, laneData);
+
+ // For DPAS_MX scales: if matrix has inst_data, scale needs adjusted
+ // inst_data. Scale inst_data is derived from matrix inst_data divided by
+ // scale factor.
+ SmallVector<int64_t> scaleInstData;
+ if (!instData.empty()) {
+ scaleInstData.assign(instData.begin(), instData.end());
+ if (isBScale)
+ scaleInstData[rank - 2] = std::max<int64_t>(
+ scaleShape[rank - 2] / (matrixShape[rank - 2] / instData[rank - 2]),
+ 1);
+ else
+ scaleInstData[rank - 1] = std::max<int64_t>(
+ scaleShape[rank - 1] / (matrixShape[rank - 1] / instData[rank - 1]),
+ 1);
}
- return nullptr;
+
+ SmallVector<int64_t> scaleLaneLayout;
+ SmallVector<int64_t> scaleLaneData;
+ if (!laneLayout.empty() && !laneData.empty()) {
+ scaleLaneLayout.assign(laneLayout.begin(), laneLayout.end());
+ scaleLaneData.assign(laneData.size(), 1);
+
+ bool isRowMajor = uArchInstruction->isLaneLayoutRowMajorOrder();
+ if (isBScale ^ isRowMajor)
+ std::swap(scaleLaneLayout[rank - 2], scaleLaneLayout[rank - 1]);
+ // Cap lane_layout by the per-instruction tile (inst_data) on each dim.
+ // Then derive lane_data = inst_data / lane_layout so the Category A
+ // invariant inst_data = lane_layout * lane_data * k (with k = 1) holds
+ // for the scale operand's load_nd consumer.
+ auto layoutCap = scaleInstData.empty() ? scaleShape : scaleInstData;
+ for (int64_t d = rank - 2; d < rank; ++d)
+ scaleLaneLayout[d] = std::min<int64_t>(layoutCap[d], scaleLaneLayout[d]);
+ }
+ return buildLayout(context, scaleSgLayout, scaleSgData, scaleInstData,
+ scaleLaneLayout, scaleLaneData, order);
}
-/// Sets up the anchor layout for a store scatter operation.
-xegpu::DistributeLayoutAttr
-xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int contigChunkSize,
- const uArch::uArch *uArch) {
+/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
+/// B_scale). The numSg and consumerLayout (optional) are only used by sg
+/// layout creation.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
+ VectorType bTy, VectorType cdTy, VectorType aScaleTy,
+ VectorType bScaleTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
+
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
+ auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction,
+ subgroupSize, /*isDpasMx=*/true);
+ if (!instDataVecs)
+ return std::nullopt;
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ auto dpasLayouts = getDpasSubgroupLayouts(
+ context, aTy, bTy, cdTy, consumerLayout, numSg, *instDataVecs);
+ if (!dpasLayouts)
+ return std::nullopt;
- const int subgroupSize = uArch->getSubgroupSize();
- ArrayRef<int64_t> srcShape = srcVecTy.getShape();
- auto context = srcVecTy.getContext();
- auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
+ auto [dpasALayout, dpasBLayout, dpasCDLayout] = *dpasLayouts;
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
- return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
- srcShape, subgroupSize);
-}
+ // Create scale layouts
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
-/// Sets up the anchor layout for a store matrix operation.
-xegpu::DistributeLayoutAttr
-xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int contigChunkSize,
- const xegpu::uArch::uArch *uArch) {
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
- const int subgroupSize = uArch->getSubgroupSize();
- ArrayRef<int64_t> srcShape = srcVecTy.getShape();
- auto context = srcVecTy.getContext();
- auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
- int maxChunkSize = std::min(
- uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
- srcShape, subgroupSize);
-}
+ auto dpasALayout =
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA);
+ auto dpasBLayout =
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB);
+ auto dpasCDLayout = buildInstDataLayoutWithLane(context, instDataCD,
+ laneLayoutCD, laneDataCD);
-/// Completes a scatter IO layout by deriving lane_layout and lane_data from
-/// inst_data when they are missing. If `consumerLayout` already has both
-/// lane_layout and lane_data, or has no inst_data, the layout is returned
-/// unchanged.
-///
-/// When lane info is absent, this function uses inst_data as the effective
-/// shape and computes the standard scatter-style lane factorization:
-/// - laneLayout[innermost] = min(subgroupSize, inst_data[innermost])
-/// - laneData[innermost] = min(inst_data[innermost] /
-/// laneLayout[innermost],
-/// maxChunkSize)
-///
-/// The returned layout carries inst_data + lane_layout + lane_data, ensuring
-/// the lane factorization is consistent with what the downstream load/store
-/// scatter anchor setup would produce.
-xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
- xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
- const xegpu::uArch::uArch *uArch) {
- if (!consumerLayout)
- return consumerLayout;
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- if (instData.empty())
- return consumerLayout;
- if (!consumerLayout.getEffectiveLaneLayoutAsInt().empty() &&
- !consumerLayout.getEffectiveLaneDataAsInt().empty())
- return consumerLayout;
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
- // Reuse the load-side setup with inst_data as the destination shape.
- const int subgroupSize = uArch->getSubgroupSize();
- auto *context = consumerLayout.getContext();
- auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- if (!uArchInstruction)
- return consumerLayout;
- int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ auto dpasALayout = buildLaneLayout(context, laneLayoutA, laneDataA);
+ auto dpasBLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
+ auto dpasCDLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
- auto [defLaneLayout, defLaneData] =
- computeScatterIOLaneLayoutAndData(instData, subgroupSize, maxChunkSize);
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
- return buildInstDataLayoutWithLane(context, instData, defLaneLayout,
- defLaneData);
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ }
+ return std::nullopt;
}
/// Generic anchor-layout setup for ND ops (load_nd, store_nd, prefetch_nd).
@@ -1848,367 +1618,660 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights,
laneLayout, laneData);
- return buildInstDataLayoutWithLane(context, *instData, laneLayout,
- laneData);
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ laneData);
+ }
+ if (layoutKind == xegpu::LayoutKind::Lane) {
+ bool validLaneLayout = true;
+ for (int dim = 0; dim < rank; ++dim) {
+ int64_t laneProduct = consumerLaneLayout[dim] * consumerLaneData[dim];
+ if (dataShape[dim] % laneProduct != 0)
+ validLaneLayout = false;
+ }
+ if (validLaneLayout) {
+ return consumerLayout;
+ } else {
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
+ hasTransform, hasTranspose);
+ return buildLaneLayout(context, laneLayout, laneData);
+ }
+ }
+ return nullptr;
+}
+
+/// Sets up the anchor layout for load gather and load matrix operation.
+/// load matrix lowers to load gather and 1d block load. All of them share the
+/// same layout setup logic.
+///
+/// For Subgroup layout, uses the consumer layout directly.
+///
+/// For InstData layout, takes consumer's inst_data as-is. lane_layout and
+/// lane_data are taken from the consumer when present; otherwise the helper
+/// derives the standard scatter-style default (subgroupSize lanes on the
+/// innermost dim, per-lane vector capped by maxChunkSize).
+///
+/// For Lane layout, lane_layout/lane_data are taken from the consumer when
+/// present; otherwise derived from the same default.
+static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
+ xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
+ xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize,
+ ArrayRef<int64_t> resShape, int subgroupSize) {
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup)
+ return consumerLayout;
+
+ SmallVector<int64_t> consumerInstData =
+ consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+
+ // Pick lane_layout / lane_data: prefer consumer's, fall back to the
+ // scatter-store default (subgroupSize lanes on innermost dim, per-lane
+ // vector capped by maxChunkSize).
+ SmallVector<int64_t> laneLayout;
+ SmallVector<int64_t> laneData;
+ if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
+ laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
+ laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
+ } else {
+ auto [LaneLayout, LaneData] =
+ computeScatterIOLaneLayoutAndData(resShape, subgroupSize, maxChunkSize);
+ }
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ // Take consumer's inst_data as-is. If the consumer doesn't have one,
+ // fall back to lane_layout * lane_data per dim.
+ SmallVector<int64_t> instData;
+ if (!consumerInstData.empty()) {
+ instData.assign(consumerInstData.begin(), consumerInstData.end());
+ } else {
+ instData.resize(resShape.size());
+ for (size_t i = 0; i < resShape.size(); ++i)
+ instData[i] = laneLayout[i] * laneData[i];
+ }
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
+ }
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return buildLaneLayout(context, laneLayout, laneData);
+ return nullptr;
+}
+
+/// Sets up the anchor layout for a load gather operation.
+xegpu::DistributeLayoutAttr xegpu::setupLoadGatherAnchorLayout(
+ xegpu::LayoutKind layoutKind, VectorType resVecTy, int contigChunkSize,
+ xegpu::DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch) {
+
+ const int subgroupSize = uArch->getSubgroupSize();
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
+ auto context = resVecTy.getContext();
+ auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
+
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
+
+ return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
+ maxChunkSize, resShape, subgroupSize);
+}
+
+/// Sets up the anchor layout for load matrix operation.
+/// TODO: enhance load matrix to indicate lowering to chunked load or not.
+xegpu::DistributeLayoutAttr
+xegpu::setupLoadMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType resVecTy, int contigChunkSize,
+ xegpu::DistributeLayoutAttr consumerLayout,
+ const xegpu::uArch::uArch *uArch) {
+
+ const int subgroupSize = uArch->getSubgroupSize();
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
+ auto context = resVecTy.getContext();
+ auto elemBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
+
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneLoadSize(elemBitWidth), contigChunkSize);
+ return setupGenericLoadAnchorLayout(layoutKind, context, consumerLayout,
+ maxChunkSize, resShape, subgroupSize);
+}
+
+/// Sets up the anchor layout for store scatter and store matrix operation.
+/// store matrix lowers to store scatter and 1d block store. All of them
+/// share the same layout setup logic. For Subgroup layout, not supported
+/// yet.
+///
+/// Lane layout is derived first via `computeScatterIOLaneLayoutAndData`;
+/// inst_data is then the element-wise product lane_layout * lane_data.
+static xegpu::DistributeLayoutAttr
+setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
+ mlir::MLIRContext *context, int maxChunkSize,
+ ArrayRef<int64_t> srcShape, int subgroupSize) {
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(true &&
+ "subgroup layout assignment not supported for storeScatter.");
+ return nullptr;
+ }
+
+ auto [laneLayout, laneData] =
+ computeScatterIOLaneLayoutAndData(srcShape, subgroupSize, maxChunkSize);
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ SmallVector<int64_t> instData(srcShape.size());
+ for (size_t i = 0; i < srcShape.size(); ++i)
+ instData[i] = laneLayout[i] * laneData[i];
+ return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- bool validLaneLayout = true;
- for (int dim = 0; dim < rank; ++dim) {
- int64_t laneProduct = consumerLaneLayout[dim] * consumerLaneData[dim];
- if (dataShape[dim] % laneProduct != 0)
- validLaneLayout = false;
- }
- if (validLaneLayout) {
- return consumerLayout;
- } else {
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
- hasTransform, hasTranspose);
- return buildLaneLayout(context, laneLayout, laneData);
- }
+ return buildLaneLayout(context, laneLayout, laneData);
}
return nullptr;
}
-/// Helper function to compute inst_data vectors for DPAS operands A, B, and
-/// C/D.
-static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
- SmallVector<int64_t>>>
-getDpasInstDataLayouts(
- VectorType aTy, VectorType bTy, VectorType cdTy,
- const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
- const int subgroupSize, bool isDpasMx = false) {
+/// Sets up the anchor layout for a store scatter operation.
+xegpu::DistributeLayoutAttr
+xegpu::setupStoreScatterAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy, int contigChunkSize,
+ const uArch::uArch *uArch) {
- // M dimension is the second-to-last dim of A (handles batch dims).
- const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
- auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
- const int maxALen =
- xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
+ const int subgroupSize = uArch->getSubgroupSize();
+ ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ auto context = srcVecTy.getContext();
+ auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
- // N dimension is the last dim of B.
- const unsigned dataBLen = bTy.getShape().back();
- auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
- const int maxBLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
+ return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
+ srcShape, subgroupSize);
+}
- auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
- const int maxCLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
- if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
- return std::nullopt;
+/// Sets up the anchor layout for a store matrix operation.
+xegpu::DistributeLayoutAttr
+xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy, int contigChunkSize,
+ const xegpu::uArch::uArch *uArch) {
- // For DPAS_MX, use getSupportedK to get the scaled K dimension.
- // assume single element in the returned vector.
- int kDimSize = subgroupSize;
- if (isDpasMx) {
- auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
- if (supportedKLen.empty())
- return std::nullopt;
- kDimSize = supportedKLen[0];
- }
+ const int subgroupSize = uArch->getSubgroupSize();
+ ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ auto context = srcVecTy.getContext();
+ auto elemBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
- SmallVector<int64_t> instDataA(aTy.getRank(), 1);
- instDataA[aTy.getRank() - 2] = maxALen;
- instDataA[aTy.getRank() - 1] = kDimSize;
- SmallVector<int64_t> instDataB(bTy.getRank(), 1);
- instDataB[bTy.getRank() - 2] = kDimSize;
- instDataB[bTy.getRank() - 1] = maxBLen;
- SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
- instDataCD[cdTy.getRank() - 2] = maxALen;
- instDataCD[cdTy.getRank() - 1] = maxCLen;
- return std::make_tuple(instDataA, instDataB, instDataCD);
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::StoreScatter));
+ int maxChunkSize = std::min(
+ uArchInstruction->getMaxLaneStoreSize(elemBitWidth), contigChunkSize);
+
+ return setupGenericStoreAnchorLayout(layoutKind, context, maxChunkSize,
+ srcShape, subgroupSize);
}
-/// Helper function to set up subgroup layouts for DPAS operands A, B, and
-/// C/D. Returns the three layouts if successful, nullopt otherwise.
-static std::optional<
- std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr>>
-getDpasSubgroupLayouts(
- mlir::MLIRContext *context, VectorType aTy, VectorType bTy, VectorType cdTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, SmallVector<int64_t>>
- instDataVecs) {
- auto [instDataA, instDataB, instDataCD] = instDataVecs;
+/// Completes a scatter IO layout by deriving lane_layout and lane_data from
+/// inst_data when they are missing. If `consumerLayout` already has both
+/// lane_layout and lane_data, or has no inst_data, the layout is returned
+/// unchanged.
+///
+/// When lane info is absent, this function uses inst_data as the effective
+/// shape and computes the standard scatter-style lane factorization:
+/// - laneLayout[innermost] = min(subgroupSize, inst_data[innermost])
+/// - laneData[innermost] = min(inst_data[innermost] /
+/// laneLayout[innermost],
+/// maxChunkSize)
+///
+/// The returned layout carries inst_data + lane_layout + lane_data, ensuring
+/// the lane factorization is consistent with what the downstream load/store
+/// scatter anchor setup would produce.
+xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::uArch *uArch) {
+ if (!consumerLayout)
+ return consumerLayout;
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ if (instData.empty())
+ return consumerLayout;
+ if (!consumerLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !consumerLayout.getEffectiveLaneDataAsInt().empty())
+ return consumerLayout;
- std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
- if (consumerLayout && consumerLayout.isForWorkgroup()) {
- consumerSgLayout = consumerLayout.getEffectiveSgLayoutAsInt();
- }
+ // Reuse the load-side setup with inst_data as the destination shape.
+ const int subgroupSize = uArch->getSubgroupSize();
+ auto *context = consumerLayout.getContext();
+ auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ if (!uArchInstruction)
+ return consumerLayout;
+ int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
- // Get all valid layouts for A, B and C/D operands
- auto layoutsA = getSgLayoutCandidates(aTy.getShape(), instDataA, numSg);
- auto layoutsB = getSgLayoutCandidates(bTy.getShape(), instDataB, numSg);
- auto layoutsCD = getSgLayoutCandidates(cdTy.getShape(), instDataCD, numSg);
- if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
- return std::nullopt;
+ auto [defLaneLayout, defLaneData] =
+ computeScatterIOLaneLayoutAndData(instData, subgroupSize, maxChunkSize);
- // Pick the best subgroup layout
- std::optional<LayoutRepresentation> bestPick;
- auto checkAlignedSgDataAB = [&](const LayoutRepresentation &sgLayout) {
- return aTy.getShape().back() / sgLayout[1] ==
- bTy.getShape().front() / sgLayout[0];
- };
- for (auto &sgLayout : layoutsB) {
- if (llvm::is_contained(layoutsA, sgLayout) &&
- llvm::is_contained(layoutsCD, sgLayout)) {
- if (!checkAlignedSgDataAB(sgLayout))
- continue;
- // Is in (A and B and CD) and matches consumer -> best pick
- if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
- bestPick = sgLayout;
- break;
- }
- // Is in (A and B and CD) layoutsB is ordered from most
- // balanced to least. So the first one we see is the most balanced one,
- // remember it and later only update if there is one that matches the
- // consumer.
- if (!bestPick)
- bestPick = sgLayout;
- }
- }
- if (!bestPick)
- return std::nullopt;
+ return buildInstDataLayoutWithLane(context, instData, defLaneLayout,
+ defLaneData);
+}
- const auto &picked = *bestPick;
+/// Sets up layout for reduction operations by creating a SliceAttr for the
+/// result.
+///
+/// Algorithm Overview:
+/// This function attempts to construct a source layout that, when sliced along
+/// reduction dimensions, produces a result layout compatible with the
+/// consumer layout.
+///
+/// For subgroup layouts, it first tries to align the source layout's subgroup
+/// layout and data with the consumer's layout on non-reduction dimensions.
+/// Then, it distributes remaining subgroups across reduction dimensions. This
+/// avoids subgroup data redistribution overhead between the reduced result and
+/// its consumer. When the consumer layout is a slice layout, it attempts to
+/// reuse the slice layout's parent layout for the source to further minimize
+/// potential data redistribution.
+///
+/// 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.
+///
+/// Examples:
+/// 1. Subgroup layout - Row reduction on 2D tensor:
+/// srcShape=[32, 128], reductionDims=[1], resShape=[32], subgroupSize=16,
+/// NumSg=32
+/// * Consumer Layout:
+/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 8]>, dims =
+/// [1]>}
+//// * Result Layout:
+/// #xegpu.slice<#xegpu.layout<sg_layout=[4, 8],sg_data=[8, 16]>, dims =
+/// [1]>}
+/// Note that the sg_layout is reused but sg_data needs to be adjusted to
+/// evenly distribute the source tensor tile among the reduction dim.
+///
+/// 2. Subgroup layout - Same example above but consumer doesn't have a
+/// reusable slice layout.
+/// * Consumer Layout:
+/// #xegpu.layout<sgLayout=[32], sgData=[1]>
+/// * Result Layout:
+/// #xegpu.slice<#xegpu.layout<sgLayout=[32,1], sgData=[1, 128]>, dims =
+/// [1]>}
+/// * Consumer Layout:
+/// #xegpu.slice<#xegpu.layout<sgLayout=[8, 2, 4], sgData=[4, 64, 32]>,
+/// dims = [1, 2]>}
+/// * Result Layout:
+/// #xegpu.slice<#xegpu.layout<sgLayout=[8,4], sgData=[4, 32]>, dims =
+/// [1]>}
+/// Note that the consumer's layout can't be directly reused as is.
+/// So the algorithm distributes all subgroups on non reduction dimensions
+/// first and then distribute remaining subgroups on the reduction
+/// dimension.
+///
+/// 3. Lane layout - Default (lanes on innermost dim):
+/// srcShape=[32, 64], reductionDims=[0], subgroupSize=16
+/// Result: laneLayout=[1, 16], laneData=[1, 1]. The innermost dim is not
+/// reduced, so lanes stay on it.
+///
+/// 4. Lane layout - Switch (lanes moved off the reduction dim):
+/// srcShape=[32, 64], reductionDims=[1], subgroupSize=16
+/// Result: laneLayout=[16, 1], laneData=[1, 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.
- auto dpasALayout = buildSgLayout(context, aTy.getShape(), picked,
- /*dimK=*/aTy.getRank() - 1);
- auto dpasBLayout = buildSgLayout(context, bTy.getShape(), picked,
- /*dimK=*/bTy.getRank() - 2);
- auto dpasCDLayout = buildSgLayout(context, cdTy.getShape(), picked);
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
-}
+xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy,
+ DistributeLayoutAttr consumerLayout, SmallVector<int64_t> reductionDims,
+ int numSg, const xegpu::uArch::uArch *uArch) {
-/// Sets up the anchor layouts for dpas operands (A, B, and C/D).
-/// The numSg and consumerLayout (optional) are only used by sg layout
-/// creation.
-std::optional<
- std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr>>
-xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
- VectorType bTy, VectorType cdTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto context = aTy.getContext();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- if (!uArchInstruction)
- return std::nullopt;
- auto subgroupSize = uArch->getSubgroupSize();
+ auto srcShape = srcVecTy.getShape();
+ int srcRank = srcShape.size();
+ auto context = srcVecTy.getContext();
- auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
- aTy.getShape(), subgroupSize,
- aTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeA());
- auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
- bTy.getShape(), subgroupSize,
- bTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
- auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
- cdTy.getShape(), subgroupSize,
- cdTy.getElementType().getIntOrFloatBitWidth(),
- cdTy.getElementType().getIntOrFloatBitWidth());
+ const int subgroupSize = uArch->getSubgroupSize();
+ int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
+ xegpu::DistributeLayoutAttr srcLayout;
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ xegpu::SliceAttr consumerSliceLayout =
+ dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ if (consumerSliceLayout &&
+ consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
+ srcLayout = consumerSliceLayout.getParent();
+ SmallVector<int64_t> sgLayoutFromConsumer =
+ srcLayout.getEffectiveSgLayoutAsInt();
+ auto srcSgData = computeShapeRatio(srcShape, sgLayoutFromConsumer);
+ if (srcSgData)
+ for (int dim = 0; dim < srcRank; dim++) {
+ if (llvm::is_contained(reductionDims, dim))
+ srcLayout =
+ srcLayout.setDimData(dim, srcSgData.value()[dim], -1, -1);
+ }
+ } else {
+ SmallVector<int64_t> consumerSgLayout =
+ consumerLayout ? consumerLayout.getEffectiveSgLayoutAsInt()
+ : SmallVector<int64_t>();
+ SmallVector<int64_t> consumerSgData =
+ consumerLayout ? consumerLayout.getEffectiveSgDataAsInt()
+ : SmallVector<int64_t>();
+ SmallVector<int64_t> consumerOrder =
+ consumerLayout ? consumerLayout.getEffectiveOrderAsInt()
+ : SmallVector<int64_t>();
+ DenseI32ArrayAttr orderAttr =
+ consumerLayout ? consumerLayout.getOrder() : nullptr;
+ SmallVector<int64_t> sgLayout(srcRank), sgData(srcRank), order(srcRank);
+ int remainingSgCount =
+ consumerLayout ? consumerLayout.getNumSubgroups() : numSg;
+ int consumerIdx = 0;
- auto instDataVecs =
- getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction, subgroupSize);
- if (!instDataVecs)
- return std::nullopt;
+ // First pass: Match consumer's layout on non-reduction dimensions
+ for (int i = 0; i < srcRank; i++) {
+ if (!llvm::is_contained(reductionDims, i) &&
+ consumerIdx < static_cast<int>(consumerSgLayout.size())) {
+ sgLayout[i] = consumerSgLayout[consumerIdx];
+ sgData[i] = consumerSgData[consumerIdx];
+ remainingSgCount /= sgLayout[i];
+ order[i] = consumerOrder[consumerIdx];
+ consumerIdx++;
+ }
+ }
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(numSg > 0 &&
- "Number of subgroups must be provided for sg layout creation.");
- return getDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
- numSg, *instDataVecs);
+ // Second pass: Distribute remaining subgroups across reduction dimensions
+ // the reduction to scalar case is handled only by this loop
+ int64_t remainOrder = consumerSgLayout.size();
+ for (int i = 0; i < srcRank; i++) {
+ if (llvm::is_contained(reductionDims, i)) {
+ sgLayout[i] =
+ std::min(srcShape[i], static_cast<int64_t>(remainingSgCount));
+ assert((srcShape[i] % sgLayout[i] == 0) &&
+ "source shape not divisible by sg_layout");
+ sgData[i] = srcShape[i] / sgLayout[i];
+ remainingSgCount /= sgLayout[i];
+ order[i] = remainOrder++;
+ }
+ }
+ DenseI32ArrayAttr resOrderAttr = DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
+ if (!orderAttr || orderAttr.empty())
+ resOrderAttr = nullptr;
+ assert(remainingSgCount == 1 && "not all subgroups distributed");
+ srcLayout = buildLayout(context, sgLayout, sgData,
+ /*instData=*/{}, /*laneLayout=*/{},
+ /*laneData=*/{}, resOrderAttr);
+ }
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- return std::make_tuple(
- buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
- buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
- buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
- laneDataCD));
+ 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 veritical 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];
+ srcLayout =
+ buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- auto aLayout = buildLaneLayout(context, laneLayoutA, laneDataA);
- auto bLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
- auto cdLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
- return std::make_tuple(aLayout, bLayout, cdLayout);
+ // 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
+ // since the inst_data propagation already insert convert_layout if
+ // the layout is not consistent
+ srcLayout = consumerSliceLayout.getParent();
+ } else {
+ bool verticalLaneLayout = consumerReductionDims.empty() &&
+ reductionDims.size() == 1 &&
+ reductionDims[0] == (srcRank - 1);
+ auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
+ verticalLaneLayout);
+ srcLayout = buildLaneLayout(context, laneLayout, laneData);
+ }
}
- return std::nullopt;
-}
-
-/// Helper to create a scale layout derived from a matrix operand layout.
-/// The scale layout is computed by mapping each dimension of the matrix
-/// layout to the corresponding scale tensor dimension using the ratio
-/// between the matrix and scale shapes.
-static xegpu::DistributeLayoutAttr
-createScaleLayout(mlir::MLIRContext *context, VectorType matrixTy,
- VectorType scaleTy, xegpu::DistributeLayoutAttr matrixLayout,
- bool isBScale, const xegpu::uArch::uArch *uArch) {
- if (!scaleTy || !matrixLayout)
- return nullptr;
- // Calculate scaling factor by dividing matrix shape by scale shape
- ArrayRef<int64_t> matrixShape = matrixTy.getShape();
- ArrayRef<int64_t> scaleShape = scaleTy.getShape();
-
- // Scale shapes can be 1D or 2D, handle both cases
- if (scaleShape.empty())
- return nullptr;
-
- auto uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
-
- int64_t rank = matrixLayout.getRank();
- assert(rank >= 2 && "dpas layouts must be at least two dimensions");
+ return xegpu::SliceAttr::get(context, srcLayout,
+ DenseI64ArrayAttr::get(context, reductionDims));
+}
- SmallVector<int64_t> sgLayout = matrixLayout.getEffectiveSgLayoutAsInt();
- SmallVector<int64_t> sgData = matrixLayout.getEffectiveSgDataAsInt();
- SmallVector<int64_t> instData = matrixLayout.getEffectiveInstDataAsInt();
- SmallVector<int64_t> laneLayout = matrixLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> laneData = matrixLayout.getEffectiveLaneDataAsInt();
- auto order = matrixLayout.getOrder();
+/// Sets up layout for Reduction operations by creating a SliceAttr for the
+/// result.
+xegpu::SliceAttr
+xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy,
+ const xegpu::uArch::uArch *uArch) {
- SmallVector<int64_t> scaleSgLayout;
- SmallVector<int64_t> scaleSgData;
- if (!sgLayout.empty() && !sgData.empty()) {
- scaleSgLayout.assign(sgLayout.begin(), sgLayout.end());
- scaleSgData.assign(sgData.begin(), sgData.end());
- scaleSgData[rank - 2] = std::max<int64_t>(
- scaleShape[rank - 2] / (matrixShape[rank - 2] / sgData[rank - 2]), 1);
- scaleSgData[rank - 1] = std::max<int64_t>(
- scaleShape[rank - 1] / (matrixShape[rank - 1] / sgData[rank - 1]), 1);
- }
+ auto srcShape = srcVecTy.getShape();
+ auto context = srcVecTy.getContext();
+ auto subgroupSize = uArch->getSubgroupSize();
+ xegpu::LayoutAttr srcLayout;
- // For DPAS_MX scales: if matrix has inst_data, scale needs adjusted
- // inst_data. Scale inst_data is derived from matrix inst_data divided by
- // scale factor.
- SmallVector<int64_t> scaleInstData;
- if (!instData.empty()) {
- scaleInstData.assign(instData.begin(), instData.end());
- if (isBScale)
- scaleInstData[rank - 2] = std::max<int64_t>(
- scaleShape[rank - 2] / (matrixShape[rank - 2] / instData[rank - 2]),
- 1);
- else
- scaleInstData[rank - 1] = std::max<int64_t>(
- scaleShape[rank - 1] / (matrixShape[rank - 1] / instData[rank - 1]),
- 1);
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(true && "subgroup layout assignment not supported for reduction (op "
+ "is not expected at this level).");
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ assert(true && "instData layout assignment not supported for reduction (op "
+ "is not expected at this level).");
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ SmallVector<int64_t> laneLayout(1), laneData(1);
+ laneLayout[0] = std::min(static_cast<int64_t>(subgroupSize), srcShape[0]);
+ laneData[0] = 1;
+ srcLayout = buildLaneLayout(context, laneLayout, laneData);
}
- SmallVector<int64_t> scaleLaneLayout;
- SmallVector<int64_t> scaleLaneData;
- if (!laneLayout.empty() && !laneData.empty()) {
- scaleLaneLayout.assign(laneLayout.begin(), laneLayout.end());
- scaleLaneData.assign(laneData.size(), 1);
+ auto result = xegpu::SliceAttr::get(context, srcLayout,
+ DenseI64ArrayAttr::get(context, 0));
+ return result;
+}
- bool isRowMajor = uArchInstruction->isLaneLayoutRowMajorOrder();
- if (isBScale ^ isRowMajor)
- std::swap(scaleLaneLayout[rank - 2], scaleLaneLayout[rank - 1]);
- // Cap lane_layout by the per-instruction tile (inst_data) on each dim.
- // Then derive lane_data = inst_data / lane_layout so the Category A
- // invariant inst_data = lane_layout * lane_data * k (with k = 1) holds
- // for the scale operand's load_nd consumer.
- auto layoutCap = scaleInstData.empty() ? scaleShape : scaleInstData;
- for (int64_t d = rank - 2; d < rank; ++d)
- scaleLaneLayout[d] = std::min<int64_t>(layoutCap[d], scaleLaneLayout[d]);
+/// Adjusts `consumerLayout`'s innermost-dim data field selected by
+/// `layoutKind` so that the source layout can be safely inferred by dividing
+/// that value by `ratio`. Doubles the value until the divisibility constraint
+/// is met, bounded above by `bound` like result-shape.
+///
+/// Used by ops whose source relates to the result by a fixed factor along the
+/// innermost dim (e.g., bitcast: bitwidth ratio; interleave: 2x).
+///
+/// Divisibility constraints per LayoutKind:
+/// - Subgroup: sgData[innermost] % ratio == 0
+/// - InstData: instData[innermost] % (laneLayout[innermost] * ratio) == 0
+/// (laneLayout falls back to subgroupSize if absent)
+/// - Lane: laneData[innermost] % ratio == 0
+static xegpu::DistributeLayoutAttr
+adjustInnermostDimForDivisibility(xegpu::DistributeLayoutAttr consumerLayout,
+ xegpu::LayoutKind layoutKind,
+ size_t innerMostDim, int ratio, int64_t bound,
+ const xegpu::uArch::uArch *uArch) {
+ SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+
+ int64_t sgDataValue = -1;
+ int64_t instDataValue = -1;
+ int64_t laneDataValue = -1;
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ sgDataValue = sgData[innerMostDim];
+ while ((sgDataValue <= bound) && (sgDataValue % ratio) != 0)
+ sgDataValue *= 2;
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
+ while ((instDataValue <= bound) &&
+ (instDataValue % (innermostDimLaneLayout * ratio) != 0))
+ instDataValue *= 2;
+ assert((bound % instDataValue) == 0 &&
+ "bound, instData, and laneLayout for innermost must be 2^n!");
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ laneDataValue = laneData[innerMostDim];
+ while ((laneDataValue <= bound) && (laneDataValue % ratio) != 0)
+ laneDataValue *= 2;
}
- return buildLayout(context, scaleSgLayout, scaleSgData, scaleInstData,
- scaleLaneLayout, scaleLaneData, order);
+
+ return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
+ laneDataValue);
}
-/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
-/// B_scale). The numSg and consumerLayout (optional) are only used by sg
-/// layout creation.
-std::optional<
- std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr>>
-xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
- VectorType bTy, VectorType cdTy, VectorType aScaleTy,
- VectorType bScaleTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto context = aTy.getContext();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- if (!uArchInstruction)
- return std::nullopt;
- auto subgroupSize = uArch->getSubgroupSize();
+/// Sets up the result layout for a bitcast operation.
+/// When casting to a smaller bitwidth, adjusts the layout dimensions (sgData,
+/// instData, or laneData) by multiplying by the bitwidth ratio to ensure the
+/// result layout can be correctly divided back to the source layout during
+/// inference.
+///
+/// Examples:
+/// 1. Casting f32 -> f16 (32-bit to 16-bit, bitWidthRatio = 2):
+/// Consumer layout: instData=[1, 16], subgroupSize=16
+/// Source shape: [8, 32]
+/// Result layout: instData=[1, 32] (16 * 2)
+/// The innermost dimension is multiplied by 2 to maintain consistency.
+///
+/// 2. Casting f32 -> i8 (32-bit to 8-bit, bitWidthRatio = 4):
+/// Consumer instData=[1, 16], subgroupSize=16
+/// Source shape: [4, 128]
+/// adjust the instData from [1, 16] to [1, 16 * 4 = 64]
+///
+/// 3. Casting i8 -> i32 (8-bit to 32-bit, bitWidthRatio = 1/4):
+/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4]
+/// No adjustment needed - returns consumer layout directly.
+///
+xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
- auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
- aTy.getShape(), subgroupSize,
- aTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeA());
- auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
- bTy.getShape(), subgroupSize,
- bTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
- auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
- cdTy.getShape(), subgroupSize,
- cdTy.getElementType().getIntOrFloatBitWidth(),
- cdTy.getElementType().getIntOrFloatBitWidth());
- auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction,
- subgroupSize, /*isDpasMx=*/true);
- if (!instDataVecs)
- return std::nullopt;
+ int srcElemTyBitWidth = srcVecTy.getElementType().getIntOrFloatBitWidth();
+ int resElemTyBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
- if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(numSg > 0 &&
- "Number of subgroups must be provided for sg layout creation.");
- auto dpasLayouts = getDpasSubgroupLayouts(
- context, aTy, bTy, cdTy, consumerLayout, numSg, *instDataVecs);
- if (!dpasLayouts)
- return std::nullopt;
+ ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
- auto [dpasALayout, dpasBLayout, dpasCDLayout] = *dpasLayouts;
+ assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
+ "laneData must be available for all dimensions");
- // Create scale layouts
- auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
+ // Casting to same/larger element type: result has fewer (or equal) elements
+ // along the innermost dim, no adjustment needed.
+ if (srcElemTyBitWidth <= resElemTyBitWidth)
+ return consumerLayout;
- auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+ // Casting to smaller element type: result has more elements along innermost
+ // dim. Adjust the innermost data field upward so the source layout can be
+ // recovered by dividing by bitWidthRatio.
+ size_t innerMostDim = srcShape.size() - 1;
+ int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
+ return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
+ innerMostDim, bitWidthRatio,
+ resShape[innerMostDim], uArch);
+}
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
- bScaleLayout);
- } else if (layoutKind == xegpu::LayoutKind::InstData) {
+/// Sets up the result layout for an interleave operation to ensure the source
+/// layout can be safely derived. Interleave doubles the innermost dimension,
+/// so the result layout must ensure that laneData is a multiple
+/// of 2, and instData must be divisible by innermostDimLaneLayout * 2.
+///
+/// Example:
+/// Interleave: vector<128x256xf4> -> vector<128x512xf4>
+/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4], instData=[1, 64]
+/// Result layout adjustment to ensure source can be safely inferred:
+/// - laneData must be >= 2 and multiple of 2 (so source = laneData/2 is
+/// valid)
+/// - instData must be divisible by (16 * 2 = 32) (so source = instData/2 is
+/// valid)
+/// - Adjusted instData: ensure (instData % 32 == 0)
+///
+xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
+ assert(consumerLayout.getRank() == static_cast<int64_t>(resShape.size()) &&
+ "consumer layout rank must match source shape rank");
- auto dpasALayout =
- buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA);
- auto dpasBLayout =
- buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB);
- auto dpasCDLayout = buildInstDataLayoutWithLane(context, instDataCD,
- laneLayoutCD, laneDataCD);
+ // Interleave doubles the innermost dimension (ratio = 2). Adjust the
+ // innermost data field so the source layout can be recovered by dividing
+ // by 2.
+ const size_t innerMostDim = resShape.size() - 1;
+ constexpr int ratio = 2;
+ return adjustInnermostDimForDivisibility(consumerLayout, layoutKind,
+ innerMostDim, ratio,
+ resShape[innerMostDim], uArch);
+}
- auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
- auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+/// Sets up the result layout for an insert strided slice operation.
+/// Creates a result layout based on the specified layout kind (InstData or
+/// Lane).
+xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVectorTy,
+ VectorType resVectorTy, xegpu::DistributeLayoutAttr consumerLayout,
+ const xegpu::uArch::uArch *uArch) {
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
- bScaleLayout);
- } else if (layoutKind == xegpu::LayoutKind::Lane) {
- auto dpasALayout = buildLaneLayout(context, laneLayoutA, laneDataA);
- auto dpasBLayout = buildLaneLayout(context, laneLayoutB, laneDataB);
- auto dpasCDLayout = buildLaneLayout(context, laneLayoutCD, laneDataCD);
+ xegpu::DistributeLayoutAttr requiredResLayout;
+ SmallVector<int64_t> consumerInstData =
+ consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ ArrayRef<int64_t> srcShape = srcVectorTy.getShape();
+ int64_t laneDataValue = -1;
- auto aScaleLayout =
- createScaleLayout(context, aTy, aScaleTy, dpasALayout, false, uArch);
- auto bScaleLayout =
- createScaleLayout(context, bTy, bScaleTy, dpasBLayout, true, uArch);
+ requiredResLayout = consumerLayout;
+ int srcRank = srcShape.size();
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
- bScaleLayout);
+ if (layoutKind == xegpu::LayoutKind::Subgroup ||
+ layoutKind == xegpu::LayoutKind::InstData) {
+ assert(true &&
+ "subgroup layout assignment not supported for insertStridedSlice.");
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ for (int dim = 0; dim < srcRank; dim++) {
+ assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
+ "srcShape must be divisible by laneLayout for all dimensions");
+ laneDataValue = std::min(srcShape[dim] / consumerLaneLayout[dim],
+ consumerLaneData[dim]);
+ requiredResLayout =
+ requiredResLayout.setDimData(dim, -1, -1, laneDataValue);
+ }
}
- return std::nullopt;
+ return requiredResLayout;
}
+/// Back-propagates a known result layout to the layout required on `operand`
+/// for a non-anchor (layout-propagating) vector op. Dispatches on the op kind —
+/// broadcast, (multi)reduction, bitcast, shape/transpose, insert/extract,
+/// interleave, etc. — applying the shape/permutation/bitwidth transform to
+/// derive the source layout; elementwise and pass-through ops reuse resLayout
+/// as-is. Returns nullptr for unknown ops or an absent result layout.
xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
OpOperand &operand, xegpu::DistributeLayoutAttr resLayout) {
if (!resLayout)
@@ -2338,6 +2401,9 @@ xegpu::DistributeLayoutAttr xegpu::inferSourceLayoutFromResultForNonAnchorOp(
return nullptr;
}
+/// Returns the layout required on `operand`: anchor ops report their declared
+/// per-operand layout directly; non-anchor ops back-derive it from their result
+/// layout via inferSourceLayoutFromResultForNonAnchorOp.
xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
Operation *op = operand.getOwner();
// Anchor ops declare the layout they
>From 5a069045a78c9ecd7f2d8df005b04c9117b0b638 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 16 Jun 2026 19:07:56 +0000
Subject: [PATCH 33/42] polish setup rules
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 136 ++++++++----------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 53 -------
2 files changed, 61 insertions(+), 128 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index bb8e96fec5f6f..ea5477d7a46ad 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -928,8 +928,7 @@ getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
/// C/D.
static std::optional<SmallVector<int64_t>>
get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
- ArrayRef<int> bHeights, ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData) {
+ ArrayRef<int> bHeights) {
int rank = dataShape.size();
// Compute inst_data from hardware block params. For Nd ops, the lane
// factorization above (laneLayout / laneData) is rigid; inst_data must be
@@ -944,16 +943,54 @@ get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
instData.back() = instWidth;
instData[rank - 2] = instHeight;
- if (instWidth == -1 || instHeight == -1) {
- instData.back() = laneLayout.back() * laneData.back();
- instData[rank - 2] = laneLayout[rank - 2] * laneData[rank - 2];
- }
- for (int dim = 0; dim < rank; ++dim)
- assert(instData[dim] % (laneLayout[dim] * laneData[dim]) == 0 &&
- "inst_data must be a multiple of lane_layout * lane_data for ND op");
return instData;
}
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D. Look up the uArch table and search for the largest supported block size
+/// that divides the data shape
+static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>>
+getDpasInstDataLayouts(
+ VectorType aTy, VectorType bTy, VectorType cdTy,
+ const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
+ const int subgroupSize, bool isDpasMx = false) {
+
+ // M dimension is the second-to-last dim of A (handles batch dims).
+ const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
+ auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
+ const int maxALen =
+ xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
+
+ // N dimension is the last dim of B.
+ const unsigned dataBLen = bTy.getShape().back();
+ auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
+ const int maxBLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
+
+ auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
+ const int maxCLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
+ if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+ return std::nullopt;
+
+ auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
+ if (supportedKLen.empty())
+ return std::nullopt;
+ auto kDimSize = supportedKLen[0];
+
+ SmallVector<int64_t> instDataA(aTy.getRank(), 1);
+ instDataA[aTy.getRank() - 2] = maxALen;
+ instDataA[aTy.getRank() - 1] = kDimSize;
+ SmallVector<int64_t> instDataB(bTy.getRank(), 1);
+ instDataB[bTy.getRank() - 2] = kDimSize;
+ instDataB[bTy.getRank() - 1] = maxBLen;
+ SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
+ instDataCD[cdTy.getRank() - 2] = maxALen;
+ instDataCD[cdTy.getRank() - 1] = maxCLen;
+ return std::make_tuple(instDataA, instDataB, instDataCD);
+}
+
/// Computes lane_layout and lane_data for scatter-style store anchor layouts
/// (store scatter, store matrix). Lanes and the per-lane vector both live on
/// the innermost dim:
@@ -1097,57 +1134,10 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
// (Lane kind only; sg/inst layouts unsupported).
//===----------------------------------------------------------------------===//
-/// Helper function to compute inst_data vectors for DPAS operands A, B, and
-/// C/D.
-static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
- SmallVector<int64_t>>>
-getDpasInstDataLayouts(
- VectorType aTy, VectorType bTy, VectorType cdTy,
- const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
- const int subgroupSize, bool isDpasMx = false) {
-
- // M dimension is the second-to-last dim of A (handles batch dims).
- const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
- auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
- const int maxALen =
- xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
-
- // N dimension is the last dim of B.
- const unsigned dataBLen = bTy.getShape().back();
- auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
- const int maxBLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
-
- auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
- const int maxCLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
- if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
- return std::nullopt;
-
- // For DPAS_MX, use getSupportedK to get the scaled K dimension.
- // assume single element in the returned vector.
- int kDimSize = subgroupSize;
- if (isDpasMx) {
- auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
- if (supportedKLen.empty())
- return std::nullopt;
- kDimSize = supportedKLen[0];
- }
-
- SmallVector<int64_t> instDataA(aTy.getRank(), 1);
- instDataA[aTy.getRank() - 2] = maxALen;
- instDataA[aTy.getRank() - 1] = kDimSize;
- SmallVector<int64_t> instDataB(bTy.getRank(), 1);
- instDataB[bTy.getRank() - 2] = kDimSize;
- instDataB[bTy.getRank() - 1] = maxBLen;
- SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
- instDataCD[cdTy.getRank() - 2] = maxALen;
- instDataCD[cdTy.getRank() - 1] = maxCLen;
- return std::make_tuple(instDataA, instDataB, instDataCD);
-}
-
/// Helper function to set up subgroup layouts for DPAS operands A, B, and
-/// C/D. Returns the three layouts if successful, nullopt otherwise.
+/// C/D. Compute subgroup layout candidates based on wgtile and instData, and
+/// then pick the best one that satisfies all operands and the consumer (if
+/// specified).
static std::optional<
std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
xegpu::DistributeLayoutAttr>>
@@ -1464,15 +1454,13 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
// Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
- auto [laneLayout, laneData] =
- compute2DBlockIOLaneLayoutAndData(dataShape, uArch->getSubgroupSize(),
- bitwidth, packingSize, /*vnni=*/false);
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, uArch->getSubgroupSize(), bitwidth, packingSize);
if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
- auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights,
- laneLayout, laneData);
+ auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights);
if (layoutKind == xegpu::LayoutKind::InstData)
return buildInstDataLayoutWithLane(context, *instData, laneLayout,
@@ -1591,14 +1579,14 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
hasTransform ? consumerLaneData[rank - 2] : consumerLaneData[rank - 1];
unsigned packingSize = packingFactor * elemTy.getIntOrFloatBitWidth();
- auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
- elemTy, hasTransform, hasTranspose,
- /*upConv=*/false);
- if (!blockWHC)
- return nullptr;
- auto [bWidths, bHeights, bCounts] = blockWHC.value();
-
if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
+ elemTy, hasTransform, hasTranspose,
+ /*upConv=*/false);
+ if (!blockWHC)
+ return nullptr;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
+
int64_t height = consumerInstData[rank - 2];
int64_t width = consumerInstData[rank - 1];
auto maxBlockCount = *llvm::max_element(bCounts);
@@ -1611,12 +1599,10 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout.getOrder());
}
}
-
+ auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights);
auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
+ *instData, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
hasTransform, hasTranspose);
- auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights,
- laneLayout, laneData);
return buildInstDataLayoutWithLane(context, *instData, laneLayout,
laneData);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 65cb48f19af68..2ed59a52287ff 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -259,59 +259,6 @@ struct LayoutInfoLattice : public Lattice<LayoutInfo> {
using Lattice::Lattice;
};
-/// Helper Functions to get default layouts. A `default layout` is a layout that
-/// is assigned to a value when the layout is not fixed by some anchor operation
-/// (like DPAS).
-
-// /// Helper Function to get the default layout for uniform values like
-// constants.
-// /// For 1D vector, lane_layout is [subgroupSize] and lane_data is [1].
-// /// For 2D vector, lane_layout is [1, subgroupSize] and lane_data is [1, 1].
-// /// For ND vector (N>2), leading dims get unit lane_layout and lane_data.
-// static LayoutInfo getDefaultSIMTLayoutInfo(mlir::MLIRContext *ctx,
-// unsigned rank,
-// const xegpu::uArch::uArch *uArch)
-// {
-// assert(rank >= 1 && "Expected at least 1D vector.");
-// if (rank == 1) {
-// return LayoutInfo(
-// xegpu::LayoutAttr::get(ctx, {uArch->getSubgroupSize()}, {1}));
-// }
-// // For rank >= 2, lane_layout is [1, ..., 1, subgroupSize] and
-// // lane_data is [1, ..., 1, 1].
-// SmallVector<int32_t> laneLayout(rank, 1);
-// SmallVector<int32_t> laneData(rank, 1);
-// laneLayout[rank - 1] = uArch->getSubgroupSize();
-// return LayoutInfo(xegpu::LayoutAttr::get(ctx, laneLayout, laneData));
-// }
-
-// /// Helper to get the default layout for 2D block operations.
-// /// For ND (N>2) types, leading dimensions get unit layout/data values.
-// template <typename Ty>
-// static LayoutInfo getSIMTLayoutInfoBlockIO(Ty ty,
-// const xegpu::uArch::uArch *uArch,
-// unsigned packingSize) {
-// // Expecting at least 1D.
-// assert(ty.getRank() >= 1 && "Expected at least 1D vector.");
-// // Expecting int or float element type.
-// assert(ty.getElementType().isIntOrFloat() &&
-// "Expected int or float element type.");
-// // If the rank is 1, then return default layout for 1D vector.
-// if (ty.getRank() == 1)
-// return getDefaultSIMTLayoutInfo(ty.getContext(), 1, uArch);
-// // Packing factor is determined by the element type bitwidth.
-// unsigned bitwidth = ty.getElementType().getIntOrFloatBitWidth();
-// int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;
-// // For rank >= 2, distribute along the last dimension with leading units.
-// unsigned rank = ty.getRank();
-// SmallVector<int32_t> laneLayout(rank, 1);
-// SmallVector<int32_t> laneData(rank, 1);
-// laneLayout[rank - 1] = uArch->getSubgroupSize();
-// laneData[rank - 1] = packingFactor;
-// return LayoutInfo(
-// xegpu::LayoutAttr::get(ty.getContext(), laneLayout, laneData));
-// }
-
//===----------------------------------------------------------------------===//
// LayoutInfoPropagation
//===----------------------------------------------------------------------===//
>From dedd802df54f5e6d4e70ab7a6e1405ec027d9398 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 16 Jun 2026 22:08:38 +0000
Subject: [PATCH 34/42] polish completeScatterIOLaneLayoutFromInstData; add
BlockIOInstructionInterface to uArch
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 7 +-
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 30 ++++--
.../mlir/Dialect/XeGPU/uArch/uArchBase.h | 19 ++++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 93 +++++++++++--------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 7 +-
5 files changed, 99 insertions(+), 57 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 7b05bfece50cc..a50cd91274e2b 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -234,9 +234,10 @@ DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
/// rule with inst_data as the destination shape. The resulting lane info is
/// merged with the consumer's inst_data so downstream setup* paths see a
/// fully-populated layout.
-DistributeLayoutAttr
-completeScatterIOLaneLayoutFromInstData(DistributeLayoutAttr consumerLayout,
- Type elemTy, const uArch::uArch *uArch);
+DistributeLayoutAttr completeScatterIOLaneLayoutFromInstData(
+ DistributeLayoutAttr userSpecifiedLayout,
+ DistributeLayoutAttr consumerLayout, Type elemTy,
+ const uArch::uArch *uArch);
/// Sets up the anchor layout for a store_nd operation. StoreNd does not
/// consider a consumer layout (it is a data sink), and picks its layout from
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index c0cd8be341adf..547708b65abe1 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -45,7 +45,8 @@ struct Xe2Plus : public uArch {
//===----------------------------------------------------------------------===//
// uArch instructions
//===----------------------------------------------------------------------===//
-struct Subgroup2DBlockStoreInstruction : public Instruction {
+struct Subgroup2DBlockStoreInstruction : public Instruction,
+ public BlockIOInstructionInterface {
Subgroup2DBlockStoreInstruction()
: Instruction(InstructionKind::Subgroup2DBlockStore,
InstructionScope::Subgroup) {}
@@ -54,9 +55,12 @@ struct Subgroup2DBlockStoreInstruction : public Instruction {
}
// Source :
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_2d_block_io.html#_add_a_new_section_5_2_x_cl_intel_subgroup_2d_block_io
+ // Stores ignore the transform / transpose / upConv flags.
std::optional<
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
- getBlockWidthHeightCount(Type elemTy) const {
+ getBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/ = false,
+ bool /*hasTranspose*/ = false,
+ bool /*upConv*/ = false) const override {
const static int kHeight[] = {1, 2, 4, 8};
const static int kWidth16[] = {16};
const static int kWidth32[] = {16};
@@ -73,10 +77,11 @@ struct Subgroup2DBlockStoreInstruction : public Instruction {
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
-struct Subgroup2DBlockLoadInstruction : public Instruction {
+struct Subgroup2DBlockLoadInstruction : public Instruction,
+ public BlockIOInstructionInterface {
Subgroup2DBlockLoadInstruction()
: Instruction(InstructionKind::Subgroup2DBlockLoad,
InstructionScope::Subgroup) {}
@@ -88,8 +93,9 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_2d_block_io.html#_add_a_new_section_5_2_x_cl_intel_subgroup_2d_block_io
std::optional<
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
- getBlockWidthHeightCount(Type elemTy, bool hasTransform, bool hasTranspose,
- bool upConv = false) const {
+ getBlockWidthHeightCount(Type elemTy, bool hasTransform = false,
+ bool hasTranspose = false,
+ bool upConv = false) const override {
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kHeightAtLeast8[] = {8, 16, 32};
static const int kHeightAtLeast16[] = {16, 32};
@@ -151,10 +157,11 @@ struct Subgroup2DBlockLoadInstruction : public Instruction {
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
-struct Subgroup2DBlockPrefetchInstruction : public Instruction {
+struct Subgroup2DBlockPrefetchInstruction : public Instruction,
+ public BlockIOInstructionInterface {
Subgroup2DBlockPrefetchInstruction()
: Instruction(InstructionKind::Subgroup2DBlockPrefetch,
InstructionScope::Subgroup) {}
@@ -163,9 +170,12 @@ struct Subgroup2DBlockPrefetchInstruction : public Instruction {
}
// Source :
// https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_buffer_prefetch.html#_add_a_new_section_6_15_x_sub_group_prefetch_functions
+ // Prefetches ignore the transform / transpose / upConv flags.
std::optional<
std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
- getBlockWidthHeightCount(Type elemTy) const {
+ getBlockWidthHeightCount(Type elemTy, bool /*hasTransform*/ = false,
+ bool /*hasTranspose*/ = false,
+ bool /*upConv*/ = false) const override {
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kWidth32[] = {32};
@@ -189,7 +199,7 @@ struct Subgroup2DBlockPrefetchInstruction : public Instruction {
return it->second;
return std::nullopt;
}
- int32_t getPackedFormatBitSize() const { return 16; }
+ int32_t getPackedFormatBitSize() const override { return 16; }
};
struct SubgroupMatrixMultiplyAcc : public Instruction,
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 147a56a52c188..61db4605e85fa 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -19,6 +19,7 @@
#include <iostream>
#include <map>
#include <mutex>
+#include <optional>
#include <shared_mutex>
#include <tuple>
@@ -255,6 +256,24 @@ struct MMAInstructionInterface {
virtual ~MMAInstructionInterface() = default;
};
+// Interface for subgroup-level 2D block instructions (load / store / prefetch).
+// All three describe the set of hardware-supported block shapes via
+// (width, height, count) tuples and share a packed-format bit size. The
+// transform / transpose / upConv flags are only meaningful for loads; store
+// and prefetch implementations ignore them.
+struct BlockIOInstructionInterface {
+ // Returns the supported (widths, heights, counts) for the given element
+ // type, or std::nullopt if the element type is unsupported.
+ virtual std::optional<
+ std::tuple<llvm::ArrayRef<int>, llvm::ArrayRef<int>, llvm::ArrayRef<int>>>
+ getBlockWidthHeightCount(Type elemTy, bool hasTransform = false,
+ bool hasTranspose = false,
+ bool upConv = false) const = 0;
+ // Bit size of the packed format used by this block instruction.
+ virtual int32_t getPackedFormatBitSize() const = 0;
+ virtual ~BlockIOInstructionInterface() = default;
+};
+
//===----------------------------------------------------------------------===//
// Common instructions (shared across architectures)
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ea5477d7a46ad..12347bc31d879 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -394,6 +394,17 @@ static xegpu::LayoutAttr buildInstDataLayoutWithLane(
orderAttr);
}
+bool isValidLaneLayout(ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ ArrayRef<int64_t> dataShape) {
+ int rank = dataShape.size();
+ for (int dim = 0; dim < rank; ++dim) {
+ int64_t laneProduct = laneLayout[dim] * laneData[dim];
+ if (dataShape[dim] % laneProduct != 0)
+ return false;
+ }
+ return true;
+}
+
static xegpu::LayoutAttr
buildLaneLayout(mlir::MLIRContext *context, ArrayRef<int64_t> laneLayout,
ArrayRef<int64_t> laneData,
@@ -1608,13 +1619,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- bool validLaneLayout = true;
- for (int dim = 0; dim < rank; ++dim) {
- int64_t laneProduct = consumerLaneLayout[dim] * consumerLaneData[dim];
- if (dataShape[dim] % laneProduct != 0)
- validLaneLayout = false;
- }
- if (validLaneLayout) {
+ if (isValidLaneLayout(consumerLaneLayout, consumerLaneData, dataShape)) {
return consumerLayout;
} else {
auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
@@ -1659,13 +1664,10 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
// vector capped by maxChunkSize).
SmallVector<int64_t> laneLayout;
SmallVector<int64_t> laneData;
- if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
- laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
- laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
- } else {
- auto [LaneLayout, LaneData] =
- computeScatterIOLaneLayoutAndData(resShape, subgroupSize, maxChunkSize);
- }
+ 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());
if (layoutKind == xegpu::LayoutKind::InstData) {
// Take consumer's inst_data as-is. If the consumer doesn't have one,
@@ -1802,47 +1804,56 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
}
/// Completes a scatter IO layout by deriving lane_layout and lane_data from
-/// inst_data when they are missing. If `consumerLayout` already has both
-/// lane_layout and lane_data, or has no inst_data, the layout is returned
-/// unchanged.
+/// `specifiedLayout`'s inst_data when they are missing. The layout is returned
+/// unchanged if `specifiedLayout` is null, carries no inst_data, or already has
+/// both lane_layout and lane_data.
///
-/// When lane info is absent, this function uses inst_data as the effective
-/// shape and computes the standard scatter-style lane factorization:
-/// - laneLayout[innermost] = min(subgroupSize, inst_data[innermost])
-/// - laneData[innermost] = min(inst_data[innermost] /
-/// laneLayout[innermost],
-/// maxChunkSize)
+/// When lane info is absent, inst_data is treated as the effective shape and
+/// the lane factorization is filled in as follows:
+/// - If `consumerLayout` is present and its lane_layout / lane_data are a
+/// valid factorization of inst_data, that consumer lane info is reused so
+/// the completed layout matches the consumer (avoiding a relayout).
+/// - Otherwise a standard scatter-style factorization is computed via
+/// `computeScatterIOLaneLayoutAndData`, bounded by `maxChunkSize` — the
+/// per-lane load width reported by the uArch's LoadGather instruction
+/// (`getMaxLaneLoadSize`).
///
-/// The returned layout carries inst_data + lane_layout + lane_data, ensuring
-/// the lane factorization is consistent with what the downstream load/store
-/// scatter anchor setup would produce.
xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout,
xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
const xegpu::uArch::uArch *uArch) {
- if (!consumerLayout)
- return consumerLayout;
- SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- if (instData.empty())
- return consumerLayout;
- if (!consumerLayout.getEffectiveLaneLayoutAsInt().empty() &&
- !consumerLayout.getEffectiveLaneDataAsInt().empty())
- return consumerLayout;
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
// Reuse the load-side setup with inst_data as the destination shape.
const int subgroupSize = uArch->getSubgroupSize();
- auto *context = consumerLayout.getContext();
+ auto *context = specifiedLayout.getContext();
auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
if (!uArchInstruction)
- return consumerLayout;
+ return specifiedLayout;
int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
-
- auto [defLaneLayout, defLaneData] =
- computeScatterIOLaneLayoutAndData(instData, subgroupSize, maxChunkSize);
-
- return buildInstDataLayoutWithLane(context, instData, defLaneLayout,
+ if (consumerLayout) {
+ auto consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ auto consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+ if (!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
+ isValidLaneLayout(consumerLaneLayout, consumerLaneData,
+ specifiedInstData))
+ return buildInstDataLayoutWithLane(context, specifiedInstData,
+ consumerLaneLayout, consumerLaneData);
+ }
+ auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, maxChunkSize);
+ return buildInstDataLayoutWithLane(context, specifiedInstData, defLaneLayout,
defLaneData);
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 2ed59a52287ff..039a94b307da4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1156,7 +1156,8 @@ void LayoutInfoPropagation::visitLoadGatherOp(
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
- anchorLayoutAttr, resVecTy.getElementType(), uArch);
+ anchorLayoutAttr, consumerLayoutAttr, resVecTy.getElementType(),
+ uArch);
load.setLayoutAttr(requiredAnchorLayoutAttr);
}
} else {
@@ -1201,7 +1202,7 @@ void LayoutInfoPropagation::visitStoreScatterOp(
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
- anchorLayoutAttr, srcVecTy.getElementType(), uArch);
+ anchorLayoutAttr, nullptr, srcVecTy.getElementType(), uArch);
storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
}
} else {
@@ -1273,7 +1274,7 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
- anchorLayoutAttr, srcVecTy.getElementType(), uArch);
+ anchorLayoutAttr, nullptr, srcVecTy.getElementType(), uArch);
storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
}
} else {
>From 8e9e69aacf2dfd506aeee6b90cd84f27fac30762 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 16 Jun 2026 22:38:17 +0000
Subject: [PATCH 35/42] refactor buildInstDataLayoutWithLane
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 156 +++++++++---------
1 file changed, 80 insertions(+), 76 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 12347bc31d879..d8f73615fbd16 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -394,8 +394,9 @@ static xegpu::LayoutAttr buildInstDataLayoutWithLane(
orderAttr);
}
-bool isValidLaneLayout(ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
- ArrayRef<int64_t> dataShape) {
+bool isValidLaneLayout(ArrayRef<int64_t> dataShape,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
int rank = dataShape.size();
for (int dim = 0; dim < rank; ++dim) {
int64_t laneProduct = laneLayout[dim] * laneData[dim];
@@ -937,10 +938,14 @@ getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
/// Helper function to compute inst_data vectors for DPAS operands A, B, and
/// C/D.
-static std::optional<SmallVector<int64_t>>
-get2DBlockIOInstDataLayout(ArrayRef<int64_t> dataShape, ArrayRef<int> bWidths,
- ArrayRef<int> bHeights) {
+static std::optional<SmallVector<int64_t>> get2DBlockIOInstDataLayout(
+ ArrayRef<int64_t> dataShape, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction) {
int rank = dataShape.size();
+ auto blockWHC = uArchInstruction->getBlockWidthHeightCount(elemTy);
+ if (!blockWHC)
+ return std::nullopt;
+ auto [bWidths, bHeights, bCounts] = blockWHC.value();
// Compute inst_data from hardware block params. For Nd ops, the lane
// factorization above (laneLayout / laneData) is rigid; inst_data must be
// a multiple of lane_layout * lane_data on each dim (Category A
@@ -965,7 +970,7 @@ static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
getDpasInstDataLayouts(
VectorType aTy, VectorType bTy, VectorType cdTy,
const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
- const int subgroupSize, bool isDpasMx = false) {
+ bool isDpasMx = false) {
// M dimension is the second-to-last dim of A (handles batch dims).
const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
@@ -1239,8 +1244,7 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
cdTy.getElementType().getIntOrFloatBitWidth(),
cdTy.getElementType().getIntOrFloatBitWidth());
- auto instDataVecs =
- getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction, subgroupSize);
+ auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction);
if (!instDataVecs)
return std::nullopt;
@@ -1380,7 +1384,7 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
cdTy.getElementType().getIntOrFloatBitWidth(),
cdTy.getElementType().getIntOrFloatBitWidth());
auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction,
- subgroupSize, /*isDpasMx=*/true);
+ /*isDpasMx=*/true);
if (!instDataVecs)
return std::nullopt;
@@ -1437,45 +1441,45 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
}
-/// Generic anchor-layout setup for ND ops (load_nd, store_nd, prefetch_nd).
-///
-/// Given hardware-supported block widths/heights, picks the largest divisor
-/// of the trailing two dims of `dataShape` as the default `inst_data`. The
-/// lane layout is the standard 2D-block-IO default (subgroupSize lanes on the
-/// innermost dim, optional packing on the innermost dim).
-///
-/// `consumerLayout` (optional) is honored when its parameters are valid w.r.t.
-/// the uArch constraints; otherwise the helper falls back to defaults.
-///
-/// For Lane kind: returns just the lane layout / lane data (consumer ignored;
-/// lane layout is fully determined by hardware).
-/// For InstData kind: returns inst_data + lane_layout/lane_data (Category A:
-/// inst_data = k * lane_layout * lane_data, k >= 1). Honors consumer's
-/// inst_data when it is uArch-valid.
-/// For Subgroup kind: if the consumer specifies a workgroup-level layout,
-/// reuses it directly; otherwise picks the most balanced sg_layout via
-/// `getSgLayoutCandidates` (requires `numSg`).
-static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
- xegpu::LayoutKind layoutKind, mlir::MLIRContext *context,
- ArrayRef<int64_t> dataShape, Type elemTy, ArrayRef<int> bWidths,
- ArrayRef<int> bHeights, unsigned packingSize, int numSg,
- const xegpu::uArch::uArch *uArch) {
- int rank = dataShape.size();
- assert(rank >= 1 && "Expected at least 1D shape for ND op");
+/// Sets up the anchor layout for a store_nd operation. StoreNd picks its
+/// own layout based on uArch block parameters (it does not take a consumer
+/// layout, since it is a data sink).
+xegpu::DistributeLayoutAttr
+xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
+ if (!uArchInstruction)
+ return nullptr;
+
+ auto context = srcVecTy.getContext();
+ Type elemTy = srcVecTy.getElementType();
+ auto subgroupSize = uArch->getSubgroupSize();
+ auto dataShape = srcVecTy.getShape();
+ int rank = srcVecTy.getRank();
+ assert(rank >= 2 && "Expected at least 2D shape for ND op");
// Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- dataShape, uArch->getSubgroupSize(), bitwidth, packingSize);
+ dataShape, subgroupSize, bitwidth,
+ uArchInstruction->getPackedFormatBitSize());
if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
- auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights);
+ auto instData =
+ get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction);
- if (layoutKind == xegpu::LayoutKind::InstData)
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ assert(instData && isValidLaneLayout(*instData, laneLayout, laneData) &&
+ "Expected the store layout to satisfy uArch block constraints");
return buildInstDataLayoutWithLane(context, *instData, laneLayout,
laneData);
+ }
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
@@ -1489,33 +1493,6 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
return nullptr;
}
-/// Sets up the anchor layout for a store_nd operation. StoreNd picks its
-/// own layout based on uArch block parameters (it does not take a consumer
-/// layout, since it is a data sink).
-xegpu::DistributeLayoutAttr
-xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
- VectorType srcVecTy, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto context = srcVecTy.getContext();
- Type elemTy = srcVecTy.getElementType();
-
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
- uArch->getInstruction(
- xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
- if (!uArchInstruction)
- return nullptr;
- auto blockWHC = uArchInstruction->getBlockWidthHeightCount(elemTy);
- if (!blockWHC)
- return nullptr;
- auto [bWidths, bHeights, bCounts] = blockWHC.value();
- unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
-
- return setupGenericNdAnchorLayout(layoutKind, context, srcVecTy.getShape(),
- elemTy, bWidths, bHeights, packingSize,
- numSg, uArch);
-}
-
/// Sets up the anchor layout for a prefetch_nd operation. PrefetchNd has no
/// consumer (it produces no value), so it picks its own layout from uArch
/// block parameters.
@@ -1523,8 +1500,6 @@ xegpu::DistributeLayoutAttr
xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
xegpu::TensorDescType tdescTy, int numSg,
const xegpu::uArch::uArch *uArch) {
- auto context = tdescTy.getContext();
- Type elemTy = tdescTy.getElementType();
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
@@ -1532,15 +1507,43 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
if (!uArchInstruction)
return nullptr;
- auto blockWHC = uArchInstruction->getBlockWidthHeightCount(elemTy);
- if (!blockWHC)
- return nullptr;
- auto [bWidths, bHeights, bCounts] = blockWHC.value();
- unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
- return setupGenericNdAnchorLayout(layoutKind, context, tdescTy.getShape(),
- elemTy, bWidths, bHeights, packingSize,
- numSg, uArch);
+ auto context = tdescTy.getContext();
+ Type elemTy = tdescTy.getElementType();
+ auto subgroupSize = uArch->getSubgroupSize();
+ auto dataShape = tdescTy.getShape();
+ int rank = tdescTy.getRank();
+ assert(rank >= 2 && "Expected at least 2D shape for ND op");
+
+ // Compute the default 2D block IO lane layout / lane data.
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ dataShape, subgroupSize, bitwidth,
+ uArchInstruction->getPackedFormatBitSize());
+
+ if (layoutKind == xegpu::LayoutKind::Lane)
+ return buildLaneLayout(context, laneLayout, laneData);
+
+ auto instData =
+ get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction);
+
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ assert(instData && isValidLaneLayout(*instData, laneLayout, laneData) &&
+ "Expected the store layout to satisfy uArch block constraints");
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout,
+ laneData);
+ }
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ auto sgLayouts = getSgLayoutCandidates(dataShape, *instData, numSg);
+ if (sgLayouts.empty())
+ return nullptr;
+ return buildSgLayout(context, dataShape, sgLayouts.front(), /*dimK=*/-1);
+ }
+
+ return nullptr;
}
/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
@@ -1610,7 +1613,8 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout.getOrder());
}
}
- auto instData = get2DBlockIOInstDataLayout(dataShape, bWidths, bHeights);
+ auto instData =
+ get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction);
auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
*instData, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
hasTransform, hasTranspose);
@@ -1619,7 +1623,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- if (isValidLaneLayout(consumerLaneLayout, consumerLaneData, dataShape)) {
+ if (isValidLaneLayout(dataShape, consumerLaneLayout, consumerLaneData)) {
return consumerLayout;
} else {
auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
>From 05a09d2a937374c2ba1c6ef6c8bae136e1ab090f Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 17 Jun 2026 03:01:43 +0000
Subject: [PATCH 36/42] add full complete lane layouts for user specified
inst_data layout, add tests
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 61 ++++-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 202 ++++++++++++++-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 242 +++++++++++++-----
.../XeGPU/propagate-layout-inst-data.mlir | 131 +++++++++-
4 files changed, 561 insertions(+), 75 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index a50cd91274e2b..c20d0bef96240 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -234,10 +234,42 @@ DistributeLayoutAttr setupStoreMatrixAnchorLayout(LayoutKind layoutKind,
/// rule with inst_data as the destination shape. The resulting lane info is
/// merged with the consumer's inst_data so downstream setup* paths see a
/// fully-populated layout.
-DistributeLayoutAttr completeScatterIOLaneLayoutFromInstData(
+/// Returns the layout unchanged when it is null, has no inst_data, or already
+/// carries lane info; returns nullopt when the derived lane factorization does
+/// not divide the user's inst_data (an invalid inst_data).
+std::optional<DistributeLayoutAttr> completeScatterLoadLaneLayoutFromInstData(
DistributeLayoutAttr userSpecifiedLayout,
DistributeLayoutAttr consumerLayout, Type elemTy,
- const uArch::uArch *uArch);
+ const xegpu::uArch::LoadGatherInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Like completeScatterLoadLaneLayoutFromInstData, but for scatter stores
+/// (store_scatter / store_matrix). A store is a data sink: lane info is derived
+/// purely from inst_data using the uArch's StoreScatter per-lane store width,
+/// with no consumer layout to reuse.
+std::optional<DistributeLayoutAttr> completeScatterStoreLaneLayoutFromInstData(
+ DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::StoreScatterInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Completes a user-provided 2D-block store_nd / prefetch_nd anchor that has
+/// only inst_data. These ops are data sinks, so lane info is derived purely
+/// from inst_data using the shared BlockIOInstructionInterface; one helper
+/// serves both store_nd and prefetch_nd.
+std::optional<DistributeLayoutAttr> completeBlockStoreLaneLayoutFromInstData(
+ DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize);
+
+/// Like completeBlockStoreLaneLayoutFromInstData, but for load_nd. The consumer
+/// layout supplies the transform / transpose / packing properties; the lane
+/// factorization is recomputed from inst_data (load-side lane counts differ
+/// from the consumer's).
+std::optional<DistributeLayoutAttr> completeBlockLoadLaneLayoutFromInstData(
+ DistributeLayoutAttr specifiedLayout, DistributeLayoutAttr consumerLayout,
+ Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize);
/// Sets up the anchor layout for a store_nd operation. StoreNd does not
/// consider a consumer layout (it is a data sink), and picks its layout from
@@ -284,6 +316,31 @@ setupDpasMxLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
DistributeLayoutAttr consumerLayout, int numSg,
const uArch::uArch *uArch);
+/// Completes user-provided DPAS A/B/C-D anchors that carry only inst_data by
+/// filling in lane_layout / lane_data derived from the operand shapes (mirrors
+/// the InstData branch of setupDpasLayout). Returns nullopt if the uArch lacks
+/// the matmul instruction.
+std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr>>
+completeDpasLaneLayoutFromInstData(DistributeLayoutAttr aLayout,
+ DistributeLayoutAttr bLayout,
+ DistributeLayoutAttr cdLayout,
+ VectorType aTy, VectorType bTy,
+ VectorType cdTy, const uArch::uArch *uArch);
+
+/// Like completeDpasLaneLayoutFromInstData, but for dpas_mx: additionally
+/// re-derives the A_scale / B_scale layouts from the completed A / B layouts.
+std::optional<
+ std::tuple<DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr, DistributeLayoutAttr>>
+completeDpasMxLaneLayoutFromInstData(DistributeLayoutAttr aLayout,
+ DistributeLayoutAttr bLayout,
+ DistributeLayoutAttr cdLayout,
+ VectorType aTy, VectorType bTy,
+ VectorType cdTy, VectorType aScaleTy,
+ VectorType bScaleTy,
+ const uArch::uArch *uArch);
+
/// Gets the expected layout for a given consumer operand. This will check if
/// the owning operation of the consumer operand is one of the special layout
/// users and determine the expected layout accordingly.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d8f73615fbd16..cd75a23667a52 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1103,7 +1103,7 @@ computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
// pick their own layout from uArch.
//
// * Derivation direction between inst_data and lane_layout/lane_data. Both
-// obey the Category-A invariant inst_data = k * lane_layout * lane_data
+// obey the invariant inst_data = k * lane_layout * lane_data
// (k >= 1, per dim), but ops solve it from opposite ends:
// - Rigid-lane ops (Nd block IO, DPAS): hardware fixes lane_layout /
// lane_data first, then inst_data is built as a multiple of their
@@ -1822,10 +1822,12 @@ xegpu::setupStoreMatrixAnchorLayout(xegpu::LayoutKind layoutKind,
/// per-lane load width reported by the uArch's LoadGather instruction
/// (`getMaxLaneLoadSize`).
///
-xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeScatterLoadLaneLayoutFromInstData(
xegpu::DistributeLayoutAttr specifiedLayout,
xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
- const xegpu::uArch::uArch *uArch) {
+ const xegpu::uArch::LoadGatherInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
if (!specifiedLayout)
return specifiedLayout;
SmallVector<int64_t> specifiedInstData =
@@ -1837,14 +1839,8 @@ xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
return specifiedLayout;
// Reuse the load-side setup with inst_data as the destination shape.
- const int subgroupSize = uArch->getSubgroupSize();
auto *context = specifiedLayout.getContext();
auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
- uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
- if (!uArchInstruction)
- return specifiedLayout;
int maxChunkSize = uArchInstruction->getMaxLaneLoadSize(elemBitWidth);
if (consumerLayout) {
auto consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
@@ -1857,10 +1853,198 @@ xegpu::DistributeLayoutAttr xegpu::completeScatterIOLaneLayoutFromInstData(
}
auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
specifiedInstData, subgroupSize, maxChunkSize);
+ if (!isValidLaneLayout(specifiedInstData, defLaneLayout, defLaneData))
+ return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, defLaneLayout,
+ defLaneData);
+}
+
+/// Like completeScatterLoadLaneLayoutFromInstData, but for scatter stores. A
+/// store is a data sink, so lane info is derived purely from inst_data (bounded
+/// by the uArch's per-lane store width); there is no consumer layout to reuse.
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeScatterStoreLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::StoreScatterInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+
+ // Reuse the store-side setup with inst_data as the source shape.
+ auto *context = specifiedLayout.getContext();
+ auto elemBitWidth = elemTy.getIntOrFloatBitWidth();
+ int maxChunkSize = uArchInstruction->getMaxLaneStoreSize(elemBitWidth);
+ auto [defLaneLayout, defLaneData] = computeScatterIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, maxChunkSize);
+ if (!isValidLaneLayout(specifiedInstData, defLaneLayout, defLaneData))
+ return std::nullopt;
return buildInstDataLayoutWithLane(context, specifiedInstData, defLaneLayout,
defLaneData);
}
+/// Completes a 2D-block store/prefetch layout from its inst_data. store_nd and
+/// prefetch_nd are data sinks, so lane info is derived purely from inst_data
+/// (no consumer to reuse). One helper serves both via
+/// BlockIOInstructionInterface.
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeBlockStoreLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+
+ auto *context = specifiedLayout.getContext();
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, elemTy.getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSize());
+ if (!isValidLaneLayout(specifiedInstData, laneLayout, laneData))
+ return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, laneLayout,
+ laneData);
+}
+
+/// Like completeBlockStoreLaneLayoutFromInstData, but for load_nd. The consumer
+/// determines transform / transpose / packing, but the lane factorization is
+/// recomputed from inst_data (load-side lane counts differ from the consumer's)
+/// — mirroring the InstData branch of setupLoadNdAnchorLayout.
+std::optional<xegpu::DistributeLayoutAttr>
+xegpu::completeBlockLoadLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr specifiedLayout,
+ xegpu::DistributeLayoutAttr consumerLayout, Type elemTy,
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ const int subgroupSize) {
+ if (!specifiedLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> specifiedInstData =
+ specifiedLayout.getEffectiveInstDataAsInt();
+ if (specifiedInstData.empty())
+ return specifiedLayout;
+ if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
+ !specifiedLayout.getEffectiveLaneDataAsInt().empty())
+ return specifiedLayout;
+
+ auto *context = specifiedLayout.getContext();
+ int rank = specifiedInstData.size();
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
+
+ // Derive the transform / transpose / packing properties from the consumer's
+ // lane info, but recompute the lane factorization itself from inst_data.
+ bool hasTransform = false;
+ bool hasTranspose = false;
+ unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
+ if (consumerLayout) {
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
+ hasTransform = consumerLaneData[rank - 2] != 1;
+ hasTranspose = consumerLaneLayout[rank - 2] != 1;
+ unsigned packingFactor = hasTransform ? consumerLaneData[rank - 2]
+ : consumerLaneData[rank - 1];
+ packingSize = packingFactor * bitwidth;
+ }
+ }
+
+ auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ specifiedInstData, subgroupSize, bitwidth, packingSize, hasTransform,
+ hasTranspose);
+ if (!isValidLaneLayout(specifiedInstData, laneLayout, laneData))
+ return std::nullopt;
+ return buildInstDataLayoutWithLane(context, specifiedInstData, laneLayout,
+ laneData);
+}
+
+/// Completes user-provided DPAS A/B/C-D anchors that carry only inst_data by
+/// filling in lane_layout / lane_data. The lane factorization mirrors the
+/// InstData branch of `setupDpasLayout` (derived from each operand's shape and
+/// matmul role, B using VNNI packing); the user's inst_data is preserved.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::completeDpasLaneLayoutFromInstData(xegpu::DistributeLayoutAttr aLayout,
+ xegpu::DistributeLayoutAttr bLayout,
+ xegpu::DistributeLayoutAttr cdLayout,
+ VectorType aTy, VectorType bTy,
+ VectorType cdTy,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (!uArchInstruction)
+ return std::nullopt;
+ auto subgroupSize = uArch->getSubgroupSize();
+
+ auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
+ aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ bTy.getShape(), subgroupSize,
+ bTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
+ auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
+ cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
+ SmallVector<int64_t> instDataA = aLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> instDataB = bLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> instDataCD = cdLayout.getEffectiveInstDataAsInt();
+ if (!isValidLaneLayout(instDataA, laneLayoutA, laneDataA) ||
+ !isValidLaneLayout(instDataB, laneLayoutB, laneDataB) ||
+ !isValidLaneLayout(instDataCD, laneLayoutCD, laneDataCD))
+ return std::nullopt;
+ return std::make_tuple(
+ buildInstDataLayoutWithLane(context, instDataA, laneLayoutA, laneDataA),
+ buildInstDataLayoutWithLane(context, instDataB, laneLayoutB, laneDataB),
+ buildInstDataLayoutWithLane(context, instDataCD, laneLayoutCD,
+ laneDataCD));
+}
+
+/// Like completeDpasLaneLayoutFromInstData, but for dpas_mx: also re-derives
+/// the A_scale / B_scale layouts from the completed A / B layouts via
+/// `createScaleLayout`, matching the default path of `setupDpasMxLayout`.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::completeDpasMxLaneLayoutFromInstData(
+ xegpu::DistributeLayoutAttr aLayout, xegpu::DistributeLayoutAttr bLayout,
+ xegpu::DistributeLayoutAttr cdLayout, VectorType aTy, VectorType bTy,
+ VectorType cdTy, VectorType aScaleTy, VectorType bScaleTy,
+ const xegpu::uArch::uArch *uArch) {
+ auto completed = completeDpasLaneLayoutFromInstData(
+ aLayout, bLayout, cdLayout, aTy, bTy, cdTy, uArch);
+ if (!completed)
+ return std::nullopt;
+ auto context = aTy.getContext();
+ auto [completedA, completedB, completedCD] = *completed;
+
+ auto aScaleLayout =
+ createScaleLayout(context, aTy, aScaleTy, completedA, false, uArch);
+ auto bScaleLayout =
+ createScaleLayout(context, bTy, bScaleTy, completedB, true, uArch);
+
+ return std::make_tuple(completedA, completedB, completedCD, aScaleLayout,
+ bScaleLayout);
+}
+
/// Sets up layout for reduction operations by creating a SliceAttr for the
/// result.
///
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 039a94b307da4..a98fbbf3e16be 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -552,15 +552,32 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo prefetchLayout;
+ const uArch *uArch = getUArch(getChipStr(prefetch).value_or(""));
+ if (!uArch)
+ return;
xegpu::DistributeLayoutAttr anchorLayout = prefetch.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
prefetchLayout = LayoutInfo(anchorLayout);
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeBlockStoreLaneLayoutFromInstData(
+ anchorLayout, prefetch.getTensorDescType().getElementType(),
+ uArchInstruction, uArch->getSubgroupSize());
+ if (!completed) {
+ prefetch.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ prefetch.setLayoutAttr(*completed);
+ prefetchLayout = LayoutInfo(*completed);
+ }
} else {
auto tdescTy = prefetch.getTensorDescType();
- const uArch *uArch = getUArch(getChipStr(prefetch).value_or(""));
- if (!uArch)
- return;
-
auto numSgOrErr = getNumSg(prefetch, uArch->getSubgroupSize());
if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
prefetch.emitWarning(
@@ -707,6 +724,13 @@ void LayoutInfoPropagation::visitDpasOp(
LayoutInfo dpasBLayout;
LayoutInfo dpasCDLayout;
+ const uArch *uArch = getUArch(getChipStr(dpas).value_or(""));
+ if (!uArch)
+ return;
+ VectorType aTy = dpas.getLhsType();
+ VectorType bTy = dpas.getRhsType();
+ VectorType cdTy = dpas.getResultType();
+
xegpu::DistributeLayoutAttr anchorLayoutCD = dpas.getLayoutCdAttr();
if (hasParamsOfLayoutKind(anchorLayoutCD)) {
xegpu::DistributeLayoutAttr anchorLayoutA = dpas.getLayoutAAttr();
@@ -718,13 +742,23 @@ void LayoutInfoPropagation::visitDpasOp(
dpasALayout = LayoutInfo(anchorLayoutA);
dpasBLayout = LayoutInfo(anchorLayoutB);
dpasCDLayout = LayoutInfo(anchorLayoutCD);
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto completed = xegpu::completeDpasLaneLayoutFromInstData(
+ anchorLayoutA, anchorLayoutB, anchorLayoutCD, aTy, bTy, cdTy, uArch);
+ if (!completed) {
+ dpas.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ auto [completedA, completedB, completedCD] = *completed;
+ dpas.setLayoutAAttr(completedA);
+ dpas.setLayoutBAttr(completedB);
+ dpas.setLayoutCdAttr(completedCD);
+ dpasALayout = LayoutInfo(completedA);
+ dpasBLayout = LayoutInfo(completedB);
+ dpasCDLayout = LayoutInfo(completedCD);
+ }
} else {
- const uArch *uArch = getUArch(getChipStr(dpas).value_or(""));
- if (!uArch)
- return;
- VectorType aTy = dpas.getLhsType();
- VectorType bTy = dpas.getRhsType();
- VectorType cdTy = dpas.getResultType();
xegpu::DistributeLayoutAttr consumerLayoutAttr = nullptr;
xegpu::DistributeLayoutAttr requiredCDLayoutAttr, requiredALayout,
@@ -784,6 +818,24 @@ void LayoutInfoPropagation::visitDpasMxOp(
xegpu::DistributeLayoutAttr anchorLayoutB = dpasMx.getLayoutBAttr();
xegpu::DistributeLayoutAttr anchorLayoutCD = dpasMx.getLayoutCdAttr();
+ const uArch *uArch = getUArch(getChipStr(dpasMx).value_or(""));
+ if (!uArch)
+ return;
+
+ VectorType aTy = dpasMx.getAType();
+ VectorType bTy = dpasMx.getBType();
+ VectorType cdTy = dpasMx.getResultType();
+
+ // Get scale types if present
+ VectorType aScaleTy;
+ VectorType bScaleTy;
+ Value scaleA = dpasMx.getScaleA();
+ Value scaleB = dpasMx.getScaleB();
+ if (scaleA)
+ aScaleTy = dyn_cast<VectorType>(scaleA.getType());
+ if (scaleB)
+ bScaleTy = dyn_cast<VectorType>(scaleB.getType());
+
// Check if all layouts are already set
if (anchorLayoutA && anchorLayoutB && anchorLayoutCD &&
hasParamsOfLayoutKind(anchorLayoutA) &&
@@ -802,26 +854,34 @@ void LayoutInfoPropagation::visitDpasMxOp(
dpasMxAScaleLayout = LayoutInfo(anchorLayoutAScale);
if (anchorLayoutBScale)
dpasMxBScaleLayout = LayoutInfo(anchorLayoutBScale);
- } else {
- // Need to compute layouts
- const uArch *uArch = getUArch(getChipStr(dpasMx).value_or(""));
- if (!uArch)
- return;
-
- VectorType aTy = dpasMx.getAType();
- VectorType bTy = dpasMx.getBType();
- VectorType cdTy = dpasMx.getResultType();
-
- // Get scale types if present
- VectorType aScaleTy;
- VectorType bScaleTy;
- Value scaleA = dpasMx.getScaleA();
- Value scaleB = dpasMx.getScaleB();
- if (scaleA)
- aScaleTy = dyn_cast<VectorType>(scaleA.getType());
- if (scaleB)
- bScaleTy = dyn_cast<VectorType>(scaleB.getType());
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto completed = xegpu::completeDpasMxLaneLayoutFromInstData(
+ anchorLayoutA, anchorLayoutB, anchorLayoutCD, aTy, bTy, cdTy,
+ aScaleTy, bScaleTy, uArch);
+ if (!completed) {
+ dpasMx.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ auto [completedA, completedB, completedCD, completedAScale,
+ completedBScale] = *completed;
+ dpasMx.setLayoutAAttr(completedA);
+ dpasMx.setLayoutBAttr(completedB);
+ dpasMx.setLayoutCdAttr(completedCD);
+ dpasMxALayout = LayoutInfo(completedA);
+ dpasMxBLayout = LayoutInfo(completedB);
+ dpasMxCDLayout = LayoutInfo(completedCD);
+ if (completedAScale) {
+ dpasMx.setLayoutAScaleAttr(completedAScale);
+ dpasMxAScaleLayout = LayoutInfo(completedAScale);
+ }
+ if (completedBScale) {
+ dpasMx.setLayoutBScaleAttr(completedBScale);
+ dpasMxBScaleLayout = LayoutInfo(completedBScale);
+ }
+ }
+ } else {
xegpu::DistributeLayoutAttr consumerLayoutAttr = nullptr;
xegpu::DistributeLayoutAttr requiredCDLayoutAttr, requiredALayout,
requiredBLayout, requiredAScaleLayout, requiredBScaleLayout;
@@ -899,14 +959,32 @@ void LayoutInfoPropagation::visitStoreNdOp(
xegpu::StoreNdOp store, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo storeLayout;
+ const uArch *uArch = getUArch(getChipStr(store).value_or(""));
+ if (!uArch)
+ return;
xegpu::DistributeLayoutAttr anchorLayout = store.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
storeLayout = LayoutInfo(anchorLayout);
- } else {
- const uArch *uArch = getUArch(getChipStr(store).value_or(""));
- if (!uArch)
- return;
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockStore));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeBlockStoreLaneLayoutFromInstData(
+ anchorLayout, store.getValueType().getElementType(), uArchInstruction,
+ uArch->getSubgroupSize());
+ if (!completed) {
+ store.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ store.setLayoutAttr(*completed);
+ storeLayout = LayoutInfo(*completed);
+ }
+ } else {
auto numSgOrErr = getNumSg(store, uArch->getSubgroupSize());
if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
store.emitWarning(
@@ -935,31 +1013,37 @@ void LayoutInfoPropagation::visitLoadNdOp(
xegpu::LoadNdOp load, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo loadLayout;
+
+ const uArch *uArch = getUArch(getChipStr(load).value_or(""));
+ if (!uArch)
+ return;
+ LayoutInfo valueLayout = results[0]->getValue();
+ if (!valueLayout.isAssigned())
+ return;
+ auto consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(valueLayout.get());
xegpu::DistributeLayoutAttr anchorLayout = load.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
loadLayout = LayoutInfo(anchorLayout);
- } else {
- LayoutInfo valueLayout = results[0]->getValue();
- if (!valueLayout.isAssigned())
- return;
-
- // LoadNdOp has a transpose effect. At this stage of analysis it is not
- // expected and should be abstracted away; emit a warning and pre-transpose
- // the consumer hint so it lines up with the load's source.
- auto consumerLayoutAttr =
- dyn_cast<xegpu::DistributeLayoutAttr>(valueLayout.get());
- if (auto transpose = load.getTranspose()) {
- load.emitWarning("Transpose effect is not expected for LoadNdOp at "
- "LayoutInfoPropagation stage.");
- LayoutInfo transposed = valueLayout.transpose(transpose.value());
- consumerLayoutAttr =
- dyn_cast<xegpu::DistributeLayoutAttr>(transposed.get());
+ if (layoutKind == xegpu::LayoutKind::InstData) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeBlockLoadLaneLayoutFromInstData(
+ anchorLayout, consumerLayoutAttr, load.getType().getElementType(),
+ uArchInstruction, uArch->getSubgroupSize());
+ if (!completed) {
+ load.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ load.setLayoutAttr(*completed);
+ loadLayout = LayoutInfo(*completed);
}
-
- const uArch *uArch = getUArch(getChipStr(load).value_or(""));
- if (!uArch)
- return;
-
+ } else {
auto numSgOrErr =
getNumSg(load, uArch->getSubgroupSize(), consumerLayoutAttr);
if (layoutKind == xegpu::LayoutKind::Subgroup && failed(numSgOrErr)) {
@@ -967,7 +1051,6 @@ void LayoutInfoPropagation::visitLoadNdOp(
"Unable to determine the number of subgroups for the operation.");
return;
}
-
auto layoutAttr = xegpu::setupLoadNdAnchorLayout(
layoutKind, load.getType(), consumerLayoutAttr, numSgOrErr.value_or(0),
uArch);
@@ -1155,9 +1238,20 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
- requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
+ const auto uArchInstruction =
+ dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
+ uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeScatterLoadLaneLayoutFromInstData(
anchorLayoutAttr, consumerLayoutAttr, resVecTy.getElementType(),
- uArch);
+ uArchInstruction, uArch->getSubgroupSize());
+ if (!completed) {
+ load.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ requiredAnchorLayoutAttr = *completed;
load.setLayoutAttr(requiredAnchorLayoutAttr);
}
} else {
@@ -1201,8 +1295,21 @@ void LayoutInfoPropagation::visitStoreScatterOp(
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
- requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
- anchorLayoutAttr, nullptr, srcVecTy.getElementType(), uArch);
+ const auto uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::StoreScatter));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeScatterStoreLaneLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArchInstruction,
+ uArch->getSubgroupSize());
+ if (!completed) {
+ storeScatter.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ requiredAnchorLayoutAttr = *completed;
storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
}
} else {
@@ -1273,8 +1380,21 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
if (layoutKind == xegpu::LayoutKind::InstData) {
- requiredAnchorLayoutAttr = xegpu::completeScatterIOLaneLayoutFromInstData(
- anchorLayoutAttr, nullptr, srcVecTy.getElementType(), uArch);
+ const auto uArchInstruction =
+ dyn_cast<xegpu::uArch::StoreScatterInstructionInterface>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::StoreScatter));
+ if (!uArchInstruction)
+ return;
+ auto completed = xegpu::completeScatterStoreLaneLayoutFromInstData(
+ anchorLayoutAttr, srcVecTy.getElementType(), uArchInstruction,
+ uArch->getSubgroupSize());
+ if (!completed) {
+ storeMatrix.emitWarning(
+ "Failed to identify lane layouts for the specified inst_data.");
+ return;
+ }
+ requiredAnchorLayoutAttr = *completed;
storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
}
} else {
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index ba169a04bcb4e..e9213dd92a557 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -199,9 +199,9 @@ func.func @insert_strided_slice_inst_data_no_packing(%arg0: memref<8x32xf32>) {
gpu.module @test {
// CHECK-LABEL: func.func @insert_strided_slice_inst_data_with_packing(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x64xi8>) {
-// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 64]>} dense<1> : vector<4x64xi8>
-// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 64]>} dense<0> : vector<8x64xi8>
-// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [8, 64]>, offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
+// CHECK: %[[CST_SMALL:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 2]>} dense<1> : vector<4x64xi8>
+// CHECK: %[[CST_LARGE:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 2]>} dense<0> : vector<8x64xi8>
+// CHECK: %[[INSERT:.*]] = vector.insert_strided_slice %[[CST_SMALL]], %[[CST_LARGE]] {layout_result_0 = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 2]>, offsets = [0, 0], strides = [1, 1]} : vector<4x64xi8> into vector<8x64xi8>
func.func @insert_strided_slice_inst_data_with_packing(%arg0: memref<8x64xi8>) {
%c0 = arith.constant 0 : index
%cst_small = arith.constant dense<1> : vector<4x64xi8>
@@ -566,3 +566,128 @@ func.func @vector_shape_cast_collapse_multi_groups(%arg0: memref<8x128xf16>) {
return
}
}
+
+// -----
+// completeBlockStoreLaneLayoutFromInstData: user supplies only inst_data on a
+// store_nd; lane_layout / lane_data are completed from it (data sink, no
+// consumer). inst_data=[8,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_store_nd_inst_data(
+// CHECK: xegpu.store_nd %{{.*}} <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+func.func @complete_store_nd_inst_data(%arg0: memref<8x32xf32>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32>
+ xegpu.store_nd %cst, %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32>
+ return
+}
+}
+
+// -----
+// completeBlockStoreLaneLayoutFromInstData (prefetch path): prefetch_nd is also
+// a data sink served by the same helper. inst_data=[8,16] -> [1,16]/[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_prefetch_nd_inst_data(
+// CHECK: xegpu.prefetch_nd %{{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<8x32xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+func.func @complete_prefetch_nd_inst_data(%arg0: memref<8x32xf32>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf32> -> !xegpu.tensor_desc<8x32xf32>
+ xegpu.prefetch_nd %0[0, 0] <{l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<inst_data = [8, 16]>}> : !xegpu.tensor_desc<8x32xf32>
+ return
+}
+}
+
+// -----
+// completeBlockLoadLaneLayoutFromInstData: load_nd feeds a DPAS, so the consumer
+// supplies the transform / transpose / packing properties while lane_layout /
+// lane_data are recomputed from inst_data. A (inst=[8,16]) -> [1,16]/[1,1];
+// B (inst=[16,16], VNNI packing from the DPAS B consumer) -> [1,16]/[2,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_load_nd_inst_data(
+// CHECK: xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+// CHECK: xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
+func.func @complete_load_nd_inst_data(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16>
+ %2 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %3 = xegpu.load_nd %1[0, 0] <{layout = #xegpu.layout<inst_data = [16, 16]>}> : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+ %4 = xegpu.dpas %2, %3, %cst : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+ %5 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+ xegpu.store_nd %4, %5[0, 0] : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
+ return
+}
+}
+
+// -----
+// completeScatterStoreLaneLayoutFromInstData: user supplies only inst_data on a
+// scatter store; lane info derived purely from inst_data (data sink).
+// inst_data=[1,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_scatter_store_inst_data(
+// CHECK: xegpu.store %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+func.func @complete_scatter_store_inst_data(%src: memref<512xf32>) {
+ %mask = arith.constant dense<1> : vector<16x32xi1>
+ %offset = arith.constant dense<12> : vector<16x32xindex>
+ %data = arith.constant dense<0.000000e+00> : vector<16x32xf32>
+ xegpu.store %data, %src[%offset], %mask <{layout = #xegpu.layout<inst_data = [1, 16]>}>
+ : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+ return
+}
+}
+
+// -----
+// completeScatterLoadLaneLayoutFromInstData: user supplies only inst_data on a
+// scatter load; with no usable consumer lane info, the scatter default is used.
+// inst_data=[1,16] -> lane_layout=[1,16], lane_data=[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_scatter_load_inst_data(
+// CHECK: xegpu.load %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+func.func @complete_scatter_load_inst_data(%src: memref<512xf32>) {
+ %mask = arith.constant dense<1> : vector<16x32xi1>
+ %offset = arith.constant dense<12> : vector<16x32xindex>
+ %0 = xegpu.load %src[%offset], %mask <{layout = #xegpu.layout<inst_data = [1, 16]>}>
+ : memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
+ xegpu.store %0, %src[%offset], %mask <{layout = #xegpu.layout<inst_data = [1, 16]>}>
+ : vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
+ return
+}
+}
+
+// -----
+// completeDpasLaneLayoutFromInstData: user supplies only inst_data on all three
+// DPAS operands; lane info is completed from each operand's shape / matmul role.
+// A=[8,16]->[1,16]/[1,1]; B=[16,16]->[1,16]/[2,1] (VNNI); CD=[8,16]->[1,16]/[1,1].
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_dpas_inst_data(
+// CHECK: xegpu.dpas %{{.*}}, %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+func.func @complete_dpas_inst_data(%arg0: vector<8x16xf16>, %arg1: vector<16x16xf16>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+ %0 = xegpu.dpas %arg0, %arg1, %cst {
+ layout_a = #xegpu.layout<inst_data = [8, 16]>,
+ layout_b = #xegpu.layout<inst_data = [16, 16]>,
+ layout_cd = #xegpu.layout<inst_data = [8, 16]>}
+ : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+ return
+}
+}
+
+// -----
+// completeDpasMxLaneLayoutFromInstData: user supplies only inst_data on A/B/C-D;
+// lane info completed from shapes and scale layouts re-derived via
+// createScaleLayout. Matches the dpas_mx_f8e5m2 default-path result.
+gpu.module @test {
+// CHECK-LABEL: func.func @complete_dpas_mx_inst_data(
+// CHECK: xegpu.dpas_mx %{{.*}}, %{{.*}}, %{{.*}} scale_a = %{{[0-9a-zA-Z]+}} scale_b = %{{[0-9a-zA-Z]+}}
+// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vector<1024x32xf8E5M2>,
+ %arg2: vector<16x32xf8E8M0FNU>, %arg3: vector<32x32xf8E8M0FNU>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
+ %0 = xegpu.dpas_mx %arg0, %arg1, %cst scale_a = %arg2 scale_b = %arg3 {
+ layout_a = #xegpu.layout<inst_data = [8, 32]>,
+ layout_a_scale = #xegpu.layout<inst_data = [8, 1]>,
+ layout_b = #xegpu.layout<inst_data = [32, 16]>,
+ layout_b_scale = #xegpu.layout<inst_data = [1, 16]>,
+ layout_cd = #xegpu.layout<inst_data = [8, 16]>}
+ : (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
+ return
+}
+}
>From 57e6ba49ab1bb047cbb3b28173e16a4c108393f0 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 17 Jun 2026 21:22:24 +0000
Subject: [PATCH 37/42] add missing parameters in convert_layout propagation,
some unit test fails
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 10 +-
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 150 +++++++++---------
2 files changed, 82 insertions(+), 78 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index cd75a23667a52..734b45991ba90 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1589,8 +1589,9 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
bool hasTransform = consumerLaneData[rank - 2] != 1;
bool hasTranspose = consumerLaneLayout[rank - 2] != 1;
- unsigned packingFactor =
- hasTransform ? consumerLaneData[rank - 2] : consumerLaneData[rank - 1];
+ unsigned packingFactor = (hasTranspose || hasTransform)
+ ? consumerLaneData[rank - 2]
+ : consumerLaneData[rank - 1];
unsigned packingSize = packingFactor * elemTy.getIntOrFloatBitWidth();
if (layoutKind == xegpu::LayoutKind::InstData) {
@@ -1955,8 +1956,9 @@ xegpu::completeBlockLoadLaneLayoutFromInstData(
if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
hasTransform = consumerLaneData[rank - 2] != 1;
hasTranspose = consumerLaneLayout[rank - 2] != 1;
- unsigned packingFactor = hasTransform ? consumerLaneData[rank - 2]
- : consumerLaneData[rank - 1];
+ unsigned packingFactor = (hasTranspose || hasTransform)
+ ? consumerLaneData[rank - 2]
+ : consumerLaneData[rank - 1];
packingSize = packingFactor * bitwidth;
}
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index a98fbbf3e16be..594b30e88c2ce 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -191,64 +191,6 @@ LayoutInfo LayoutInfo::join(const LayoutInfo &lhs, const LayoutInfo &rhs) {
llvm_unreachable("Join should not be triggered by layout propagation.");
}
-/// Construct a new layout with the transposed inst_data or lane_layout,
-/// lane_data.
-LayoutInfo LayoutInfo::transpose(ArrayRef<int64_t> permutation) const {
- if (!isAssigned())
- return {};
- // Check if the permutation is valid.
- llvm::SmallSet<int64_t, 4> seen(permutation.begin(), permutation.end());
- bool hasDuplicates = seen.size() != permutation.size();
- bool withinRange = llvm::all_of(permutation, [&](int64_t idx) {
- return idx >= 0 && idx < static_cast<int64_t>(permutation.size());
- });
-
- if (!withinRange || hasDuplicates) {
- assert(false && "Invalid permutation for transpose.");
- return {};
- }
-
- SmallVector<int32_t> laneLayout;
- SmallVector<int32_t> laneData;
- SmallVector<int32_t> instData;
- SmallVector<int32_t> sgLayout;
- SmallVector<int32_t> sgData;
- SmallVector<int32_t> order;
-
- for (int64_t idx : permutation) {
- if (getLaneLayout().size()) {
- laneLayout.push_back(static_cast<int32_t>(getLaneLayout()[idx]));
- laneData.push_back(static_cast<int32_t>(getLaneData()[idx]));
- }
- if (getInstData().size())
- instData.push_back(static_cast<int32_t>(getInstData()[idx]));
- if (getSgData().size()) {
- sgLayout.push_back(static_cast<int32_t>(getSgLayout()[idx]));
- sgData.push_back(static_cast<int32_t>(getSgData()[idx]));
- }
- if (getOrder().size()) {
- order.push_back(static_cast<int32_t>(getOrder()[idx]));
- }
- }
- auto orderAttr = order.size()
- ? DenseI32ArrayAttr::get(storage.getContext(), order)
- : nullptr;
- xegpu::LayoutAttr layoutAttr;
- if (getLaneLayout().size())
- layoutAttr =
- xegpu::LayoutAttr::get(storage.getContext(), laneLayout, laneData);
- if (getInstData().size())
- layoutAttr = xegpu::LayoutAttr::get(storage.getContext(), instData);
- if (getSgData().size())
- layoutAttr = xegpu::LayoutAttr::get(
- storage.getContext(),
- DenseI32ArrayAttr::get(storage.getContext(), sgLayout),
- DenseI32ArrayAttr::get(storage.getContext(), sgData),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, orderAttr);
- return LayoutInfo(layoutAttr);
-}
-
//===----------------------------------------------------------------------===//
// LayoutInfoLattice
//===----------------------------------------------------------------------===//
@@ -1070,6 +1012,48 @@ void LayoutInfoPropagation::visitLoadNdOp(
void LayoutInfoPropagation::visitConvertLayoutOp(
xegpu::ConvertLayoutOp convert, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
+
+ LayoutInfo resultLayout = results[0]->getValue();
+ if (!resultLayout.isAssigned())
+ return;
+
+ // TODO: fix if one of the layouts is a slice layout
+ auto targetLayout =
+ dyn_cast<xegpu::LayoutAttr>(convert.getTargetLayoutAttr());
+
+ // The result's propagated layout is authoritative for the converted value.
+ // Fill the lane_layout / lane_data / order parameters the target_layout is
+ // missing from it (sg_layout / sg_data / inst_data are left as-is), so the
+ // target stays consistent with what is actually propagated downstream.
+ auto resultLayoutAttr = dyn_cast<xegpu::LayoutAttr>(resultLayout.get());
+ if (resultLayoutAttr && targetLayout) {
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !targetLayout.getLaneLayout()) {
+ targetLayout = xegpu::LayoutAttr::get(
+ convert.getContext(), targetLayout.getSgLayout(),
+ targetLayout.getSgData(), targetLayout.getInstData(),
+ resultLayoutAttr.getLaneLayout(), resultLayoutAttr.getLaneData(),
+ resultLayoutAttr.getOrder());
+ convert.setTargetLayoutAttr(targetLayout);
+ }
+ }
+
+ // Fill only the lane_layout / lane_data / order parameters the input_layout
+ // is missing from the target_layout (sg_layout / sg_data / inst_data are left
+ // as-is), so the producer side receives a fully-populated lane layout.
+ auto inputLayout = dyn_cast<xegpu::LayoutAttr>(convert.getInputLayoutAttr());
+ if (inputLayout && targetLayout) {
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !inputLayout.getLaneLayout()) {
+ auto merged = xegpu::LayoutAttr::get(
+ convert.getContext(), inputLayout.getSgLayout(),
+ inputLayout.getSgData(), inputLayout.getInstData(),
+ targetLayout.getLaneLayout(), targetLayout.getLaneData(),
+ targetLayout.getOrder());
+ convert.setInputLayoutAttr(merged);
+ }
+ }
+
xegpu::DistributeLayoutAttr anchorLayout = convert.getInputLayoutAttr();
LayoutInfo convertLayout(anchorLayout);
// Propagate the new layout to the tensor descriptor operand.
@@ -1128,9 +1112,9 @@ void LayoutInfoPropagation::visitVectorBitcastOp(
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
-/// For vector::InterleaveOp, the result has double the innermost dimension size
-/// compared to each source operand. The layout is propagated from result to
-/// sources, adjusting for the 2x size increase.
+/// For vector::InterleaveOp, the result has double the innermost dimension
+/// size compared to each source operand. The layout is propagated from result
+/// to sources, adjusting for the 2x size increase.
void LayoutInfoPropagation::visitVectorInterleaveOp(
vector::InterleaveOp interleave, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -1178,8 +1162,8 @@ void LayoutInfoPropagation::visitVectorDeinterleaveOp(
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
- // Derive the source layout from the result layout (double the innermost dim)
- // No setup function needed - just infer directly
+ // Derive the source layout from the result layout (double the innermost
+ // dim) No setup function needed - just infer directly
auto srcLayoutAttr = xegpu::inferDeinterleaveSourceLayout(consumerLayoutAttr);
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
@@ -1216,8 +1200,8 @@ void LayoutInfoPropagation::visitInsertStridedSliceOp(
operands[1]->meet(LayoutInfo(requiredResLayoutAttr)));
}
-/// Propagate the layout of the result to the tensor descriptor, mask and offset
-/// operands in LoadGatherOp.
+/// Propagate the layout of the result to the tensor descriptor, mask and
+/// offset operands in LoadGatherOp.
void LayoutInfoPropagation::visitLoadGatherOp(
xegpu::LoadGatherOp load, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -1278,8 +1262,8 @@ void LayoutInfoPropagation::visitLoadGatherOp(
propagateIfChanged(operands[2], operands[2]->meet(maskLayoutInfo));
}
-/// Set the layout for the value, tensor descriptor, offset and mask operands in
-/// the StoreScatterOp.
+/// Set the layout for the value, tensor descriptor, offset and mask operands
+/// in the StoreScatterOp.
void LayoutInfoPropagation::visitStoreScatterOp(
xegpu::StoreScatterOp storeScatter, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -1500,9 +1484,9 @@ namespace {
// ResolveLayoutConflicts
//===----------------------------------------------------------------------===//
-/// Helper to get the defining CreateNdDescOp of a tensor descriptor value. This
-/// function tries to find the defining CreateNdDescOp recursively accross
-/// control-flow boundaries.
+/// Helper to get the defining CreateNdDescOp of a tensor descriptor value.
+/// This function tries to find the defining CreateNdDescOp recursively
+/// accross control-flow boundaries.
static xegpu::CreateNdDescOp getDefiningCreateNdDescOp(Value tdescValue) {
// Try to get the defining CreateNdDescOp of the tensor descriptor.
auto definingOp = tdescValue.getDefiningOp<xegpu::CreateNdDescOp>();
@@ -1541,9 +1525,9 @@ LogicalResult ResolveLayoutConflicts::run() {
// Scan all operations in the parent op and resolve layout conflicts at
// tensor descriptor and vector use points.
auto r = parentOp->walk([&](Operation *op) -> WalkResult {
- // if the operation inputs vector and output scalar, like multi-reduction we
- // need to check if the result has layout and add a convert_layout to serve
- // as anchor op for the reduction op's layout.
+ // if the operation inputs vector and output scalar, like multi-reduction
+ // we need to check if the result has layout and add a convert_layout to
+ // serve as anchor op for the reduction op's layout.
if (isa<vector::MultiDimReductionOp>(op) || isa<vector::ReductionOp>(op)) {
for (OpResult result : op->getResults()) {
if (result.getType().isIntOrFloat()) {
@@ -1610,7 +1594,8 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
if (auto vectorTy = dyn_cast<VectorType>(vectorValue.getType());
vectorTy && vectorTy.getRank() > 1)
consumerOp->emitWarning("Expected layout for non-1D vectors.");
- return success(); // uniform non-tensor-data vector does not require layout
+ return success(); // uniform non-tensor-data vector does not require
+ // layout
}
// Region branch ops (e.g. scf.for) and their terminators (e.g. scf.yield)
// forward their operands to successor region inputs / parent op results;
@@ -1629,6 +1614,23 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
if (consumerLayout.isEqualTo(producerLayout))
return success();
+ // Consumer is a convert_layout: retarget its input_layout to the producer
+ // instead of chaining a second convert. Always safe (single source
+ // operand).
+ if (auto consumerConvert = dyn_cast<xegpu::ConvertLayoutOp>(consumerOp)) {
+ consumerConvert.setInputLayoutAttr(producerLayout);
+ return success();
+ }
+
+ // Producer is a convert_layout feeding only this use: retarget its
+ // target_layout to the consumer instead of appending another convert.
+ if (auto producerConvert =
+ vectorValue.getDefiningOp<xegpu::ConvertLayoutOp>();
+ producerConvert && vectorValue.hasOneUse()) {
+ producerConvert.setTargetLayoutAttr(consumerLayout);
+ return success();
+ }
+
// If the producer is trivially rematerializable (e.g. `vector.step`, splat
// `arith.constant`), clone it and stamp the consumer's expected layout on
// the clone instead of inserting a `xegpu.convert_layout`. The convert
>From a7c192ba89162c6449df543cb6999b9a4f7d60f2 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 18 Jun 2026 20:59:02 +0000
Subject: [PATCH 38/42] fixed bugs in load_nd layout seting and convert_layout
handling
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 148 +++++++++++-------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 83 +++-------
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 2 +-
.../XeGPU/resolve-layout-conflicts.mlir | 8 +-
4 files changed, 114 insertions(+), 127 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 734b45991ba90..d7e3f91579dd1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -940,9 +940,11 @@ getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
/// C/D.
static std::optional<SmallVector<int64_t>> get2DBlockIOInstDataLayout(
ArrayRef<int64_t> dataShape, Type elemTy,
- const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction) {
+ const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction,
+ bool transform = false, bool transpose = false) {
int rank = dataShape.size();
- auto blockWHC = uArchInstruction->getBlockWidthHeightCount(elemTy);
+ auto blockWHC =
+ uArchInstruction->getBlockWidthHeightCount(elemTy, transform, transpose);
if (!blockWHC)
return std::nullopt;
auto [bWidths, bHeights, bCounts] = blockWHC.value();
@@ -1032,19 +1034,15 @@ computeScatterIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
int64_t subgroupSize, int64_t bitwidth,
- int64_t packingSize, bool vnni = false,
- bool transpose = false) {
+ int64_t packingSize, bool transform = false) {
int64_t rank = instShape.size();
SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- int64_t packingDim = vnni ? rank - 2 : rank - 1;
- laneData[packingDim] = bitwidth < packingSize ? packingSize / bitwidth : 1;
- assert(
- !(vnni && transpose) &&
- "transpose and VNNI cannot be enabled at the same time for 2D block IO");
- if (transpose)
- laneLayout[rank - 2] = subgroupSize;
- else
- laneLayout.back() = subgroupSize;
+ int kDim = transform ? rank - 2 : rank - 1;
+ unsigned vnniFactor = packingSize / bitwidth;
+ laneData[kDim] = bitwidth < packingSize ? vnniFactor : 1;
+ laneLayout.back() =
+ std::min(subgroupSize, instShape.back() / laneData.back());
+
// assert that the lane layout and data fit in the inst shape
for (int64_t i = 0; i < rank; ++i) {
int64_t laneProduct = laneLayout[i] * laneData[i];
@@ -1566,6 +1564,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
auto context = resVecTy.getContext();
Type elemTy = resVecTy.getElementType();
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
auto subgroupSize = uArch->getSubgroupSize();
auto dataShape = resVecTy.getShape();
const auto *uArchInstruction =
@@ -1582,17 +1581,22 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
- SmallVector<int64_t> consumerOrder = consumerLayout.getEffectiveOrderAsInt();
+ auto consumerOrderAttr = consumerLayout.getOrder();
assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
"Expected consumer layout to have lane_layout and lane_data");
- bool hasTransform = consumerLaneData[rank - 2] != 1;
- bool hasTranspose = consumerLaneLayout[rank - 2] != 1;
- unsigned packingFactor = (hasTranspose || hasTransform)
- ? consumerLaneData[rank - 2]
- : consumerLaneData[rank - 1];
- unsigned packingSize = packingFactor * elemTy.getIntOrFloatBitWidth();
+ // vertical lane layout means that the blockload must be transposed
+ // note scaleA on PVC has vertical lan layout even without transposed order
+ // attr
+ bool hasTranspose =
+ consumerLaneLayout[rank - 2] > 1 && consumerLaneLayout[rank - 1] == 1;
+ bool hasTransform = !hasTranspose && 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");
+ unsigned packingfactor =
+ std::max(consumerLaneData[rank - 2], consumerLaneData[rank - 1]);
if (layoutKind == xegpu::LayoutKind::InstData) {
auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
@@ -1602,6 +1606,18 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
auto [bWidths, bHeights, bCounts] = blockWHC.value();
+ SmallVector<int64_t> laneLayout;
+ // set the laneLayout to use consumer's LaneLayout as base, but adjust its
+ // size to match the subgroupsize in case its orignal value is larger than 1
+ for (int i = 0; i < rank; i++) {
+ if (consumerLaneLayout[i] > 1)
+ laneLayout.push_back(std::max(static_cast<int64_t>(subgroupSize),
+ consumerLaneLayout[i]));
+ else
+ laneLayout.push_back(1);
+ }
+
+ // See whether the consumer's inst_data satisfies the block constraints.
int64_t height = consumerInstData[rank - 2];
int64_t width = consumerInstData[rank - 1];
auto maxBlockCount = *llvm::max_element(bCounts);
@@ -1610,28 +1626,27 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
(width % maxWidth == 0 && width / maxWidth < maxBlockCount)) {
if (llvm::is_contained(bHeights, static_cast<int>(height))) {
return buildInstDataLayoutWithLane(context, consumerInstData,
- consumerLaneLayout, consumerLaneData,
- consumerLayout.getOrder());
+ laneLayout, consumerLaneData,
+ consumerOrderAttr);
}
}
- auto instData =
- get2DBlockIOInstDataLayout(dataShape, elemTy, uArchInstruction);
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- *instData, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
- hasTransform, hasTranspose);
+ // if consumer instData size too small, try the larger one. like DPAS_MX's
+ // scale is smaller than block load
+ auto instData = get2DBlockIOInstDataLayout(
+ dataShape, elemTy, uArchInstruction, hasTransform, hasTranspose);
+ // assert instData is valid against consumer layout since
+ // transform/transpose attribute are derived from consumer layout
+ assert(instData &&
+ isValidLaneLayout(*instData, laneLayout, consumerLaneData) &&
+ "Expected the store layout to satisfy uArch block constraints");
return buildInstDataLayoutWithLane(context, *instData, laneLayout,
- laneData);
+ consumerLaneData, consumerOrderAttr);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
- if (isValidLaneLayout(dataShape, consumerLaneLayout, consumerLaneData)) {
- return consumerLayout;
- } else {
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
- hasTransform, hasTranspose);
- return buildLaneLayout(context, laneLayout, laneData);
- }
+ assert(isValidLaneLayout(dataShape, consumerLaneLayout, consumerLaneData) &&
+ "Expected the lane layout to satisfy uArch block constraints");
+ return consumerLayout;
}
return nullptr;
}
@@ -1938,38 +1953,53 @@ xegpu::completeBlockLoadLaneLayoutFromInstData(
if (!specifiedLayout.getEffectiveLaneLayoutAsInt().empty() &&
!specifiedLayout.getEffectiveLaneDataAsInt().empty())
return specifiedLayout;
+ if (!consumerLayout)
+ return specifiedLayout;
+ SmallVector<int64_t> consumerLaneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> consumerLaneData =
+ consumerLayout.getEffectiveLaneDataAsInt();
+ if (consumerLaneLayout.empty() || consumerLaneData.empty())
+ return specifiedLayout;
auto *context = specifiedLayout.getContext();
int rank = specifiedInstData.size();
- unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
- // Derive the transform / transpose / packing properties from the consumer's
- // lane info, but recompute the lane factorization itself from inst_data.
- bool hasTransform = false;
- bool hasTranspose = false;
- unsigned packingSize = uArchInstruction->getPackedFormatBitSize();
- if (consumerLayout) {
- SmallVector<int64_t> consumerLaneLayout =
- consumerLayout.getEffectiveLaneLayoutAsInt();
- SmallVector<int64_t> consumerLaneData =
- consumerLayout.getEffectiveLaneDataAsInt();
- if (!consumerLaneLayout.empty() && !consumerLaneData.empty()) {
- hasTransform = consumerLaneData[rank - 2] != 1;
- hasTranspose = consumerLaneLayout[rank - 2] != 1;
- unsigned packingFactor = (hasTranspose || hasTransform)
- ? consumerLaneData[rank - 2]
- : consumerLaneData[rank - 1];
- packingSize = packingFactor * bitwidth;
+ SmallVector<int64_t> laneLayout;
+ // set the laneLayout to use consumer's LaneLayout as base, but adjust its
+ // size to match the subgroupsize in case its orignal value is larger than 1
+ for (int i = 0; i < rank; i++) {
+ if (consumerLaneLayout[i] > 1) {
+ laneLayout.push_back(
+ std::max(static_cast<int64_t>(subgroupSize), consumerLaneLayout[i]));
+ } else {
+ laneLayout.push_back(1);
}
}
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- specifiedInstData, subgroupSize, bitwidth, packingSize, hasTransform,
- hasTranspose);
- if (!isValidLaneLayout(specifiedInstData, laneLayout, laneData))
+ // // Derive the transform / transpose / packing properties from the
+ // consumer's
+ // // lane info, but recompute the lane factorization itself from inst_data.
+ // auto consumerOrder = consumerLayout.getEffectiveOrderAsInt();
+ // // if consumerOrder is like [0, 1, 2, 3], then set hasTranspose to true
+ // bool hasTranspose =
+ // (consumerOrder == llvm::to_vector(llvm::seq<int64_t>(0, rank)));
+ // int kDim = hasTranspose ? rank - 1 : rank - 2;
+ // unsigned vnniFactor = consumerLaneData[kDim];
+ // unsigned packingSize = vnniFactor * bitwidth;
+ // bool hasTransform = vnniFactor != 1;
+ // // scaleA on PVC has vertical lan layout even without transpose
+ // bool hasVerticalLayout = consumerLaneLayout[rank - 2] > 1 &&
+ // consumerLaneLayout[rank - 1] == 1;
+
+ // auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ // specifiedInstData, subgroupSize, bitwidth, packingSize, hasTransform,
+ // hasTranspose, hasVerticalLayout);
+ if (!isValidLaneLayout(specifiedInstData, laneLayout, consumerLaneData))
return std::nullopt;
return buildInstDataLayoutWithLane(context, specifiedInstData, laneLayout,
- laneData);
+ consumerLaneData,
+ consumerLayout.getOrder());
}
/// Completes user-provided DPAS A/B/C-D anchors that carry only inst_data by
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 594b30e88c2ce..1b4a949f0e853 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -130,48 +130,6 @@ struct LayoutInfo {
void set(const xegpu::DistributeLayoutAttr &layout) { storage = layout; }
};
-SmallVector<int> LayoutInfo::getLaneLayout() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveLaneLayoutAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getLaneData() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveLaneDataAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getInstData() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveInstDataAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getSgLayout() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveSgLayoutAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getSgData() const {
- if (!isAssigned())
- return {};
- return llvm::map_to_vector(storage.getEffectiveSgDataAsInt(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
-SmallVector<int> LayoutInfo::getOrder() const {
- if (!isAssigned() || !storage.getOrder())
- return {};
- return llvm::map_to_vector(storage.getOrder().asArrayRef(),
- [](int64_t val) { return static_cast<int>(val); });
-}
-
void LayoutInfo::print(raw_ostream &os) const {
if (isAssigned()) {
os << storage;
@@ -967,7 +925,8 @@ void LayoutInfoPropagation::visitLoadNdOp(
xegpu::DistributeLayoutAttr anchorLayout = load.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
loadLayout = LayoutInfo(anchorLayout);
- if (layoutKind == xegpu::LayoutKind::InstData) {
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !consumerLayoutAttr.getEffectiveLaneLayoutAsInt().empty()) {
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
uArch->getInstruction(
@@ -1014,42 +973,43 @@ void LayoutInfoPropagation::visitConvertLayoutOp(
ArrayRef<const LayoutInfoLattice *> results) {
LayoutInfo resultLayout = results[0]->getValue();
- if (!resultLayout.isAssigned())
- return;
// TODO: fix if one of the layouts is a slice layout
- auto targetLayout =
+ auto targetLayoutAttr =
dyn_cast<xegpu::LayoutAttr>(convert.getTargetLayoutAttr());
+ auto inputLayoutAttr =
+ dyn_cast<xegpu::LayoutAttr>(convert.getInputLayoutAttr());
// The result's propagated layout is authoritative for the converted value.
// Fill the lane_layout / lane_data / order parameters the target_layout is
// missing from it (sg_layout / sg_data / inst_data are left as-is), so the
// target stays consistent with what is actually propagated downstream.
- auto resultLayoutAttr = dyn_cast<xegpu::LayoutAttr>(resultLayout.get());
- if (resultLayoutAttr && targetLayout) {
+ auto resultLayoutAttr = resultLayout.isAssigned()
+ ? dyn_cast<xegpu::LayoutAttr>(resultLayout.get())
+ : nullptr;
+ if (resultLayoutAttr && targetLayoutAttr) {
if (layoutKind == xegpu::LayoutKind::InstData &&
- !targetLayout.getLaneLayout()) {
- targetLayout = xegpu::LayoutAttr::get(
- convert.getContext(), targetLayout.getSgLayout(),
- targetLayout.getSgData(), targetLayout.getInstData(),
+ !targetLayoutAttr.getLaneLayout()) {
+ targetLayoutAttr = xegpu::LayoutAttr::get(
+ convert.getContext(), targetLayoutAttr.getSgLayout(),
+ targetLayoutAttr.getSgData(), targetLayoutAttr.getInstData(),
resultLayoutAttr.getLaneLayout(), resultLayoutAttr.getLaneData(),
resultLayoutAttr.getOrder());
- convert.setTargetLayoutAttr(targetLayout);
+ convert.setTargetLayoutAttr(targetLayoutAttr);
}
}
// Fill only the lane_layout / lane_data / order parameters the input_layout
// is missing from the target_layout (sg_layout / sg_data / inst_data are left
// as-is), so the producer side receives a fully-populated lane layout.
- auto inputLayout = dyn_cast<xegpu::LayoutAttr>(convert.getInputLayoutAttr());
- if (inputLayout && targetLayout) {
+ if (inputLayoutAttr && targetLayoutAttr) {
if (layoutKind == xegpu::LayoutKind::InstData &&
- !inputLayout.getLaneLayout()) {
+ !inputLayoutAttr.getLaneLayout()) {
auto merged = xegpu::LayoutAttr::get(
- convert.getContext(), inputLayout.getSgLayout(),
- inputLayout.getSgData(), inputLayout.getInstData(),
- targetLayout.getLaneLayout(), targetLayout.getLaneData(),
- targetLayout.getOrder());
+ convert.getContext(), inputLayoutAttr.getSgLayout(),
+ inputLayoutAttr.getSgData(), inputLayoutAttr.getInstData(),
+ targetLayoutAttr.getLaneLayout(), targetLayoutAttr.getLaneData(),
+ targetLayoutAttr.getOrder());
convert.setInputLayoutAttr(merged);
}
}
@@ -1221,7 +1181,8 @@ void LayoutInfoPropagation::visitLoadGatherOp(
if (hasParamsOfLayoutKind(anchorLayoutAttr)) {
requiredAnchorLayoutAttr = anchorLayoutAttr;
- if (layoutKind == xegpu::LayoutKind::InstData) {
+ if (layoutKind == xegpu::LayoutKind::InstData &&
+ !consumerLayoutAttr.getEffectiveLaneLayoutAsInt().empty()) {
const auto uArchInstruction =
dyn_cast<xegpu::uArch::LoadGatherInstructionInterface>(
uArch->getInstruction(xegpu::uArch::InstructionKind::LoadGather));
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 55bfc317f89fe..a9c73b3b84025 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -988,7 +988,7 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
Value newSource = op.getSource();
SmallVector<Value> newOps;
- if (inputLayout && targetLayout) {
+ if (inputLayout && targetLayout && !inputLayout.isEqualTo(targetLayout)) {
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
SmallVector<Value> convertedValues =
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 40e907be6e4a1..55d0e64bb2c65 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -254,8 +254,7 @@ func.func @conflict_postop() {
}
// CHECK-LABEL: func.func @convert_layout
-// CHECK: %[[V0:.*]] = xegpu.convert_layout %[[CST:.*]] <{input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>, target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>}> : vector<32x128xf32>
-// CHECK: %[[V1:.*]] = xegpu.convert_layout %[[V0]] <{input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>, target_layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 16]>}> : vector<32x128xf32>
+// CHECK: %[[V1:.*]] = xegpu.convert_layout %[[V0:.*]] <{input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [4, 32]>, target_layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 16]>}> : vector<32x128xf32>
func.func @convert_layout() {
%src0 = arith.constant
{layout_result_0 = #xegpu.layout<sg_layout=[8, 4], sg_data=[4, 32]>}
@@ -354,10 +353,7 @@ func.func @extract_source_conflict_with_order() -> vector<16x32xf16> {
// CHECK-LABEL: func.func @convert_layout_bridge_input_mismatch
// CHECK: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
// CHECK-NEXT: %[[BRIDGE:.*]] = xegpu.convert_layout %[[V0]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
-// CHECK-SAME: : vector<32x32xf16>
-// CHECK-NEXT: %[[CVT:.*]] = xegpu.convert_layout %[[BRIDGE]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [32, 16]>}>
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [32, 16]>}>
// CHECK-SAME: : vector<32x32xf16>
gpu.module @test_convert_layout_bridge {
func.func @convert_layout_bridge_input_mismatch() {
>From 63325893dfbf8544acef4855d66136f3edb6293f Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 18 Jun 2026 21:29:04 +0000
Subject: [PATCH 39/42] fix format
---
.../Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 9 +++++----
1 file changed, 5 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 1b4a949f0e853..10ff008df9704 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -272,18 +272,19 @@ class LayoutInfoPropagation
visitOperation(Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) override;
- void visitBranchOperand(OpOperand &operand) override{};
+ void visitBranchOperand(OpOperand &operand) override {};
- void visitCallOperand(OpOperand &operand) override{};
+ void visitCallOperand(OpOperand &operand) override {};
void
visitNonControlFlowArguments(RegionSuccessor &successor,
- ArrayRef<BlockArgument> arguments) override{};
+ ArrayRef<BlockArgument> arguments) override {};
void
visitExternalCall(CallOpInterface call,
ArrayRef<LayoutInfoLattice *> operands,
- ArrayRef<const LayoutInfoLattice *> results) override{};
+ ArrayRef<const LayoutInfoLattice *> results) override {
+ };
void setToExitState(LayoutInfoLattice *lattice) override {
(void)lattice->meet(LayoutInfo());
>From 045976a935d058bf078116503630c17120348f45 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 18 Jun 2026 21:34:34 +0000
Subject: [PATCH 40/42] polish
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 3 ---
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 7 +++----
2 files changed, 3 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ba754ac3d9fb3..61f353c7e23b9 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1564,7 +1564,6 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
auto context = resVecTy.getContext();
Type elemTy = resVecTy.getElementType();
- unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
auto subgroupSize = uArch->getSubgroupSize();
auto dataShape = resVecTy.getShape();
const auto *uArchInstruction =
@@ -1595,8 +1594,6 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
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");
- unsigned packingfactor =
- std::max(consumerLaneData[rank - 2], consumerLaneData[rank - 1]);
if (layoutKind == xegpu::LayoutKind::InstData) {
auto blockWHC = uArchInstruction->getBlockWidthHeightCount(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 10ff008df9704..3aad86369c16f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -280,10 +280,9 @@ class LayoutInfoPropagation
visitNonControlFlowArguments(RegionSuccessor &successor,
ArrayRef<BlockArgument> arguments) override {};
- void
- visitExternalCall(CallOpInterface call,
- ArrayRef<LayoutInfoLattice *> operands,
- ArrayRef<const LayoutInfoLattice *> results) override {
+ void visitExternalCall(CallOpInterface call,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) override {
};
void setToExitState(LayoutInfoLattice *lattice) override {
>From 07ce0a98158cd2677ad94c51886d39f02df3f6ad Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 18 Jun 2026 22:27:21 +0000
Subject: [PATCH 41/42] fix uArch table
---
.../mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h | 62 ++++++++-----------
1 file changed, 25 insertions(+), 37 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index 547708b65abe1..ff80a77b28d37 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -99,61 +99,49 @@ struct Subgroup2DBlockLoadInstruction : public Instruction,
static const int kHeightAtLeast1[] = {1, 2, 4, 8, 16, 32};
static const int kHeightAtLeast8[] = {8, 16, 32};
static const int kHeightAtLeast16[] = {16, 32};
- static const int kHeightAtLeast32[] = {32};
+ static const int kHeight32[] = {32};
+ static const int kHeight64[] = {64};
+ static const int kWidth64[] = {64};
static const int kWidth32[] = {32};
static const int kWidth16[] = {16};
static const int kWidthAtLeast16[] = {16, 32};
+ static const int kWidthAtLeast32[] = {32, 64};
static const int kWidth8[] = {8};
static const int32_t kCount1[] = {1};
static const int32_t kCount2[] = {1, 2};
static const int32_t kCount4[] = {1, 2, 4};
static const int32_t kCount4Only[] = {4};
- // (elemBytes, transform, transpose, upConvert)
+ // (elemBits, transform, transpose, upConvert)
using Key = std::tuple<int, uint8_t, uint8_t, uint8_t>;
// (widths, heights, counts)
using Value = std::tuple<llvm::ArrayRef<int32_t>, llvm::ArrayRef<int32_t>,
llvm::ArrayRef<int32_t>>;
+ // The table is keyed on element bit width so sub-byte elements can be
+ // expressed directly. 4-bit elements are packed two-per-byte, so their
+ // widths (or heights, when transformed) are double the 8-bit rows.
static const llvm::DenseMap<Key, Value> kMap = {
- {{1, false, false, false}, {kWidthAtLeast16, kHeightAtLeast1, kCount2}},
- {{1, false, false, true}, {kWidth16, kHeightAtLeast8, kCount4Only}},
- {{2, false, false, false}, {kWidth16, kHeightAtLeast1, kCount2}},
- {{4, false, false, false}, {kWidth16, kHeightAtLeast1, kCount1}},
+ {{8, false, false, false}, {kWidthAtLeast16, kHeightAtLeast1, kCount2}},
+ {{8, false, false, true}, {kWidth16, kHeightAtLeast8, kCount4Only}},
+ {{16, false, false, false}, {kWidth16, kHeightAtLeast1, kCount2}},
+ {{32, false, false, false}, {kWidth16, kHeightAtLeast1, kCount1}},
// Block Loads with Transform:
- {{1, true, false, false}, {kWidth16, kHeightAtLeast32, kCount4}},
- {{2, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
+ {{8, true, false, false}, {kWidth16, kHeight32, kCount4}},
+ {{16, true, false, false}, {kWidth16, kHeightAtLeast16, kCount2}},
// Block Loads with Transpose:
- {{1, false, true, false}, {kWidth32, kHeightAtLeast16, kCount1}},
- {{2, false, true, false}, {kWidth16, kHeightAtLeast16, kCount1}},
- {{4, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}}};
- int elemByteSize = elemTy.getIntOrFloatBitWidth() / 8;
- // handle sub-byte elements by treating them as 1 byte elements
- if (elemByteSize == 0)
- elemByteSize = 1;
- auto it = kMap.find({elemByteSize, hasTransform, hasTranspose, upConv});
- if (it != kMap.end()) {
- // for sub-byte elements, need to double width retrieved from map since
- // the map is based on byte-sized elements
- if (elemTy.getIntOrFloatBitWidth() < 8) {
- int subByteElemCount = 8 / elemTy.getIntOrFloatBitWidth();
- auto [widths, heights, counts] = it->second;
- if (hasTransform) {
- llvm::SmallVector<int, 8> newHeights;
- for (int h : heights)
- newHeights.push_back(h * subByteElemCount);
- return std::make_tuple(widths, llvm::ArrayRef<int>(newHeights),
- counts);
- } else {
- llvm::SmallVector<int, 8> newWidths;
- for (int w : widths)
- newWidths.push_back(w * subByteElemCount);
- return std::make_tuple(llvm::ArrayRef<int>(newWidths), heights,
- counts);
- }
- }
+ {{8, false, true, false}, {kWidth32, kHeightAtLeast16, kCount1}},
+ {{16, false, true, false}, {kWidth16, kHeightAtLeast16, kCount1}},
+ {{32, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}},
+ // 4-bit elements (sub-byte):
+ {{4, false, false, false}, {kWidthAtLeast32, kHeightAtLeast1, kCount2}},
+ {{4, false, false, true}, {kWidth32, kHeightAtLeast8, kCount4Only}},
+ {{4, true, false, false}, {kWidth16, kHeight64, kCount4}},
+ {{4, false, true, false}, {kWidth64, kHeightAtLeast16, kCount1}}};
+ int elemBitSize = elemTy.getIntOrFloatBitWidth();
+ auto it = kMap.find({elemBitSize, hasTransform, hasTranspose, upConv});
+ if (it != kMap.end())
return it->second;
- }
return std::nullopt;
}
>From e5fb64f6a59c0f6ccf8c2a0c4aec740013d29f47 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 18 Jun 2026 22:50:12 +0000
Subject: [PATCH 42/42] address claude review comments
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 67 +++++++------------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 7 ++
2 files changed, 32 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 61f353c7e23b9..435c39eb01ce0 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -394,9 +394,9 @@ static xegpu::LayoutAttr buildInstDataLayoutWithLane(
orderAttr);
}
-bool isValidLaneLayout(ArrayRef<int64_t> dataShape,
- ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData) {
+static bool isValidLaneLayout(ArrayRef<int64_t> dataShape,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
int rank = dataShape.size();
for (int dim = 0; dim < rank; ++dim) {
int64_t laneProduct = laneLayout[dim] * laneData[dim];
@@ -971,8 +971,7 @@ static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
SmallVector<int64_t>>>
getDpasInstDataLayouts(
VectorType aTy, VectorType bTy, VectorType cdTy,
- const xegpu::uArch::MMAInstructionInterface *uArchInstruction,
- bool isDpasMx = false) {
+ const xegpu::uArch::MMAInstructionInterface *uArchInstruction) {
// M dimension is the second-to-last dim of A (handles batch dims).
const unsigned dataALen = aTy.getShape()[aTy.getRank() - 2];
@@ -1381,8 +1380,7 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
cdTy.getShape(), subgroupSize,
cdTy.getElementType().getIntOrFloatBitWidth(),
cdTy.getElementType().getIntOrFloatBitWidth());
- auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction,
- /*isDpasMx=*/true);
+ auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction);
if (!instDataVecs)
return std::nullopt;
@@ -1527,7 +1525,7 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
if (layoutKind == xegpu::LayoutKind::InstData) {
assert(instData && isValidLaneLayout(*instData, laneLayout, laneData) &&
- "Expected the store layout to satisfy uArch block constraints");
+ "Expected the prefetch layout to satisfy uArch block constraints");
return buildInstDataLayoutWithLane(context, *instData, laneLayout,
laneData);
}
@@ -1586,7 +1584,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
"Expected consumer layout to have lane_layout and lane_data");
// vertical lane layout means that the blockload must be transposed
- // note scaleA on PVC has vertical lan layout even without transposed order
+ // note scaleA on PVC has vertical lane layout even without transposed order
// attr
bool hasTranspose =
consumerLaneLayout[rank - 2] > 1 && consumerLaneLayout[rank - 1] == 1;
@@ -1605,7 +1603,8 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
SmallVector<int64_t> laneLayout;
// set the laneLayout to use consumer's LaneLayout as base, but adjust its
- // size to match the subgroupsize in case its orignal value is larger than 1
+ // size to match the subgroupsize in case its original value is larger than
+ // 1
for (int i = 0; i < rank; i++) {
if (consumerLaneLayout[i] > 1)
laneLayout.push_back(std::max(static_cast<int64_t>(subgroupSize),
@@ -1636,7 +1635,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
// transform/transpose attribute are derived from consumer layout
assert(instData &&
isValidLaneLayout(*instData, laneLayout, consumerLaneData) &&
- "Expected the store layout to satisfy uArch block constraints");
+ "Expected the load layout to satisfy uArch block constraints");
return buildInstDataLayoutWithLane(context, *instData, laneLayout,
consumerLaneData, consumerOrderAttr);
}
@@ -1759,7 +1758,7 @@ setupGenericStoreAnchorLayout(xegpu::LayoutKind layoutKind,
ArrayRef<int64_t> srcShape, int subgroupSize) {
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true &&
+ assert(false &&
"subgroup layout assignment not supported for storeScatter.");
return nullptr;
}
@@ -1931,10 +1930,10 @@ xegpu::completeBlockStoreLaneLayoutFromInstData(
laneData);
}
-/// Like completeBlockStoreLaneLayoutFromInstData, but for load_nd. The consumer
-/// determines transform / transpose / packing, but the lane factorization is
-/// recomputed from inst_data (load-side lane counts differ from the consumer's)
-/// — mirroring the InstData branch of setupLoadNdAnchorLayout.
+/// Like completeBlockStoreLaneLayoutFromInstData, but for load_nd. The
+/// consumer's lane_data and order are reused as-is; lane_layout is rebuilt from
+/// the consumer's lane_layout, bumping every non-unit dim up to the subgroup
+/// size. The user-provided inst_data is preserved.
std::optional<xegpu::DistributeLayoutAttr>
xegpu::completeBlockLoadLaneLayoutFromInstData(
xegpu::DistributeLayoutAttr specifiedLayout,
@@ -1964,7 +1963,7 @@ xegpu::completeBlockLoadLaneLayoutFromInstData(
SmallVector<int64_t> laneLayout;
// set the laneLayout to use consumer's LaneLayout as base, but adjust its
- // size to match the subgroupsize in case its orignal value is larger than 1
+ // size to match the subgroupsize in case its original value is larger than 1
for (int i = 0; i < rank; i++) {
if (consumerLaneLayout[i] > 1) {
laneLayout.push_back(
@@ -1974,24 +1973,6 @@ xegpu::completeBlockLoadLaneLayoutFromInstData(
}
}
- // // Derive the transform / transpose / packing properties from the
- // consumer's
- // // lane info, but recompute the lane factorization itself from inst_data.
- // auto consumerOrder = consumerLayout.getEffectiveOrderAsInt();
- // // if consumerOrder is like [0, 1, 2, 3], then set hasTranspose to true
- // bool hasTranspose =
- // (consumerOrder == llvm::to_vector(llvm::seq<int64_t>(0, rank)));
- // int kDim = hasTranspose ? rank - 1 : rank - 2;
- // unsigned vnniFactor = consumerLaneData[kDim];
- // unsigned packingSize = vnniFactor * bitwidth;
- // bool hasTransform = vnniFactor != 1;
- // // scaleA on PVC has vertical lan layout even without transpose
- // bool hasVerticalLayout = consumerLaneLayout[rank - 2] > 1 &&
- // consumerLaneLayout[rank - 1] == 1;
-
- // auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- // specifiedInstData, subgroupSize, bitwidth, packingSize, hasTransform,
- // hasTranspose, hasVerticalLayout);
if (!isValidLaneLayout(specifiedInstData, laneLayout, consumerLaneData))
return std::nullopt;
return buildInstDataLayoutWithLane(context, specifiedInstData, laneLayout,
@@ -2229,7 +2210,7 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
consumerSliceLayout
? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
: SmallVector<int64_t>({});
- // A[i] reduced from A[i, j] is stored out directly, use veritical Lane
+ // 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 &&
@@ -2290,11 +2271,13 @@ xegpu::setupReductionResultLayout(xegpu::LayoutKind layoutKind,
xegpu::LayoutAttr srcLayout;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- assert(true && "subgroup layout assignment not supported for reduction (op "
- "is not expected at this level).");
+ assert(false &&
+ "subgroup layout assignment not supported for reduction (op "
+ "is not expected at this level).");
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- assert(true && "instData layout assignment not supported for reduction (op "
- "is not expected at this level).");
+ assert(false &&
+ "instData layout assignment not supported for reduction (op "
+ "is not expected at this level).");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
SmallVector<int64_t> laneLayout(1), laneData(1);
laneLayout[0] = std::min(static_cast<int64_t>(subgroupSize), srcShape[0]);
@@ -2465,8 +2448,8 @@ xegpu::DistributeLayoutAttr xegpu::setupInsertStridedSliceResultLayout(
if (layoutKind == xegpu::LayoutKind::Subgroup ||
layoutKind == xegpu::LayoutKind::InstData) {
- assert(true &&
- "subgroup layout assignment not supported for insertStridedSlice.");
+ assert(false && "subgroup/instData layout assignment not supported for "
+ "insertStridedSlice.");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
for (int dim = 0; dim < srcRank; dim++) {
assert(srcShape[dim] % consumerLaneLayout[dim] == 0 &&
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 3aad86369c16f..8600492e4bf41 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -612,6 +612,13 @@ void LayoutInfoPropagation::visitShapeCastOp(
xegpu::DistributeLayoutAttr srcLayoutAttr =
xegpu::inferShapeCastSourceLayout(resultLayoutAttr, resShape, srcShape);
+ // TODO: turn this into a real pass failure once propagation failures are
+ // wired to signalPassFailure().
+ if (!srcLayoutAttr) {
+ shapeCast.emitWarning("Failed to infer source layout for shape_cast; "
+ "unsupported shape-cast pattern.");
+ return;
+ }
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
More information about the Mlir-commits
mailing list