[Mlir-commits] [mlir] [MLIR][XeGPU] Support Layout propagation for interleave and deintereleave op (PR #194966)
Jianhui Li
llvmlistbot at llvm.org
Mon May 4 15:52:41 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/194966
>From e73d36e1f4f837b6a49768674c19cfb73f0c55a2 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 23 Apr 2026 23:06:03 +0000
Subject: [PATCH 1/5] [mlir][XeGPU] Add layout propagation for bitcast,
interleave, and deinterleave ops
This patch adds layout inference and propagation support for vector bitcast,
interleave, and deinterleave operations in the XeGPU dialect.
Layout inference functions:
- inferBitCastSourceLayout: Infers source layout by scaling innermost dimension
based on element type bitwidth ratios
- inferInterleaveSourceLayout: Infers source layout for interleave (halves
innermost dim)
- inferDeinterleaveSourceLayout: Infers source layout for deinterleave (doubles
innermost dim)
Layout setup functions:
- setupBitCastResultLayout: Sets up result layout for bitcast operations
- setupInterleaveResultLayout: Ensures source layout can be safely derived after
interleave by adjusting laneData/instData to be divisible by 2
These functions enable proper layout tracking through vector transformation
operations during workgroup-to-subgroup distribution.
Co-Authored-By: Claude Sonnet 4.5 <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 18 +
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 376 ++++++++++++------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 95 ++++-
3 files changed, 376 insertions(+), 113 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 83eb939cf1bec..76ca3bb225792 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -103,6 +103,16 @@ DistributeLayoutAttr inferBitCastSourceLayout(DistributeLayoutAttr resLayout,
int resElemTyBitWidth,
int srcElemTyBitWidth);
+/// Infers the source layout attribute for an interleave operation given the
+/// result layout attribute. Interleave doubles the innermost dimension size.
+DistributeLayoutAttr
+inferInterleaveSourceLayout(DistributeLayoutAttr resLayout);
+
+/// Infers the source layout attribute for a deinterleave operation given the
+/// result layout attribute. Deinterleave halves the innermost dimension size.
+DistributeLayoutAttr
+inferDeinterleaveSourceLayout(DistributeLayoutAttr resLayout);
+
/// Infers the source layout attribute for a shape cast operation given the
/// result layout attribute, result shape, and source shape.
DistributeLayoutAttr inferShapeCastSourceLayout(DistributeLayoutAttr resLayout,
@@ -161,6 +171,14 @@ DistributeLayoutAttr setupBitCastResultLayout(
LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy,
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
+/// Sets up the result layout for an interleave operation to ensure the source
+/// layout can be safely derived. Interleave doubles the innermost dimension,
+/// so the result layout must ensure that laneData is at least 2 (or a multiple
+/// of 2), and instData must be divisible by innermostDimLaneLayout * 2.
+DistributeLayoutAttr setupInterleaveResultLayout(
+ LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy,
+ DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
+
/// Sets up the result layout for an insert strided slice operation.
/// Creates a result layout based on the specified layout kind (InstData or
/// Lane).
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 7d48315eec6ff..037aa736c69c8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -458,6 +458,85 @@ xegpu::inferBitCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return finalSrcLayout;
}
+/// Infers the source layout attribute for an interleave operation given the
+/// result layout attribute. Interleave doubles the size of the innermost
+/// dimension, so the layout inference is similar to bitcast where the source
+/// element type is larger than the result element type (ratio = 2).
+xegpu::DistributeLayoutAttr
+xegpu::inferInterleaveSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
+
+ SmallVector<int64_t> sgData = resLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = resLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneData = resLayout.getEffectiveLaneDataAsInt();
+ size_t sgDataSize = sgData.size();
+ size_t instDataSize = instData.size();
+ size_t laneDataSize = laneData.size();
+ int64_t sgDataValue = -1;
+ int64_t instDataValue = -1;
+ int64_t laneDataValue = -1;
+ int64_t dim = resLayout.getRank() - 1;
+
+ // Interleave doubles the innermost dimension, so we need to halve the
+ // layout values (similar to bitcast with ratio = 2)
+ int ratio = 2;
+ if (sgDataSize) {
+ assert((sgData.back() % ratio) == 0 &&
+ "sgData not divisible by interleave ratio");
+ sgDataValue = sgData.back() / ratio;
+ }
+ if (instDataSize) {
+ assert((instData.back() % ratio) == 0 &&
+ "instData not divisible by interleave ratio");
+ instDataValue = instData.back() / ratio;
+ }
+ if (laneDataSize) {
+ assert((laneData.back() % ratio) == 0 &&
+ "laneData not divisible by interleave ratio");
+ laneDataValue = laneData.back() / ratio;
+ }
+
+ xegpu::DistributeLayoutAttr finalSrcLayout;
+ finalSrcLayout =
+ resLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
+
+ return finalSrcLayout;
+}
+
+/// Infers the source layout attribute for a deinterleave operation given the
+/// result layout attribute. Deinterleave halves the size of the innermost
+/// dimension, so the layout inference is similar to bitcast where the source
+/// element type is smaller than the result element type (ratio = 2).
+xegpu::DistributeLayoutAttr
+xegpu::inferDeinterleaveSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
+
+ SmallVector<int64_t> sgData = resLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = resLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneData = resLayout.getEffectiveLaneDataAsInt();
+ size_t sgDataSize = sgData.size();
+ size_t instDataSize = instData.size();
+ size_t laneDataSize = laneData.size();
+ int64_t sgDataValue = -1;
+ int64_t instDataValue = -1;
+ int64_t laneDataValue = -1;
+ int64_t dim = resLayout.getRank() - 1;
+
+ // Deinterleave halves the innermost dimension, so we need to double the
+ // layout values (similar to bitcast with ratio = 2)
+ int ratio = 2;
+ if (sgDataSize)
+ sgDataValue = sgData.back() * ratio;
+ if (instDataSize)
+ instDataValue = instData.back() * ratio;
+ if (laneDataSize)
+ laneDataValue = laneData.back() * ratio;
+
+ xegpu::DistributeLayoutAttr finalSrcLayout;
+ finalSrcLayout =
+ resLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
+
+ return finalSrcLayout;
+}
+
/// Infers the source layout attribute for an insert strided slice operation
/// given the result layout attribute, result shape, and source shape. Removes
/// leading dimensions from the result layout to match the source shape size.
@@ -877,6 +956,71 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
return consumerLayout;
}
+/// Sets up the result layout for an interleave operation to ensure the source
+/// layout can be safely derived. Interleave doubles the innermost dimension,
+/// so the result layout must ensure that laneData is at least 2 (or a multiple
+/// of 2), and instData must be divisible by innermostDimLaneLayout * 2.
+///
+/// Example:
+/// Interleave: vector<128x256xf4> -> vector<128x512xf4>
+/// Consumer layout: laneLayout=[1, 16], laneData=[1, 4], instData=[1, 64]
+/// Result layout adjustment to ensure source can be safely inferred:
+/// - laneData must be >= 2 and multiple of 2 (so source = laneData/2 is
+/// valid)
+/// - instData must be divisible by (16 * 2 = 32) (so source = instData/2 is
+/// valid)
+/// - Adjusted instData: ensure (instData % 32 == 0)
+///
+xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
+
+ ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+
+ assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
+ "consumer layout rank must match source shape rank");
+ size_t dim = srcShape.size() - 1;
+ int64_t sgDataValue = -1;
+ int64_t instDataValue = -1;
+ int64_t laneDataValue = -1;
+ const int subgroupSize = uArch->getSubgroupSize();
+
+ // Interleave doubles the innermost dimension (ratio = 2)
+ int ratio = 2;
+ int innermostDimLaneLayout = subgroupSize;
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ sgDataValue = sgData[dim];
+ // Ensure sgDataValue is divisible by ratio so source sgData can be inferred
+ while ((sgDataValue <= srcShape[dim]) && (sgDataValue % ratio != 0))
+ sgDataValue *= 2;
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ instDataValue = instData[dim];
+ // Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
+ // ratio) when inferring the source layout
+ while ((instDataValue <= srcShape[dim]) &&
+ (instDataValue % (innermostDimLaneLayout * ratio) != 0))
+ instDataValue *= 2;
+ assert((srcShape[dim] % instDataValue) == 0 &&
+ "srcShape, instData, and laneLayout for innermost must be 2^n!");
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ laneDataValue = laneData[dim];
+ // Ensure laneDataValue is at least 2 and divisible by ratio
+ // so that source laneData = laneDataValue/2 is valid
+ while ((laneDataValue <= srcShape[dim]) && (laneDataValue % ratio != 0))
+ laneDataValue *= 2;
+ }
+
+ xegpu::DistributeLayoutAttr resLayout;
+ resLayout =
+ consumerLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
+
+ return resLayout;
+}
+
/// Sets up the result layout for an insert strided slice operation.
/// Creates a result layout based on the specified layout kind (InstData or
/// Lane).
@@ -1179,6 +1323,122 @@ getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
return candidates;
}
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D.
+static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>>
+getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
+ const xegpu::uArch::uArch *uArch) {
+ const int subgroupSize = uArch->getSubgroupSize();
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+
+ const unsigned dataALen = aTy.getShape().front();
+ auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
+ const int maxALen =
+ xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
+
+ const unsigned dataBLen = bTy.getShape().back();
+ auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
+ const int maxBLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
+
+ auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
+ const int maxCLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
+ if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+ return std::nullopt;
+
+ SmallVector<int64_t> instDataA(aTy.getRank(), 1);
+ instDataA[aTy.getRank() - 2] = maxALen;
+ instDataA[aTy.getRank() - 1] = subgroupSize;
+ SmallVector<int64_t> instDataB(bTy.getRank(), 1);
+ instDataB[bTy.getRank() - 2] = subgroupSize;
+ instDataB[bTy.getRank() - 1] = maxBLen;
+ SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
+ instDataCD[cdTy.getRank() - 2] = maxALen;
+ instDataCD[cdTy.getRank() - 1] = maxCLen;
+ return std::make_tuple(instDataA, instDataB, instDataCD);
+}
+
+/// Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.
+/// Returns the three layouts if successful, nullopt otherwise.
+static std::optional<std::tuple<xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
+ VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
+ if (!instDataVecs)
+ return std::nullopt;
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+ assert(instDataA.size() == 2 && instDataB.size() == 2 &&
+ instDataCD.size() == 2 &&
+ "Sg layout creation expects valid 2D inst data");
+
+ std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
+ if (consumerLayout && consumerLayout.isForWorkgroup()) {
+ SmallVector<int64_t> sgLayoutD =
+ consumerLayout.getEffectiveSgLayoutAsInt();
+ consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
+ }
+
+ // Get all valid layouts for A, B and C/D operands
+ auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
+ if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
+ return std::nullopt;
+
+ // Pick the best subgroup layout
+ llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
+ llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
+ layoutsCD.end());
+ std::optional<LayoutRepresentation> bestPick;
+ for (auto &sgLayout : layoutsB) {
+ if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
+ if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
+ bestPick = sgLayout;
+ break;
+ }
+ if (!bestPick)
+ bestPick = sgLayout;
+ }
+ }
+ if (!bestPick)
+ return std::nullopt;
+
+ SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
+ static_cast<int>(bestPick->second)};
+ SmallVector<int> sgDataA = {
+ static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(aTy.getShape()[1] / sgLayout[1])};
+ SmallVector<int> sgDataB = {
+ static_cast<int>(bTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
+ SmallVector<int> sgDataCD = {
+ static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
+
+ auto dpasALayout = xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgDataA), nullptr, nullptr, nullptr,
+ nullptr);
+ auto dpasBLayout = xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgDataB), nullptr, nullptr, nullptr,
+ nullptr);
+ auto dpasCDLayout = xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgDataCD), nullptr, nullptr, nullptr,
+ nullptr);
+
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
+}
+
/// Sets up the anchor layouts for dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout
/// creation.
@@ -1194,122 +1454,13 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
- auto getInstDataVectors = [&]()
- -> std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
- SmallVector<int64_t>>> {
- const int subgroupSize = uArch->getSubgroupSize();
- const unsigned dataALen = aTy.getShape().front();
- auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
- const int maxALen =
- xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
-
- const unsigned dataBLen = bTy.getShape().back();
- auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
- const int maxBLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
-
- auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
- const int maxCLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
- if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
- return std::nullopt;
-
- SmallVector<int64_t> instDataA(aTy.getRank(), 1);
- instDataA[aTy.getRank() - 2] = maxALen;
- instDataA[aTy.getRank() - 1] = subgroupSize;
- SmallVector<int64_t> instDataB(bTy.getRank(), 1);
- instDataB[bTy.getRank() - 2] = subgroupSize;
- instDataB[bTy.getRank() - 1] = maxBLen;
- SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
- instDataCD[cdTy.getRank() - 2] = maxALen;
- instDataCD[cdTy.getRank() - 1] = maxCLen;
- return std::make_tuple(instDataA, instDataB, instDataCD);
- };
-
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
- auto instDataVecs = getInstDataVectors();
- if (!instDataVecs)
- return std::nullopt;
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- assert(instDataA.size() == 2 && instDataB.size() == 2 &&
- instDataCD.size() == 2 &&
- "Sg layout creation expects valid 2D inst data");
-
- std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
- if (consumerLayout && consumerLayout.isForWorkgroup()) {
- SmallVector<int64_t> sgLayoutD =
- consumerLayout.getEffectiveSgLayoutAsInt();
- consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
- }
-
- // Step 1. Get all valid layouts for A, B and C/D operands.
- // Order them from most balanced to least balanced.
- auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
- auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
- auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
- if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
- return std::nullopt;
-
- // Step 2. If the consumer layout can be reused for all operands, that
- // layout is chosen. Otherwise, pick the most balanced subgroup layout
- // that is valid for A, B and C (if present) operands
- llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
- llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
- layoutsCD.end());
- std::optional<LayoutRepresentation> bestPick;
- for (auto &sgLayout : layoutsB) {
- if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
- // Is in (A and B and CD) and matches consumer -> best pick
- if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
- bestPick = sgLayout;
- break;
- }
- // Is in (A and B and CD) layoutsB is ordered from most
- // balanced to least. So the first one we see is the most balanced
- // one, remember it and later only update if there is one that matches
- // the consumer.
- if (!bestPick)
- bestPick = sgLayout;
- }
- }
- // Step 3. If there is no subgroup layout compatible with A, B and C (if
- // present) operands, we fail.
- if (!bestPick)
- return std::nullopt;
- SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
- static_cast<int>(bestPick->second)};
- SmallVector<int> sgDataA = {
- static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(aTy.getShape()[1] / sgLayout[1])};
- SmallVector<int> sgDataB = {
- static_cast<int>(bTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
- SmallVector<int> sgDataCD = {
- static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
-
- auto dpasALayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataA),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/nullptr);
-
- auto dpasBLayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataB),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/nullptr);
-
- auto dpasCDLayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataCD),
- /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
- /*lane_data =*/nullptr, /*order =*/nullptr);
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
+ return getupDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
+ numSg, uArch);
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto instDataVecs = getInstDataVectors();
+ auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
if (!instDataVecs)
return std::nullopt;
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
@@ -1332,6 +1483,7 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
}
+
xegpu::DistributeLayoutAttr
xegpu::inferSourceLayoutFromResult(OpOperand &operand,
xegpu::DistributeLayoutAttr resLayout) {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 43998ed41f7aa..dddd8f49be9a5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -316,6 +316,7 @@ class LayoutInfoPropagation
void visitDpasOp(xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
+
void visitStoreNdOp(xegpu::StoreNdOp store,
ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
@@ -340,6 +341,14 @@ class LayoutInfoPropagation
ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
+ void visitVectorInterleaveOp(vector::InterleaveOp interleave,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results);
+
+ void visitVectorDeinterleaveOp(vector::DeinterleaveOp deinterleave,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results);
+
void visitUpdateNdOffsetOp(xegpu::UpdateNdOffsetOp updateNdOffset,
ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
@@ -450,6 +459,12 @@ LogicalResult LayoutInfoPropagation::visitOperation(
.Case([&](vector::BitCastOp bitcastOp) {
visitVectorBitcastOp(bitcastOp, operands, results);
})
+ .Case([&](vector::InterleaveOp interleaveOp) {
+ visitVectorInterleaveOp(interleaveOp, operands, results);
+ })
+ .Case([&](vector::DeinterleaveOp deinterleaveOp) {
+ visitVectorDeinterleaveOp(deinterleaveOp, operands, results);
+ })
.Case([&](vector::MultiDimReductionOp reductionOp) {
visitVectorMultiReductionOp(reductionOp, operands, results);
})
@@ -816,7 +831,6 @@ void LayoutInfoPropagation::visitDpasOp(
propagateIfChanged(operands[2], operands[2]->meet(dpasCDLayout));
}
-/// Set the layout for the value and tensor descriptor operands in StoreNdOp.
void LayoutInfoPropagation::visitStoreNdOp(
xegpu::StoreNdOp store, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
@@ -947,10 +961,24 @@ void LayoutInfoPropagation::visitTransposeOp(
LayoutInfo resultLayout = results[0]->getValue();
if (!resultLayout.isAssigned())
return;
+
+ llvm::dbgs() << "[DEBUG visitTransposeOp] transpose op: " << transpose << "\n";
+ llvm::dbgs() << "[DEBUG visitTransposeOp] resultLayout (consumer): " << resultLayout.get() << "\n";
+ llvm::dbgs() << "[DEBUG visitTransposeOp] permutation: [";
+ auto perm = transpose.getPermutation();
+ for (size_t i = 0; i < perm.size(); ++i) {
+ if (i > 0) llvm::dbgs() << ", ";
+ llvm::dbgs() << perm[i];
+ }
+ llvm::dbgs() << "]\n";
+
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resultLayout.get());
auto srcLayoutAttr = xegpu::inferTransposeSourceLayout(
consumerLayoutAttr, transpose.getPermutation());
+
+ llvm::dbgs() << "[DEBUG visitTransposeOp] srcLayoutAttr (propagated to operand): " << srcLayoutAttr << "\n";
+
// Propagate the new layout to the vector operand.
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
@@ -988,6 +1016,71 @@ void LayoutInfoPropagation::visitVectorBitcastOp(
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
+/// For vector::InterleaveOp, the result has double the innermost dimension size
+/// compared to each source operand. The layout is propagated from result to
+/// sources, adjusting for the 2x size increase.
+void LayoutInfoPropagation::visitVectorInterleaveOp(
+ vector::InterleaveOp interleave, ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) {
+ // Need the layout of interleave result to propagate to the operands.
+ LayoutInfo resLayoutInfo = results[0]->getValue();
+ if (!resLayoutInfo.isAssigned())
+ return;
+
+ auto srcVecType = interleave.getSourceVectorType();
+ auto resVecType = interleave.getResultVectorType();
+
+ auto consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
+ const uArch *uArch = getUArch(xegpu::getChipStr(interleave).value_or(""));
+ if (!uArch)
+ return;
+
+ // Setup the result layout to ensure the source layout can be safely derived
+ auto requiredResLayoutAttr = setupInterleaveResultLayout(
+ layoutKind, srcVecType, resVecType, consumerLayoutAttr, uArch);
+
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] interleave op: " << interleave << "\n";
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] srcVecType: " << srcVecType << "\n";
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] resVecType: " << resVecType << "\n";
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] consumerLayoutAttr: " << consumerLayoutAttr << "\n";
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] requiredResLayoutAttr (before inferInterleaveSourceLayout): " << requiredResLayoutAttr << "\n";
+
+ xegpu::setTemporaryLayout(interleave->getResult(0), requiredResLayoutAttr);
+
+ // Derive the source layout from the result layout (halve the innermost dim)
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] About to call inferInterleaveSourceLayout...\n";
+ auto srcLayoutAttr = xegpu::inferInterleaveSourceLayout(requiredResLayoutAttr);
+ llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] After inferInterleaveSourceLayout, srcLayoutAttr: " << srcLayoutAttr << "\n";
+
+ // Both operands (lhs and rhs) get the same source layout
+ propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[1], operands[1]->meet(LayoutInfo(srcLayoutAttr)));
+}
+
+/// For vector::DeinterleaveOp, the source has double the innermost dimension
+/// size compared to each result. The layout is propagated from results to
+/// source, adjusting for the 2x size decrease in results.
+void LayoutInfoPropagation::visitVectorDeinterleaveOp(
+ vector::DeinterleaveOp deinterleave, ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) {
+ // Need the layout of deinterleave results to propagate to the operand.
+ // Use the first result's layout (both results should have the same layout)
+ LayoutInfo resLayoutInfo = results[0]->getValue();
+ if (!resLayoutInfo.isAssigned())
+ return;
+
+ auto consumerLayoutAttr =
+ dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
+
+ // Derive the source layout from the result layout (double the innermost dim)
+ // No setup function needed - just infer directly
+ auto srcLayoutAttr =
+ xegpu::inferDeinterleaveSourceLayout(consumerLayoutAttr);
+
+ propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+}
+
void LayoutInfoPropagation::visitInsertStridedSliceOp(
vector::InsertStridedSliceOp insertStridedSlice,
ArrayRef<LayoutInfoLattice *> operands,
>From 999b3c6963702b4ca4723fa159a626eb423efe4b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 29 Apr 2026 19:04:56 +0000
Subject: [PATCH 2/5] remove related code
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 253 +++++++++---------
1 file changed, 133 insertions(+), 120 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 037aa736c69c8..19e544000b11a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1323,122 +1323,6 @@ getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
return candidates;
}
-/// Helper function to compute inst_data vectors for DPAS operands A, B, and
-/// C/D.
-static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
- SmallVector<int64_t>>>
-getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
- const xegpu::uArch::uArch *uArch) {
- const int subgroupSize = uArch->getSubgroupSize();
- const auto *uArchInstruction =
- dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
- xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
-
- const unsigned dataALen = aTy.getShape().front();
- auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
- const int maxALen =
- xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
-
- const unsigned dataBLen = bTy.getShape().back();
- auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
- const int maxBLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
-
- auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
- const int maxCLen =
- xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
- if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
- return std::nullopt;
-
- SmallVector<int64_t> instDataA(aTy.getRank(), 1);
- instDataA[aTy.getRank() - 2] = maxALen;
- instDataA[aTy.getRank() - 1] = subgroupSize;
- SmallVector<int64_t> instDataB(bTy.getRank(), 1);
- instDataB[bTy.getRank() - 2] = subgroupSize;
- instDataB[bTy.getRank() - 1] = maxBLen;
- SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
- instDataCD[cdTy.getRank() - 2] = maxALen;
- instDataCD[cdTy.getRank() - 1] = maxCLen;
- return std::make_tuple(instDataA, instDataB, instDataCD);
-}
-
-/// Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.
-/// Returns the three layouts if successful, nullopt otherwise.
-static std::optional<std::tuple<xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr,
- xegpu::DistributeLayoutAttr>>
-getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
- VectorType bTy, VectorType cdTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
- auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
- if (!instDataVecs)
- return std::nullopt;
- auto [instDataA, instDataB, instDataCD] = *instDataVecs;
- assert(instDataA.size() == 2 && instDataB.size() == 2 &&
- instDataCD.size() == 2 &&
- "Sg layout creation expects valid 2D inst data");
-
- std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
- if (consumerLayout && consumerLayout.isForWorkgroup()) {
- SmallVector<int64_t> sgLayoutD =
- consumerLayout.getEffectiveSgLayoutAsInt();
- consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
- }
-
- // Get all valid layouts for A, B and C/D operands
- auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
- auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
- auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
- if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
- return std::nullopt;
-
- // Pick the best subgroup layout
- llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
- llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
- layoutsCD.end());
- std::optional<LayoutRepresentation> bestPick;
- for (auto &sgLayout : layoutsB) {
- if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
- if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
- bestPick = sgLayout;
- break;
- }
- if (!bestPick)
- bestPick = sgLayout;
- }
- }
- if (!bestPick)
- return std::nullopt;
-
- SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
- static_cast<int>(bestPick->second)};
- SmallVector<int> sgDataA = {
- static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(aTy.getShape()[1] / sgLayout[1])};
- SmallVector<int> sgDataB = {
- static_cast<int>(bTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
- SmallVector<int> sgDataCD = {
- static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
- static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
-
- auto dpasALayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataA), nullptr, nullptr, nullptr,
- nullptr);
- auto dpasBLayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataB), nullptr, nullptr, nullptr,
- nullptr);
- auto dpasCDLayout = xegpu::LayoutAttr::get(
- context, DenseI32ArrayAttr::get(context, sgLayout),
- DenseI32ArrayAttr::get(context, sgDataCD), nullptr, nullptr, nullptr,
- nullptr);
-
- return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
-}
-
/// Sets up the anchor layouts for dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout
/// creation.
@@ -1454,13 +1338,122 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ auto getInstDataVectors = [&]()
+ -> std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>> {
+ const int subgroupSize = uArch->getSubgroupSize();
+ const unsigned dataALen = aTy.getShape().front();
+ auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
+ const int maxALen =
+ xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
+
+ const unsigned dataBLen = bTy.getShape().back();
+ auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
+ const int maxBLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
+
+ auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
+ const int maxCLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
+ if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+ return std::nullopt;
+
+ SmallVector<int64_t> instDataA(aTy.getRank(), 1);
+ instDataA[aTy.getRank() - 2] = maxALen;
+ instDataA[aTy.getRank() - 1] = subgroupSize;
+ SmallVector<int64_t> instDataB(bTy.getRank(), 1);
+ instDataB[bTy.getRank() - 2] = subgroupSize;
+ instDataB[bTy.getRank() - 1] = maxBLen;
+ SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
+ instDataCD[cdTy.getRank() - 2] = maxALen;
+ instDataCD[cdTy.getRank() - 1] = maxCLen;
+ return std::make_tuple(instDataA, instDataB, instDataCD);
+ };
+
if (layoutKind == xegpu::LayoutKind::Subgroup) {
assert(numSg > 0 &&
"Number of subgroups must be provided for sg layout creation.");
- return getupDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout,
- numSg, uArch);
+ auto instDataVecs = getInstDataVectors();
+ if (!instDataVecs)
+ return std::nullopt;
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+ assert(instDataA.size() == 2 && instDataB.size() == 2 &&
+ instDataCD.size() == 2 &&
+ "Sg layout creation expects valid 2D inst data");
+
+ std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
+ if (consumerLayout && consumerLayout.isForWorkgroup()) {
+ SmallVector<int64_t> sgLayoutD =
+ consumerLayout.getEffectiveSgLayoutAsInt();
+ consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
+ }
+
+ // Step 1. Get all valid layouts for A, B and C/D operands.
+ // Order them from most balanced to least balanced.
+ auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
+ if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
+ return std::nullopt;
+
+ // Step 2. If the consumer layout can be reused for all operands, that
+ // layout is chosen. Otherwise, pick the most balanced subgroup layout
+ // that is valid for A, B and C (if present) operands
+ llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
+ llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
+ layoutsCD.end());
+ std::optional<LayoutRepresentation> bestPick;
+ for (auto &sgLayout : layoutsB) {
+ if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
+ // Is in (A and B and CD) and matches consumer -> best pick
+ if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
+ bestPick = sgLayout;
+ break;
+ }
+ // Is in (A and B and CD) layoutsB is ordered from most
+ // balanced to least. So the first one we see is the most balanced
+ // one, remember it and later only update if there is one that matches
+ // the consumer.
+ if (!bestPick)
+ bestPick = sgLayout;
+ }
+ }
+ // Step 3. If there is no subgroup layout compatible with A, B and C (if
+ // present) operands, we fail.
+ if (!bestPick)
+ return std::nullopt;
+ SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
+ static_cast<int>(bestPick->second)};
+ SmallVector<int> sgDataA = {
+ static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(aTy.getShape()[1] / sgLayout[1])};
+ SmallVector<int> sgDataB = {
+ static_cast<int>(bTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
+ SmallVector<int> sgDataCD = {
+ static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
+
+ auto dpasALayout = xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgDataA),
+ /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
+ /*lane_data =*/nullptr, /*order =*/nullptr);
+
+ auto dpasBLayout = xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgDataB),
+ /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
+ /*lane_data =*/nullptr, /*order =*/nullptr);
+
+ auto dpasCDLayout = xegpu::LayoutAttr::get(
+ context, DenseI32ArrayAttr::get(context, sgLayout),
+ DenseI32ArrayAttr::get(context, sgDataCD),
+ /*inst_data =*/nullptr, /*lane_layout =*/nullptr,
+ /*lane_data =*/nullptr, /*order =*/nullptr);
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout);
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
+ auto instDataVecs = getInstDataVectors();
if (!instDataVecs)
return std::nullopt;
auto [instDataA, instDataB, instDataCD] = *instDataVecs;
@@ -1483,7 +1476,6 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
}
-
xegpu::DistributeLayoutAttr
xegpu::inferSourceLayoutFromResult(OpOperand &operand,
xegpu::DistributeLayoutAttr resLayout) {
@@ -1555,6 +1547,27 @@ xegpu::inferSourceLayoutFromResult(OpOperand &operand,
transpose.getPermutation());
}
+ // For vector::BitCastOp, infer source layout from result layout using
+ // element type bitwidths.
+ if (auto bitcast = dyn_cast<vector::BitCastOp>(op)) {
+ int resElemBitWidth =
+ bitcast.getResultVectorType().getElementType().getIntOrFloatBitWidth();
+ int srcElemBitWidth =
+ bitcast.getSourceVectorType().getElementType().getIntOrFloatBitWidth();
+ return xegpu::inferBitCastSourceLayout(resLayout, resElemBitWidth,
+ srcElemBitWidth);
+ }
+
+ // for vector::interleave
+ if (auto interleave = dyn_cast<vector::InterleaveOp>(op)) {
+ return xegpu::inferInterleaveSourceLayout(resLayout);
+ }
+
+ // for vector::deinterleave
+ if (auto deinterleave = dyn_cast<vector::DeinterleaveOp>(op)) {
+ return xegpu::inferDeinterleaveSourceLayout(resLayout);
+ }
+
// For vector::ExtractStridedSliceOp, simply return result layout
if (dyn_cast<vector::ExtractStridedSliceOp>(op))
return resLayout;
>From 1bb951112a4cf7894e5ad662221b502188b72a2d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 29 Apr 2026 19:05:52 +0000
Subject: [PATCH 3/5] format
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 19e544000b11a..e56f17867cb81 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1559,13 +1559,13 @@ xegpu::inferSourceLayoutFromResult(OpOperand &operand,
}
// for vector::interleave
- if (auto interleave = dyn_cast<vector::InterleaveOp>(op)) {
- return xegpu::inferInterleaveSourceLayout(resLayout);
+ if (auto interleave = dyn_cast<vector::InterleaveOp>(op)) {
+ return xegpu::inferInterleaveSourceLayout(resLayout);
}
// for vector::deinterleave
if (auto deinterleave = dyn_cast<vector::DeinterleaveOp>(op)) {
- return xegpu::inferDeinterleaveSourceLayout(resLayout);
+ return xegpu::inferDeinterleaveSourceLayout(resLayout);
}
// For vector::ExtractStridedSliceOp, simply return result layout
>From fd2660636ee6995ae6a7ef9b253e734734a99bc1 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 29 Apr 2026 21:18:15 +0000
Subject: [PATCH 4/5] add tests
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 34 ---------------
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 42 +++++++++++++++++++
2 files changed, 42 insertions(+), 34 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 91554029769f9..a9236187eb77b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1079,28 +1079,11 @@ void LayoutInfoPropagation::visitTransposeOp(
if (!resultLayout.isAssigned())
return;
- llvm::dbgs() << "[DEBUG visitTransposeOp] transpose op: " << transpose
- << "\n";
- llvm::dbgs() << "[DEBUG visitTransposeOp] resultLayout (consumer): "
- << resultLayout.get() << "\n";
- llvm::dbgs() << "[DEBUG visitTransposeOp] permutation: [";
- auto perm = transpose.getPermutation();
- for (size_t i = 0; i < perm.size(); ++i) {
- if (i > 0)
- llvm::dbgs() << ", ";
- llvm::dbgs() << perm[i];
- }
- llvm::dbgs() << "]\n";
-
auto consumerLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resultLayout.get());
auto srcLayoutAttr = xegpu::inferTransposeSourceLayout(
consumerLayoutAttr, transpose.getPermutation());
- llvm::dbgs()
- << "[DEBUG visitTransposeOp] srcLayoutAttr (propagated to operand): "
- << srcLayoutAttr << "\n";
-
// Propagate the new layout to the vector operand.
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
@@ -1162,28 +1145,11 @@ void LayoutInfoPropagation::visitVectorInterleaveOp(
auto requiredResLayoutAttr = setupInterleaveResultLayout(
layoutKind, srcVecType, resVecType, consumerLayoutAttr, uArch);
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] interleave op: "
- << interleave << "\n";
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] srcVecType: " << srcVecType
- << "\n";
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] resVecType: " << resVecType
- << "\n";
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] consumerLayoutAttr: "
- << consumerLayoutAttr << "\n";
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] requiredResLayoutAttr "
- "(before inferInterleaveSourceLayout): "
- << requiredResLayoutAttr << "\n";
-
xegpu::setTemporaryLayout(interleave->getResult(0), requiredResLayoutAttr);
// Derive the source layout from the result layout (halve the innermost dim)
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] About to call "
- "inferInterleaveSourceLayout...\n";
auto srcLayoutAttr =
xegpu::inferInterleaveSourceLayout(requiredResLayoutAttr);
- llvm::dbgs() << "[DEBUG visitVectorInterleaveOp] After "
- "inferInterleaveSourceLayout, srcLayoutAttr: "
- << srcLayoutAttr << "\n";
// Both operands (lhs and rhs) get the same source layout
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 72d066d516540..fe4637170bd15 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -1064,3 +1064,45 @@ func.func @dpas_mx_fp4(%arg0: memref<8x64xf4E2M1FN>, %arg1: memref<64x16xf4E2M1F
return
}
}
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @vector_interleave_f16(
+// CHECK: %[[LOAD1:.*]] = xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
+// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+// CHECK: %[[LOAD2:.*]] = xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
+// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
+// CHECK-NEXT: %{{.*}} = vector.interleave %[[LOAD1]], %[[LOAD2]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}
+// CHECK-SAME: vector<8x16xf16> -> vector<8x32xf16>
+func.func @vector_interleave_f16(%arg0: memref<8x16xf16>, %arg1: memref<8x16xf16>, %arg2: memref<8x32xf16>) {
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %1 = xegpu.create_nd_tdesc %arg1 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %2 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %3 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %4 = vector.interleave %2, %3 : vector<8x16xf16> -> vector<8x32xf16>
+ %5 = xegpu.create_nd_tdesc %arg2 : memref<8x32xf16> -> !xegpu.tensor_desc<8x32xf16>
+ xegpu.store_nd %4, %5[0, 0] : vector<8x32xf16>, !xegpu.tensor_desc<8x32xf16>
+ return
+}
+}
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @vector_deinterleave_f16(
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}>
+// CHECK-SAME: !xegpu.tensor_desc<8x32xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<8x32xf16>
+// CHECK-NEXT: %{{.*}}, %{{.*}} = vector.deinterleave %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_result_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+// CHECK-SAME: vector<8x32xf16> -> vector<8x16xf16>
+func.func @vector_deinterleave_f16(%arg0: memref<8x32xf16>, %arg1: memref<8x16xf16>, %arg2: memref<8x16xf16>) {
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<8x32xf16> -> !xegpu.tensor_desc<8x32xf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<8x32xf16> -> vector<8x32xf16>
+ %2:2 = vector.deinterleave %1 : vector<8x32xf16> -> vector<8x16xf16>
+ %3 = xegpu.create_nd_tdesc %arg1 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %4 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
+ xegpu.store_nd %2#0, %3[0, 0] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
+ xegpu.store_nd %2#1, %4[0, 0] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
+ return
+}
+}
>From 5f5008190b00a46fbd3667c666bb50448fe05e78 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 4 May 2026 21:37:04 +0000
Subject: [PATCH 5/5] address feedback
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 56 ++++++++-----------
1 file changed, 23 insertions(+), 33 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 1c9a81ac8c915..02aa831cb3b61 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -478,7 +478,7 @@ xegpu::inferInterleaveSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
// Interleave doubles the innermost dimension, so we need to halve the
// layout values (similar to bitcast with ratio = 2)
- int ratio = 2;
+ constexpr int ratio = 2;
if (sgDataSize) {
assert((sgData.back() % ratio) == 0 &&
"sgData not divisible by interleave ratio");
@@ -495,11 +495,7 @@ xegpu::inferInterleaveSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
laneDataValue = laneData.back() / ratio;
}
- xegpu::DistributeLayoutAttr finalSrcLayout;
- finalSrcLayout =
- resLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
-
- return finalSrcLayout;
+ return resLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
}
/// Infers the source layout attribute for a deinterleave operation given the
@@ -522,7 +518,7 @@ xegpu::inferDeinterleaveSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
// Deinterleave halves the innermost dimension, so we need to double the
// layout values (similar to bitcast with ratio = 2)
- int ratio = 2;
+ constexpr int ratio = 2;
if (sgDataSize)
sgDataValue = sgData.back() * ratio;
if (instDataSize)
@@ -530,11 +526,7 @@ xegpu::inferDeinterleaveSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
if (laneDataSize)
laneDataValue = laneData.back() * ratio;
- xegpu::DistributeLayoutAttr finalSrcLayout;
- finalSrcLayout =
- resLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
-
- return finalSrcLayout;
+ return resLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
}
/// Infers the source layout attribute for an insert strided slice operation
@@ -958,8 +950,8 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
/// Sets up the result layout for an interleave operation to ensure the source
/// layout can be safely derived. Interleave doubles the innermost dimension,
-/// so the result layout must ensure that laneData is at least 2 (or a multiple
-/// of 2), and instData must be divisible by innermostDimLaneLayout * 2.
+/// so the result layout must ensure that laneData is a multiple
+/// of 2, and instData must be divisible by innermostDimLaneLayout * 2.
///
/// Example:
/// Interleave: vector<128x256xf4> -> vector<128x512xf4>
@@ -982,43 +974,41 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
"consumer layout rank must match source shape rank");
- size_t dim = srcShape.size() - 1;
+ const size_t innerMostDim = srcShape.size() - 1;
int64_t sgDataValue = -1;
int64_t instDataValue = -1;
int64_t laneDataValue = -1;
- const int subgroupSize = uArch->getSubgroupSize();
// Interleave doubles the innermost dimension (ratio = 2)
- int ratio = 2;
- int innermostDimLaneLayout = subgroupSize;
+ constexpr int ratio = 2;
+ int innermostDimLaneLayout = uArch->getSubgroupSize();
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[dim];
+ sgDataValue = sgData[innerMostDim];
// Ensure sgDataValue is divisible by ratio so source sgData can be inferred
- while ((sgDataValue <= srcShape[dim]) && (sgDataValue % ratio != 0))
- sgDataValue *= 2;
+ while ((sgDataValue <= srcShape[innerMostDim]) &&
+ (sgDataValue % ratio != 0))
+ sgDataValue *= ratio;
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[dim];
+ instDataValue = instData[innerMostDim];
// Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
// ratio) when inferring the source layout
- while ((instDataValue <= srcShape[dim]) &&
+ while ((instDataValue <= srcShape[innerMostDim]) &&
(instDataValue % (innermostDimLaneLayout * ratio) != 0))
- instDataValue *= 2;
- assert((srcShape[dim] % instDataValue) == 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[dim];
+ laneDataValue = laneData[innerMostDim];
// Ensure laneDataValue is at least 2 and divisible by ratio
// so that source laneData = laneDataValue/2 is valid
- while ((laneDataValue <= srcShape[dim]) && (laneDataValue % ratio != 0))
- laneDataValue *= 2;
+ while ((laneDataValue <= srcShape[innerMostDim]) &&
+ (laneDataValue % ratio != 0))
+ laneDataValue *= ratio;
}
- xegpu::DistributeLayoutAttr resLayout;
- resLayout =
- consumerLayout.setDimData(dim, sgDataValue, instDataValue, laneDataValue);
-
- return resLayout;
+ return consumerLayout.setDimData(innerMostDim, sgDataValue, instDataValue,
+ laneDataValue);
}
/// Sets up the result layout for an insert strided slice operation.
More information about the Mlir-commits
mailing list