[Mlir-commits] [mlir] [MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (PR #202868)
Jianhui Li
llvmlistbot at llvm.org
Tue Jun 9 23:47:34 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/202868
Motivation
Enhance setup* rules in layout propagation to pass lane_layout, and lane_data information during inst_data propagation, so that the propagation can have lane level information when choosing an optimal inst_data. This branch makes that relationship explicit and uniform across all setup rules.
Invariant
All setup rules now produce layouts that satisfy:
- Nd ops + dpas/dpas_mx: inst_data = k * (lane_layout * lane_data), k ≥ 1
- Scatter/matrix ops + non-anchor ops: inst_data = lane_layout * lane_data
Key changes in XeGPULayoutImpl
- New setupStoreNdAnchorLayout, setupPrefetchNdAnchorLayout, setupLoadNdAnchorLayout (Nd ops have rigid lane info; inst_data fits the lane factorization).
- setupMultiReductionResultLayout reorganized so InstData and Lane branches share the same lane-layout logic.
- setupGenericLoadAnchorLayout (scatter) takes consumer's inst_data as-is and consumer's lane info when present.
- New complete*LayoutFromInstData helpers that fill in lane info on user-provided anchors that specify only inst_data — by
re-running the op's Lane setup with inst_data as the destination shape.
- inferShapeCastSourceLayout preserves lane info through 1D ↔ ND collapse-style casts.
- createScaleLayout (dpas_mx) caps scale lane_layout by inst_data so the scale operand load_nd satisfies the multiple-of
invariant.
Assisted-by-claude
>From c88c24241cc506912b689ce2b48e5e54ca3d026e 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/20] [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 11b36f56efa30..e566e013c6bff 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -712,6 +712,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);
@@ -797,6 +812,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.
///
@@ -813,9 +904,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:
@@ -943,22 +1038,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));
}
@@ -999,6 +1107,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
@@ -1030,53 +1190,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
@@ -1098,52 +1228,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.
@@ -1162,21 +1258,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 f84d29aa51164..e5d4bd9f1ff2a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -966,8 +966,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 6f587959b697d..aacf07222ee4a 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 a6907fa630d93..a1009d1d66ab8 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 87ada366684a74e713359f8203367705435a6575 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/20] 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 e566e013c6bff..cc62aafd75e85 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -854,10 +854,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";
@@ -1040,10 +1042,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);
@@ -1063,10 +1068,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));
}
@@ -1294,8 +1302,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;
@@ -1309,31 +1317,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;
}
@@ -1354,15 +1347,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) {
@@ -1374,10 +1366,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.
@@ -1393,9 +1385,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);
@@ -1407,28 +1398,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;
}
@@ -1447,15 +1426,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();
@@ -1466,9 +1446,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 aacf07222ee4a..6df6693419579 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 a1009d1d66ab8..90f69fc34db1e 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 dd4db0d647ec720f3dbf967896880b05b365852d 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/20] 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 cc62aafd75e85..114a8708ddd67 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -854,15 +854,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;
@@ -873,15 +864,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)
@@ -1118,7 +1100,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).
@@ -1288,18 +1270,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,
@@ -1313,20 +1313,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;
}
@@ -1372,42 +1370,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;
}
@@ -1578,8 +1591,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>>
@@ -1709,9 +1722,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,
@@ -1802,8 +1815,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,
@@ -1929,8 +1942,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(
@@ -1999,8 +2012,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 3ed2ddaff632e9555e919a366fb3893f3e5fa942 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/20] 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 114a8708ddd67..3903465f6a5cc 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -827,6 +827,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
@@ -1036,12 +1071,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.
@@ -1270,26 +1301,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
@@ -1312,19 +1345,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;
}
@@ -1377,16 +1416,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};
}
@@ -1415,12 +1455,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;
}
@@ -1466,13 +1510,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.
@@ -1482,17 +1528,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:
@@ -1691,7 +1747,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.");
@@ -1703,19 +1764,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;
@@ -1854,12 +1910,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 71b0b0e0875c371a3a8d734339ff171820dcb91a 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/20] 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 3903465f6a5cc..608b49816d395 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1510,6 +1510,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
@@ -1558,7 +1705,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 bb0079b4eac3bb43ec02889cbe8dc05e9ad88e6d 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/20] 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 608b49816d395..a0207cacbb938 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1516,25 +1516,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");
@@ -1550,8 +1597,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);
@@ -1566,8 +1620,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 &&
@@ -1579,12 +1642,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;
@@ -1627,7 +1690,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
@@ -1654,7 +1717,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 df472290d19df1220b959062de068a95249b3ce4 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/20] 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 | 101 ++++++++++++++++--
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 28 +++--
3 files changed, 132 insertions(+), 14 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 a0207cacbb938..db036ef993c87 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -757,8 +757,10 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// Extract layout info from result's innermost dimension and apply to
// source's innermost dimension while setting all other dimensions to 1.
- // The inferred layout is restricted by srcShape to ensure it fits within
- // the source dimensions.
+ // Whichever fields the result has (inst_data, lane_layout, lane_data) are
+ // all carried over together so downstream consumers see a consistent
+ // layout (no lane info lost during the collapse-style cast).
+ //
// Examples 1:
// srcShape=[8, 16, 32], resShape=[1, 4096]
// resInstData=[1, 16]
@@ -768,9 +770,12 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// resLaneLayout=[16], resLaneData=[2]
// -> inferredLaneLayout=[1, 1, 16]
// -> inferredLaneData=[1, 1, min(2, 64/16)]=[1, 1, 2]
+ auto toI32Attr = [&](ArrayRef<int> v) {
+ return DenseI32ArrayAttr::get(context, v);
+ };
+ DenseI32ArrayAttr instDataAttr;
if (resInstData.size() != 0) {
- // assert resInstData must be 1 for all but the innermost dim
for (int i = 0; i < resShapeSize - 1; i++) {
assert(resInstData[i] == 1 &&
"only innermost dim can have non-unit instData");
@@ -778,9 +783,11 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
SmallVector<int> inferredInstData(srcShapeSize, 1);
inferredInstData[srcShapeSize - 1] =
std::min(resInstData[resShapeSize - 1], srcShape[srcShapeSize - 1]);
- return xegpu::LayoutAttr::get(context, inferredInstData);
+ instDataAttr = toI32Attr(inferredInstData);
}
+ DenseI32ArrayAttr laneLayoutAttr;
+ DenseI32ArrayAttr laneDataAttr;
if (resLaneLayout.size() != 0) {
for (int i = 0; i < resShapeSize - 1; i++) {
assert(resLaneData[i] == 1 &&
@@ -793,9 +800,15 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
inferredLaneLayout.back() = resLaneLayout.back();
inferredLaneData.back() = std::min(
resLaneData.back(), srcShape.back() / inferredLaneLayout.back());
- return xegpu::LayoutAttr::get(context, inferredLaneLayout,
- inferredLaneData);
+ laneLayoutAttr = toI32Attr(inferredLaneLayout);
+ laneDataAttr = toI32Attr(inferredLaneData);
}
+
+ if (instDataAttr || laneLayoutAttr)
+ return xegpu::LayoutAttr::get(context, /*sg_layout=*/nullptr,
+ /*sg_data=*/nullptr, instDataAttr,
+ laneLayoutAttr, laneDataAttr,
+ /*order=*/nullptr);
}
llvm_unreachable("running into unsupported shape cast scenarios");
return nullptr;
@@ -1510,6 +1523,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 8367937a167e2b84ae3239012605197475a98470 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/20] 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 db036ef993c87..2df1519780da5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1683,6 +1683,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 340b809f5231750ffcf8b2b5235453d9a155dcdb 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/20] 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 2df1519780da5..eea3bc1b51bb6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -845,9 +845,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);
@@ -857,12 +860,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);
@@ -872,7 +876,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
@@ -1707,8 +1711,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
@@ -1752,7 +1760,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)
@@ -1859,16 +1868,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 0884a3fc0b00bca7d40b475a0b8a03c2cca7deaf 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/20] 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 fe1e99cf611590840a794e7747439aebb95cb5a8 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/20] 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 6df6693419579..f073ffed91dd6 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 895da336d590d3bcdecbb2bc0dbea4a8b68e772e 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/20] 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 eea3bc1b51bb6..a1574ee89066a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1341,15 +1341,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,
@@ -1360,28 +1369,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 f073ffed91dd6..397ae7f50ce5c 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 0fb1c3ddac193185865bc3821a07c3eff741644f 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/20] 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 a1574ee89066a..6c8bef93ac9e1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1748,19 +1748,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.
@@ -2251,15 +2258,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 c2cb1d930227412ceabcdf04cd150864d0589ae2 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/20] 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 397ae7f50ce5c..6419c5bc555f3 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 93899a4e6719f808b1c34f5acdee309c1a5e3f30 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/20] 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 6419c5bc555f3..e910e39c14a50 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 b465edeb44358dda31a1d0f65abaec613fcc7c97 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/20] 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 e910e39c14a50..9ef323169dca7 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 5f22456ec28c3d4ae2a715c528576e58ae389e0d 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/20] 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 9ef323169dca7..6c45d4c5d32ec 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 52c39b30a41d6610427612df6dc0d0d01ce66c25 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/20] 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 6c8bef93ac9e1..5cbb5cbc70c96 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1701,13 +1701,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 6c45d4c5d32ec..68df00a80337b 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 40368f7b315aebee2486063998bd7bd711b873cc 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/20] 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 90f69fc34db1e..4be18158bdfcf 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 b89947f8da967764dc2fd741d59105eb6c99b9d0 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/20] 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 4be18158bdfcf..97d0768e590c0 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
}
}
More information about the Mlir-commits
mailing list