[Mlir-commits] [mlir] [MLIR][XeGPU] Enable WG-level mxfp GEMM via generalized shape_cast collapse inference (PR #201496)
Jianhui Li
llvmlistbot at llvm.org
Wed Jun 3 21:48:17 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/201496
>From 52420372ba5ffecedc0587332ead8c73c59fab34 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 4 Jun 2026 01:01:01 +0000
Subject: [PATCH 1/5] [MLIR][XeGPU] Generalize shape_cast layout collapse
inference
Replace the narrow matchCollapseToInnermostDim helper in
inferShapeCastSourceLayout with a general matchDimCollapse utility, and
extend the inference logic to handle arbitrary collapse patterns where
multiple consecutive src dims fold into one or more dst dims.
For each dst dim produced by a collapsed group of src dims:
- sg_layout distributes outer-to-inner so each subgroup owns a
contiguous run in the collapsed dst dim's row-major linearization;
the value spreads inward when a single dim cannot hold it.
- lane_layout distributes inner-to-outer so lanes vectorize along the
fastest-varying axis; the value spreads outward when needed.
- sg_data, inst_data, and lane_data fill innermost-first and spill
outward, capped per dim by srcShape[d] and any layout placed on d.
- order is propagated by walking dst order fastest-first and emitting
each group's src dims from innermost to outermost.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 9 +
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 293 ++++++++++++++----
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 53 ++++
3 files changed, 302 insertions(+), 53 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 1b594f17e15ec..57401b170326a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -234,6 +234,15 @@ bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
SmallVector<SmallVector<int64_t>> &splitDimGroups);
+// Checks if dst shape is a collapse of src shape where each dimension in dst is
+// produced by one or more consecutive dimensions in src whose product equals
+// the dst dimension. Populates collapseDims with groups of src indices that are
+// collapsed into each dst dimension. Leading or trailing unit dst dimensions
+// (with no backing src dim) result in empty groups. Example: src=[8,16,32],
+// dst=[1,4096] -> true, collapseDims=[[],[0,1,2]].
+bool matchDimCollapse(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
+ SmallVector<SmallVector<int64_t>> &collapseDims);
+
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 7384f6be8d051..b008a88ca9e8f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -425,10 +425,16 @@ xegpu::inferReductionSourceLayout(xegpu::DistributeLayoutAttr resLayout) {
/// Infers the source layout attribute for a transpose operation given the
/// result layout attribute and permutation.
+///
+/// vector.transpose semantics is `result[i] = source[permutation[i]]`, so
+/// `result_layout[i] = source_layout[permutation[i]]`. To recover the source
+/// layout from the result layout we must apply the inverse permutation.
xegpu::DistributeLayoutAttr
xegpu::inferTransposeSourceLayout(xegpu::DistributeLayoutAttr resLayout,
ArrayRef<int64_t> permutation) {
- return resLayout.transposeDims(permutation);
+ SmallVector<int64_t> inversePermutation =
+ invertPermutationVector(permutation);
+ return resLayout.transposeDims(inversePermutation);
}
/// Infers the source layout attribute for a bitcast operation given the
@@ -679,6 +685,37 @@ xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return resLayout;
}
+/// Walk `srcDims` in the requested direction and assign each src dim
+/// `min(remaining, perDimCap[d])`, spilling the leftover into the next dim.
+/// `srcDims` is always given outer-to-inner; `innerToOuter` selects the
+/// iteration direction. `perDimCap[d]` is the maximum value that may be
+/// placed on src dim `d`; callers compute it as either the full available
+/// extent (e.g. `srcShape[d]`) or the per-sg / per-lane share of that
+/// extent (e.g. `srcShape[d] / inferredSgLayout[d]`) when a layout has
+/// already been placed on that dim.
+static void distributeAcrossSrcDims(int64_t total, ArrayRef<int64_t> srcDims,
+ bool innerToOuter,
+ ArrayRef<int64_t> perDimCap,
+ SmallVectorImpl<int64_t> &out) {
+ int64_t remaining = total;
+ auto step = [&](int64_t d) {
+ if (remaining == 1)
+ return;
+ int64_t take = std::min(remaining, perDimCap[d]);
+ assert(take > 0 && "distribution must not be zero");
+ assert(remaining % take == 0 && "must divide evenly across dims");
+ out[d] = take;
+ remaining /= take;
+ };
+ if (innerToOuter)
+ for (int64_t d : llvm::reverse(srcDims))
+ step(d);
+ else
+ for (int64_t d : srcDims)
+ step(d);
+ assert(remaining == 1 && "total must fit within collapsed src dims");
+}
+
/// Infers the source layout attribute for a shape cast operation given the
/// result layout attribute, result shape, and source shape.
xegpu::DistributeLayoutAttr
@@ -719,68 +756,218 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return srcLayout;
}
- // Use case 3: Collaspse to innermost dim, for cross-sg reduction to SLM
- auto matchCollapseToInnermostDim = [&](ArrayRef<int64_t> src,
- ArrayRef<int64_t> dst) -> bool {
- // only one non-unit dim in dst which is the innermost dim
- if ((dst.size() != 2) && (dst.size() != 1))
- return false;
- int64_t srcSize = std::accumulate(src.begin(), src.end(), 1LL,
- std::multiplies<int64_t>());
- if (dst.size() == 1)
- return (dst[0] == srcSize);
- return (dst[0] == 1) && (dst[1] == srcSize);
- };
-
- if (matchCollapseToInnermostDim(srcShape, resShape)) {
+ // Use case 3: General dim collapse, for cross-sg reduction to SLM and other
+ // shape casts where consecutive src dims fold into a single dst dim.
+ // For each dst dim that is produced by collapsing >=2 src dims, distribute
+ // the consumer layout/data of that dst dim across the src dim group so that
+ // each sg/lane owns a contiguous run in the collapsed dst dim:
+ // - sg_layout: distribute starting from the OUTERMOST src dim of the
+ // group, spreading INWARD. Each dim takes min(remaining, srcShape[d]);
+ // leftover spills into the next inner dim. Subgroup partitioning along
+ // the slowest-varying axis keeps each sg's tile contiguous in the dst.
+ // - lane_layout: distribute starting from the OUTERMOST src dim of the
+ // group, spreading INWARD. Each dim takes min(remaining, srcShape[d]);
+ // leftover spills into the next inner dim. (Same direction as
+ // sg_layout; lane_data is the one that fills innermost-first to
+ // match the dst's row-major linearization.)
+ // - sg_data / inst_data / lane_data: fill from the innermost src dim
+ // outward, capped per dim by srcShape[d] (or srcShape[d]/sg_layout[d] /
+ // srcShape[d]/lane_layout[d] when a layout is placed on that dim).
+ // Examples:
+ // srcShape=[8, 16, 32], resShape=[1, 4096], inst_data=[1, 16]
+ // -> inferredInstData=[1, 1, 16]
+ // srcShape=[4, 8, 64], resShape=[2048], lane_layout=[16], lane_data=[2]
+ // -> inferredLaneLayout=[16, 1, 1], inferredLaneData=[1, 1, 2]
+ // srcShape=[8, 16, 32], resShape=[4096], sg_layout=[8], sg_data=[512]
+ // -> inferredSgLayout=[8, 1, 1] (outermost holds 8), inferredSgData=
+ // [1, 16, 32]
+ // srcShape=[2, 8, 32], resShape=[512], sg_layout=[16], sg_data=[32]
+ // -> outer dim 0 holds min(16, 2)=2, leftover 8 spills inward to dim 1:
+ // inferredSgLayout=[2, 8, 1]; inferredSgData fills innermost first:
+ // [1, 1, 32]
+ // srcShape=[64, 4, 2], resShape=[512], lane_layout=[16], lane_data=[2]
+ // -> outer dim 0 holds min(16, 64)=16: inferredLaneLayout=[16, 1, 1];
+ // inferredLaneData=[1, 1, 2] (innermost dim 2 fills first).
+ SmallVector<SmallVector<int64_t>> collapseDims;
+ if (xegpu::matchDimCollapse(srcShape, resShape, collapseDims)) {
int srcShapeSize = srcShape.size();
- int resShapeSize = resShape.size();
auto context = resLayout.getContext();
+ auto resSgLayout = resLayout.getEffectiveSgLayoutAsInt();
+ auto resSgData = resLayout.getEffectiveSgDataAsInt();
auto resInstData = resLayout.getEffectiveInstDataAsInt();
auto resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
auto resLaneData = resLayout.getEffectiveLaneDataAsInt();
- // 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.
- // Examples 1:
- // srcShape=[8, 16, 32], resShape=[1, 4096]
- // resInstData=[1, 16]
- // -> inferredInstData=[1, 1, min(16, 32)]=[1, 1, 16]
- // Examples 2:
- // srcShape=[4, 8, 64], resShape=[2048]
- // resLaneLayout=[16], resLaneData=[2]
- // -> inferredLaneLayout=[1, 1, 16]
- // -> inferredLaneData=[1, 1, min(2, 64/16)]=[1, 1, 2]
-
- 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");
+ SmallVector<int64_t> inferredSgLayout(
+ resSgLayout.empty() ? 0 : srcShapeSize, 1);
+ SmallVector<int64_t> inferredSgData(resSgData.empty() ? 0 : srcShapeSize,
+ 1);
+ SmallVector<int64_t> inferredInstData(
+ resInstData.empty() ? 0 : srcShapeSize, 1);
+ SmallVector<int64_t> inferredLaneLayout(
+ resLaneLayout.empty() ? 0 : srcShapeSize, 1);
+ SmallVector<int64_t> inferredLaneData(
+ resLaneData.empty() ? 0 : srcShapeSize, 1);
+
+ for (size_t dstIdx = 0; dstIdx < collapseDims.size(); ++dstIdx) {
+ ArrayRef<int64_t> srcDims = collapseDims[dstIdx];
+ if (srcDims.empty())
+ continue;
+
+ // Order matters: each *_data step depends on the matching *_layout
+ // having been computed first (the layout values are used as per-dim
+ // divisor caps when distributing the data). So we interleave:
+ // sg_layout -> sg_data -> lane_layout -> lane_data -> inst_data
+ // sg_data / lane_data / inst_data all fill innermost-first and spill
+ // outward; sg_data is capped per dim by srcShape[d]/sg_layout[d],
+ // and inst_data is seeded from lane_layout*lane_data (see below).
+ //
+ // After sg_data is distributed, `srcShape` is rebound to point at
+ // `inferredSgData` so that the subsequent lane_layout / lane_data /
+ // inst_data steps see the per-subgroup tile rather than the full
+ // source. This is what implicitly gives lane_data a cap of
+ // sgData[d]/lane_layout[d] when sg_data is available (and the
+ // original srcShape[d]/lane_layout[d] otherwise).
+ //
+ // When the consumer layout replicates the dst dim across
+ // subgroups/lanes (i.e. sg_layout[dstIdx] * sg_data[dstIdx] >
+ // resShape[dstIdx]), each subgroup/lane owns the full extent of the
+ // collapsed src dims, so the per-dim cap drops the layout divisor.
+
+ // Helper: build a per-src-dim cap by dividing each entry of `base` by
+ // the corresponding entry of `layout` (the consumer's already-placed
+ // partitioning on that dim). `layout` may be empty, meaning "no
+ // partition on this dim, so the cap is just the base extent".
+ auto perDimCap = [&](ArrayRef<int64_t> base,
+ ArrayRef<int64_t> layout) -> SmallVector<int64_t> {
+ SmallVector<int64_t> cap(base.begin(), base.end());
+ if (!layout.empty())
+ for (int64_t d : srcDims)
+ cap[d] /= layout[d];
+ return cap;
+ };
+
+ // sg_layout: outer-to-inner so the outermost src dim of the group fills
+ // first; leftover spreads inward when a single dim can't hold the value.
+ if (!resSgLayout.empty())
+ distributeAcrossSrcDims(resSgLayout[dstIdx], srcDims,
+ /*innerToOuter=*/false,
+ /*perDimCap=*/srcShape, inferredSgLayout);
+
+ // sg_data: innermost-first, capped per dim by srcShape[d]/sgLayout[d]
+ // unless the dst dim is sg-replicated (then each sg owns the full
+ // src extent).
+ if (!resSgData.empty()) {
+ bool sgReplicated =
+ !resSgLayout.empty() &&
+ resSgLayout[dstIdx] * resSgData[dstIdx] > resShape[dstIdx];
+ SmallVector<int64_t> cap =
+ sgReplicated
+ ? SmallVector<int64_t>(srcShape.begin(), srcShape.end())
+ : perDimCap(srcShape, inferredSgLayout);
+ distributeAcrossSrcDims(resSgData[dstIdx], srcDims,
+ /*innerToOuter=*/true, cap, inferredSgData);
}
- SmallVector<int> inferredInstData(srcShapeSize, 1);
- inferredInstData[srcShapeSize - 1] =
- std::min(resInstData[resShapeSize - 1], srcShape[srcShapeSize - 1]);
- return xegpu::LayoutAttr::get(context, inferredInstData);
- }
- if (resLaneLayout.size() != 0) {
- for (int i = 0; i < resShapeSize - 1; i++) {
- assert(resLaneData[i] == 1 &&
- "only innermost dim can have non-unit instData");
+ // Use a per-subgroup view for the remaining steps (lane_layout /
+ // lane_data / inst_data): they describe how a single subgroup's tile is
+ // partitioned across lanes, so their caps must be relative to
+ // inferredSgData rather than the full source. This is a local view; the
+ // function-scope `srcShape` must stay pointing at the full source so the
+ // next dst dim's sg_layout/sg_data steps are computed correctly.
+ ArrayRef<int64_t> laneSrcShape = resSgData.empty()
+ ? ArrayRef<int64_t>(srcShape)
+ : ArrayRef<int64_t>(inferredSgData);
+
+ // lane_layout: outer-to-inner so the outermost src dim of the group
+ // fills first; leftover spreads inward when a single dim is too small.
+ // Computed AFTER sg_data and BEFORE lane_data/inst_data so that the
+ // following data steps can use inferredLaneLayout as their cap/seed.
+ if (!resLaneLayout.empty())
+ distributeAcrossSrcDims(resLaneLayout[dstIdx], srcDims,
+ /*innerToOuter=*/false,
+ /*perDimCap=*/laneSrcShape, inferredLaneLayout);
+
+ // lane_data: innermost-first, capped per dim by
+ // (per-sg) srcShape[d] / inferredLaneLayout[d], unless the dst dim is
+ // lane-replicated (then each lane owns the full per-sg extent).
+ if (!resLaneData.empty()) {
+ bool laneReplicated =
+ !resLaneLayout.empty() &&
+ resLaneLayout[dstIdx] * resLaneData[dstIdx] > resShape[dstIdx];
+ SmallVector<int64_t> cap =
+ laneReplicated
+ ? SmallVector<int64_t>(laneSrcShape.begin(), laneSrcShape.end())
+ : perDimCap(laneSrcShape, inferredLaneLayout);
+ distributeAcrossSrcDims(resLaneData[dstIdx], srcDims,
+ /*innerToOuter=*/true, cap, inferredLaneData);
}
- assert(srcShape.back() % resLaneLayout.back() == 0 &&
- "source innermost dim must be >= result lane layout");
- SmallVector<int> inferredLaneLayout(srcShapeSize, 1);
- SmallVector<int> inferredLaneData(srcShapeSize, 1);
- inferredLaneLayout.back() = resLaneLayout.back();
- inferredLaneData.back() = std::min(
- resLaneData.back(), srcShape.back() / inferredLaneLayout.back());
- return xegpu::LayoutAttr::get(context, inferredLaneLayout,
- inferredLaneData);
+
+ // inst_data[d] must be a multiple of lane_layout[d] * lane_data[d] on
+ // each src dim of the collapsed group; otherwise the consumer's lane
+ // partitioning cannot be evenly mapped onto the source. Seed each src
+ // dim with the per-lane atomic unit (inferredLaneLayout[d] *
+ // inferredLaneData[d]) -- both already computed above -- then
+ // distribute the remaining factor innermost-first, capped per dim by
+ // srcShape[d] / inferredInstData[d].
+ if (!resInstData.empty()) {
+ // When the consumer layout has no lane_layout/lane_data, there is no
+ // per-lane atomic unit to seed; distribute inst_data the same way as
+ // sg_data/lane_data (innermost-first, no divisor cap).
+ if (resLaneLayout.empty() || resLaneData.empty()) {
+ distributeAcrossSrcDims(resInstData[dstIdx], srcDims,
+ /*innerToOuter=*/true,
+ /*perDimCap=*/laneSrcShape, inferredInstData);
+ } else {
+ int64_t laneAtom = resLaneLayout[dstIdx] * resLaneData[dstIdx];
+ for (int64_t d : srcDims)
+ inferredInstData[d] = inferredLaneLayout[d] * inferredLaneData[d];
+ int64_t remaining = resInstData[dstIdx] / laneAtom;
+ for (int64_t d : llvm::reverse(srcDims)) {
+ if (remaining == 1)
+ break;
+ int64_t cap = laneSrcShape[d] / inferredInstData[d];
+ int64_t take = std::min(remaining, cap);
+ assert(take > 0 && "inst_data distribution must not be zero");
+ assert(remaining % take == 0 &&
+ "inst_data must divide evenly across dims");
+ inferredInstData[d] *= take;
+ remaining /= take;
+ }
+ assert(remaining == 1 &&
+ "inst_data must fit within collapsed src dims");
+ }
+ }
+ }
+
+ auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
+ if (v.empty())
+ return DenseI32ArrayAttr();
+ SmallVector<int32_t> v32(v.begin(), v.end());
+ return DenseI32ArrayAttr::get(context, v32);
+ };
+
+ // Propagate order: for each dst dim taken in dst-order (fastest first),
+ // emit its collapsed src dims from innermost to outermost. Unit dst dims
+ // with no backing src (empty groups) contribute nothing.
+ // Example: src=[n1,n2,n3,n4,n5], dst=[m1,n3,m2] with collapse groups
+ // [[0,1],[2],[3,4]] and dst order=[1,2,0] -> src order=[2,4,3,1,0].
+ DenseI32ArrayAttr srcOrderAttr;
+ if (DenseI32ArrayAttr resOrder = resLayout.getOrder();
+ resOrder && !resOrder.empty()) {
+ SmallVector<int64_t> resOrderVec = resLayout.getEffectiveOrderAsInt();
+ SmallVector<int64_t> srcOrder;
+ srcOrder.reserve(srcShapeSize);
+ for (int64_t dstIdx : resOrderVec)
+ for (int64_t d : llvm::reverse(collapseDims[dstIdx]))
+ srcOrder.push_back(d);
+ srcOrderAttr = toAttr(srcOrder);
}
+
+ return xegpu::LayoutAttr::get(
+ context, toAttr(inferredSgLayout), toAttr(inferredSgData),
+ toAttr(inferredInstData), toAttr(inferredLaneLayout),
+ toAttr(inferredLaneData), srcOrderAttr);
}
llvm_unreachable("running into unsupported shape cast scenarios");
return nullptr;
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 0fb0ac6e3416d..b0b86fb979182 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -991,3 +991,56 @@ bool xegpu::matchSplitDimExpansion(
}
return srcIdx == src.size();
}
+
+// Checks if dst shape is a collapse of src shape where each dim in dst is
+// produced by one or more consecutive dims in src whose product equals the dst
+// dim. Populates collapseDims with one group per dst dim listing the src
+// indices collapsed into it. Unit dims in dst that have no backing src dim
+// (leading or trailing) get empty groups.
+// Examples:
+// src=[8,16,32], dst=[1,4096] -> true, collapseDims=[[],[0,1,2]]
+// src=[2,3,4], dst=[6,4] -> true, collapseDims=[[0,1],[2]]
+// src=[64], dst=[64] -> true, collapseDims=[[0]]
+bool xegpu::matchDimCollapse(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
+ SmallVector<SmallVector<int64_t>> &collapseDims) {
+ collapseDims.clear();
+ collapseDims.resize(dst.size());
+
+ size_t dstIdx = 0;
+ size_t srcIdx = 0;
+ int64_t accumulatedSize = 1;
+ SmallVector<int64_t> currentSrcDims;
+
+ // Skip any leading unit dst dims; they have no backing src dim.
+ while (dstIdx < dst.size() && dst[dstIdx] == 1)
+ dstIdx++;
+
+ while (srcIdx < src.size()) {
+ // Trailing src unit dims are absorbed into the most-recent group, or
+ // skipped if they appear before any group has started.
+ if (dstIdx >= dst.size()) {
+ if (src[srcIdx] != 1)
+ return false;
+ if (!collapseDims.empty() && !collapseDims.back().empty())
+ collapseDims.back().push_back(srcIdx);
+ srcIdx++;
+ continue;
+ }
+ accumulatedSize *= src[srcIdx];
+ currentSrcDims.push_back(srcIdx);
+ srcIdx++;
+
+ if (accumulatedSize == dst[dstIdx]) {
+ collapseDims[dstIdx] = currentSrcDims;
+ currentSrcDims.clear();
+ accumulatedSize = 1;
+ dstIdx++;
+ // Skip any subsequent unit dst dims; they have no backing src dim.
+ while (dstIdx < dst.size() && dst[dstIdx] == 1)
+ dstIdx++;
+ } else if (accumulatedSize > dst[dstIdx]) {
+ return false;
+ }
+ }
+ return dstIdx == dst.size() && currentSrcDims.empty();
+}
>From 3764009d4aec00604d7c13a39ea588a783dd2135 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 4 Jun 2026 01:05:35 +0000
Subject: [PATCH 2/5] [MLIR][XeGPU] Fix isTransposeOf, harden load/store_matrix
unroll, and add mxfp WG integration tests
- LayoutAttr::isTransposeOf: correct the per-dim check to match
vector.transpose semantics (dst[i] = src[perm[i]]); previously the
comparison indexed src and dst inversely, which could let mismatched
layouts compare equal.
- UnrollLoadMatrixOp / UnrollStoreMatrixOp: stop assuming the op's
layout is always a LayoutAttr. Use DistributeLayoutAttr and guard
dropInstData() with a null check so SliceAttr or missing-layout cases
no longer crash during unrolling.
- propagate-layout.mlir: update the lane_layout expectation for the
256 -> 2x4x32 shape_cast to match the generalized collapse inference
(lane_layout spreads across multiple src dims).
- Add WG-level integration tests for simple mxfp GEMM with quantizeA/F4
and dequantizeB/F4.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 5 +-
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 11 +-
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 2 +-
.../WG/simple_mxfp_gemm_dequantizeB_F4.mlir | 71 +++++++++++
.../WG/simple_mxfp_gemm_quantizeA_F4.mlir | 110 ++++++++++++++++++
5 files changed, 193 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir
create mode 100644 mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index e92b109c2223e..25b100d706e5f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -689,10 +689,13 @@ bool LayoutAttr::isTransposeOf(const xegpu::DistributeLayoutAttr &other,
return false;
if (!isPermutationVector(perm))
return false;
+ // vector.transpose semantics: dst[i] = src[perm[i]]. So `this` (= dst) is a
+ // transpose of `other` (= src) via `perm` iff for all i, dst[i] ==
+ // src[perm[i]].
auto checkTranspose = [](ArrayRef<int64_t> dst, ArrayRef<int64_t> src,
ArrayRef<int64_t> perm) {
for (const auto &ta : llvm::enumerate(perm)) {
- if (src[ta.index()] != dst[ta.value()])
+ if (dst[ta.index()] != src[ta.value()])
return false;
}
return true;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index aab36b79845e4..929b20b3b8d6e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -683,7 +683,7 @@ struct UnrollLoadMatrixOp : public UnrollPattern<xegpu::LoadMatrixOp> {
Type elemTy = valueTy.getElementType();
ArrayRef<int64_t> shape = valueTy.getShape();
- auto layout = dyn_cast<xegpu::LayoutAttr>(op.getLayoutAttr());
+ xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);
@@ -698,7 +698,8 @@ struct UnrollLoadMatrixOp : public UnrollPattern<xegpu::LoadMatrixOp> {
}
SmallVector<Value> newOps;
- layout = layout.dropInstData();
+ if (layout)
+ layout = layout.dropInstData();
for (SmallVector<OpFoldResult> offsets : offsetsList) {
auto newOp = xegpu::LoadMatrixOp::create(
rewriter, op.getLoc(), newValueTy, op.getMemDesc(), offsets, layout);
@@ -722,7 +723,9 @@ struct UnrollStoreMatrixOp : public UnrollPattern<xegpu::StoreMatrixOp> {
VectorType valueTy = llvm::dyn_cast<VectorType>(op.getData().getType());
assert(valueTy && "the value type must be vector type!");
ArrayRef<int64_t> shape = valueTy.getShape();
- auto layout = dyn_cast<xegpu::LayoutAttr>(op.getLayoutAttr());
+ xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
+ if (layout)
+ layout = layout.dropInstData();
SmallVector<Type> convertedValTypes =
getUnrolledTypes(valueTy, *targetShape);
@@ -741,7 +744,7 @@ struct UnrollStoreMatrixOp : public UnrollPattern<xegpu::StoreMatrixOp> {
for (auto [v, offsets] : llvm::zip_equal(convertedValues, offsetsList))
xegpu::StoreMatrixOp::create(rewriter, loc, v, op.getMemDesc(), offsets,
- layout.dropInstData());
+ layout);
rewriter.eraseOp(op);
return success();
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index a6907fa630d93..bad956d45d186 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -814,7 +814,7 @@ gpu.module @test {
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} dense<true> : vector<256xi1>
// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<256xindex>
// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
-// CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
+// CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [2, 4, 2], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
// CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>} : vector<2x4x32xf16> to vector<1x256xf16>
// CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<1x256xf16> to vector<256xf16>
// CHECK: xegpu.store %[[CAST_2]], %arg1[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}> : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir
new file mode 100644
index 0000000000000..9100cfa9b5b56
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir
@@ -0,0 +1,71 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=workgroup zebin-chip=cri" \
+// RUN: | mlir-runner \
+// RUN: --shared-libs=%mlir_levelzero_runtime \
+// RUN: --shared-libs=%mlir_runner_utils \
+// RUN: --shared-libs=%mlir_c_runner_utils \
+// RUN: --entry-point-result=void \
+// RUN: | FileCheck %s
+
+// XFAIL: *
+#a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#b_packed_ui8 = #xegpu.layout<sg_layout = [8, 8], sg_data = [256, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
+#b_f4 = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+#b_f16 = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16], inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>
+#c = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [2, 16], lane_layout = [1, 16], lane_data = [2, 1]>
+
+gpu.module @test {
+ gpu.func @gemm_quantized_b(%arg0: memref<1024x4096xbf16>, %arg1: memref<2048x1024xui8>, %arg2: memref<128x1024xf8E8M0FNU>, %arg3: memref<1024x1024xf32>) {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c128 = arith.constant 128 : index
+ %c1024 = arith.constant 1024 : index
+ %block_id_x = gpu.block_id x
+ %block_id_y = gpu.block_id y
+ %0 = arith.muli %block_id_x, %c128 : index
+ %1 = arith.muli %block_id_y, %c128 : index
+
+ %a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<1024x4096xbf16> -> !xegpu.tensor_desc<128x512xbf16>
+ %bp_tdesc = xegpu.create_nd_tdesc %arg1 : memref<2048x1024xui8> -> !xegpu.tensor_desc<256x128xui8>
+ // load_nd with offset
+ %a = xegpu.load_nd %a_tdesc[%0, %c0] {layout = #a}: !xegpu.tensor_desc<128x512xbf16> -> vector<128x512xbf16>
+ %bp = xegpu.load_nd %bp_tdesc[%c0, %1] {layout = #b_packed_ui8}: !xegpu.tensor_desc<256x128xui8> -> vector<256x128xui8>
+
+ // Bitcast to fp4: 256x128 uint8 -> 256x256 fp4 (each uint8 holds 2 fp4 values)
+ %b_bitcast = vector.bitcast %bp : vector<256x128xui8> to vector<256x256xf4E2M1FN>
+
+ // De-interleave: extract even and odd columns
+ // Even columns (indices 0, 2, 4, ..., 254) -> first half
+ // Odd columns (indices 1, 3, 5, ..., 255) -> second half
+ %b_even, %b_odd = vector.deinterleave %b_bitcast : vector<256x256xf4E2M1FN> -> vector<256x128xf4E2M1FN>
+
+ // Reconstruct 512x128 by interleaving even/odd rows:
+ // Transpose to move the row dim to trailing position, interleave, transpose back.
+ %b_even_t = vector.transpose %b_even, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_odd_t = vector.transpose %b_odd, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_interleaved = vector.interleave %b_even_t, %b_odd_t : vector<128x256xf4E2M1FN> -> vector<128x512xf4E2M1FN>
+ %b = vector.transpose %b_interleaved, [1, 0] : vector<128x512xf4E2M1FN> to vector<512x128xf4E2M1FN>
+
+ %cd_tdesc = xegpu.create_nd_tdesc %arg3 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32, #c>
+ %c = xegpu.load_nd %cd_tdesc[%0, %1] {layout = #c}: !xegpu.tensor_desc<128x128xf32, #c> -> vector<128x128xf32>
+
+ %b_scale_tdesc = xegpu.create_nd_tdesc %arg2 : memref<128x1024xf8E8M0FNU> -> !xegpu.tensor_desc<16x128xf8E8M0FNU>
+ %scale_b = xegpu.load_nd %b_scale_tdesc[%c0, %1] {layout = #b_scale}: !xegpu.tensor_desc<16x128xf8E8M0FNU> -> vector<16x128xf8E8M0FNU>
+
+ // Broadcast scale_b from <16x128> to <512x128>: each scale value applies to
+ // 32 consecutive K rows of B.
+ %scale_b_bcast = vector.broadcast %scale_b : vector<16x128xf8E8M0FNU> to vector<32x16x128xf8E8M0FNU>
+ %scale_b_t = vector.transpose %scale_b_bcast, [1, 0, 2] : vector<32x16x128xf8E8M0FNU> to vector<16x32x128xf8E8M0FNU>
+ %scale_b_full = vector.shape_cast %scale_b_t : vector<16x32x128xf8E8M0FNU> to vector<512x128xf8E8M0FNU>
+
+ // Dequantize B from f4E2M1FN to bf16 using scale_b.
+ %b_bf16 = arith.scaling_extf %b, %scale_b_full : vector<512x128xf4E2M1FN>, vector<512x128xf8E8M0FNU> to vector<512x128xbf16>
+
+ %d = xegpu.dpas %a, %b_bf16, %c {layout_a = #a, layout_b = #b_f16, layout_cd = #c}
+ : vector<128x512xbf16>, vector<512x128xbf16>, vector<128x128xf32> -> vector<128x128xf32>
+
+ // store_nd with offset
+ xegpu.store_nd %d, %cd_tdesc[%0, %1] {layout = #c} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #c>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
new file mode 100644
index 0000000000000..48144d89cd2f4
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
@@ -0,0 +1,110 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=workgroup zebin-chip=cri" \
+// RUN: | mlir-runner \
+// RUN: --shared-libs=%mlir_levelzero_runtime \
+// RUN: --shared-libs=%mlir_runner_utils \
+// RUN: --shared-libs=%mlir_c_runner_utils \
+// RUN: --entry-point-result=void \
+// RUN: | FileCheck %s
+
+// XFAIL: *
+#a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
+#b_packed = #xegpu.layout<sg_layout = [8, 8], sg_data = [256, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
+#b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+#c = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#b_scale_ld = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#a_scale_dpas = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 2]>
+#b_scale_dpas = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+
+gpu.module @test {
+ // A is loaded as bf16 and quantized in-place to mx-fp4 (fp4 + f8E8M0 scale)
+ // along the K dimension with block size 32. B and its scale are passed in
+ // pre-quantized (packed ui8 fp4 and f8E8M0). The quantized values are then
+ // consumed by xegpu.dpas_mx.
+ gpu.func @gemm_mxfp(%arg0: memref<1024x4096xbf16>, %arg1: memref<2048x1024xui8>, %arg2: memref<128x1024xf8E8M0FNU>, %arg3: memref<1024x1024xf32>) {
+ %c0 = arith.constant 0 : index
+ %c128 = arith.constant 128 : index
+ %block_id_x = gpu.block_id x
+ %block_id_y = gpu.block_id y
+ %0 = arith.muli %block_id_x, %c128 : index
+ %1 = arith.muli %block_id_y, %c128 : index
+
+ // -------- Load A (bf16) --------
+ %a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<1024x4096xbf16> -> !xegpu.tensor_desc<128x512xbf16>
+ %a_bf16 = xegpu.load_nd %a_tdesc[%0, %c0] {layout = #a}: !xegpu.tensor_desc<128x512xbf16> -> vector<128x512xbf16>
+
+ // -------- Quantize A: bf16 -> fp4 + f8E8M0 scale (block_size=32 along K) --------
+ // 1) abs and reduce-max per block of 32 along K dim using vector ops.
+ %a_abs = math.absf %a_bf16 : vector<128x512xbf16>
+ %a_abs_r = vector.shape_cast %a_abs : vector<128x512xbf16> to vector<128x16x32xbf16>
+ %a_neg_inf_i = arith.constant dense<0xFF80> : vector<128x16xi16>
+ %a_neg_inf = arith.bitcast %a_neg_inf_i : vector<128x16xi16> to vector<128x16xbf16>
+ %a_amax = vector.multi_reduction <maximumf>, %a_abs_r, %a_neg_inf [2]
+ : vector<128x16x32xbf16> to vector<128x16xbf16>
+
+ // 2) Largest power-of-two <= amax: mask out mantissa bits of bf16.
+ %a_amax_i16 = arith.bitcast %a_amax : vector<128x16xbf16> to vector<128x16xi16>
+ %a_exp_mask = arith.constant dense<0x7F80> : vector<128x16xi16>
+ %a_pow2_i16 = arith.andi %a_amax_i16, %a_exp_mask : vector<128x16xi16>
+ %a_pow2 = arith.bitcast %a_pow2_i16 : vector<128x16xi16> to vector<128x16xbf16>
+
+ // 3) Divide by largest power-of-two representable by E2M1 (= 4.0).
+ %a_e2m1_max = arith.constant dense<4.000000e+00> : vector<128x16xbf16>
+ %a_scale_bf16 = arith.divf %a_pow2, %a_e2m1_max : vector<128x16xbf16>
+
+ // 4) Truncate scale to f8E8M0FNU.
+ %a_scale = arith.truncf %a_scale_bf16 : vector<128x16xbf16> to vector<128x16xf8E8M0FNU>
+
+ // 5) Broadcast the per-block scale across the block (32 elements along K).
+ // vector.broadcast can only prepend leading dims, so we broadcast onto a
+ // leading 32 dim, transpose it to the trailing position, then shape_cast.
+ %a_scale_lead = vector.broadcast %a_scale
+ : vector<128x16xf8E8M0FNU> to vector<32x128x16xf8E8M0FNU>
+ %a_scale_t = vector.transpose %a_scale_lead, [1, 2, 0]
+ : vector<32x128x16xf8E8M0FNU> to vector<128x16x32xf8E8M0FNU>
+ %a_scale_full = vector.shape_cast %a_scale_t
+ : vector<128x16x32xf8E8M0FNU> to vector<128x512xf8E8M0FNU>
+
+ // 6) Scaled truncf to fp4 (to_nearest_even).
+ %a = arith.scaling_truncf %a_bf16, %a_scale_full
+ : vector<128x512xbf16>, vector<128x512xf8E8M0FNU> to vector<128x512xf4E2M1FN>
+
+ // -------- Load B (packed ui8 fp4) and unpack to f4E2M1FN --------
+ %bp_tdesc = xegpu.create_nd_tdesc %arg1 : memref<2048x1024xui8> -> !xegpu.tensor_desc<256x128xui8>
+ %bp = xegpu.load_nd %bp_tdesc[%c0, %1] {layout = #b_packed}: !xegpu.tensor_desc<256x128xui8> -> vector<256x128xui8>
+
+ // Bitcast to fp4: 256x128 uint8 -> 256x256 fp4 (each uint8 holds 2 fp4 values)
+ %b_bitcast = vector.bitcast %bp : vector<256x128xui8> to vector<256x256xf4E2M1FN>
+
+ // De-interleave: extract even and odd columns
+ %b_even, %b_odd = vector.deinterleave %b_bitcast : vector<256x256xf4E2M1FN> -> vector<256x128xf4E2M1FN>
+
+ // Reconstruct 512x128 by interleaving even/odd rows.
+ %b_even_t = vector.transpose %b_even, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_odd_t = vector.transpose %b_odd, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_interleaved = vector.interleave %b_even_t, %b_odd_t : vector<128x256xf4E2M1FN> -> vector<128x512xf4E2M1FN>
+ %b = vector.transpose %b_interleaved, [1, 0] : vector<128x512xf4E2M1FN> to vector<512x128xf4E2M1FN>
+
+ // -------- Load B scale --------
+ %b_scale_tdesc = xegpu.create_nd_tdesc %arg2 : memref<128x1024xf8E8M0FNU> -> !xegpu.tensor_desc<16x128xf8E8M0FNU>
+ %b_scale = xegpu.load_nd %b_scale_tdesc[%c0, %1] {layout = #b_scale_ld}: !xegpu.tensor_desc<16x128xf8E8M0FNU> -> vector<16x128xf8E8M0FNU>
+
+ // -------- Load C and run dpas_mx --------
+ %cd_tdesc = xegpu.create_nd_tdesc %arg3 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32, #c>
+ %c = xegpu.load_nd %cd_tdesc[%0, %1] {layout = #c}: !xegpu.tensor_desc<128x128xf32, #c> -> vector<128x128xf32>
+
+ %d = xegpu.dpas_mx %a, %b, %c scale_a = %a_scale scale_b = %b_scale
+ {layout_a = #a,
+ layout_b = #b,
+ layout_cd = #c,
+ layout_a_scale = #a_scale_dpas,
+ layout_b_scale = #b_scale_dpas}
+ : (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>,
+ vector<128x128xf32>,
+ vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>)
+ -> vector<128x128xf32>
+
+ // store_nd with offset
+ xegpu.store_nd %d, %cd_tdesc[%0, %1] {layout = #c} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #c>
+ gpu.return
+ }
+}
>From 36154b4267f5f16e777a8260771d452289063cba Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 4 Jun 2026 02:22:05 +0000
Subject: [PATCH 3/5] [MLIR][XeGPU] Add tests for shape_cast collapse layout
inference
Cover three patterns at both subgroup and inst-data levels:
- Plain collapse: all src dims fold into a single innermost dst dim.
- Layout spill: the consumer layout (sg_layout / lane_layout) is larger
than a single src dim, so the value spreads across adjacent src dims.
- Multi-group: every dst dim collapses >=2 src dims, exercising the
general matchDimCollapse path.
Inst-data tests use the canonical anchor combining inst_data with
lane_layout=[..,subgroupSize] and lane_data=[..,1], which is the
primary target of the new code.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../XeGPU/propagate-layout-inst-data.mlir | 75 +++++++++++++++++++
.../XeGPU/propagate-layout-subgroup.mlir | 69 +++++++++++++++++
2 files changed, 144 insertions(+)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 6f587959b697d..cace6b736fd81 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -530,3 +530,78 @@ func.func @dpas_mx_f4e2m1(%arg0: memref<16x128xf4E2M1FN>, %arg1: memref<128x32xf
return
}
}
+
+// -----
+// shape_cast that collapses all src dims into a single innermost dst dim.
+// Consumer carries inst_data + lane_layout=[..,subgroupSize] + lane_data=[..,1]
+// (the canonical inst-level anchor for this code path).
+// srcShape=[8, 16, 32], resShape=[4096], consumer inst_data=[32],
+// lane_layout=[16], lane_data=[1]
+// lane_layout outer-to-inner: dim0 take=min(16,8)=8 (rem=2);
+// dim1 take=min(2,16)=2 (rem=1) -> [8, 2, 1].
+// lane_data innermost-first: total=1 -> [1, 1, 1].
+// inst_data: laneAtom=16, seed [8, 2, 1]; remaining=32/16=2 spreads
+// innermost-first: dim2 cap=32/1=32, take=2 -> inst_data[2]=2 -> [8, 2, 2].
+gpu.module @test {
+// CHECK-LABEL: func.func @vector_shape_cast_collapse_innermost(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 2, 2], lane_layout = [8, 2, 1], lane_data = [1, 1, 1]>} dense<0.000000e+00> : vector<8x16x32xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [1]>} : vector<8x16x32xf16> to vector<4096xf16>
+func.func @vector_shape_cast_collapse_innermost(%arg0: memref<4096xf16>) {
+ %cst_v = arith.constant dense<0.000000e+00> : vector<8x16x32xf16>
+ %0 = vector.shape_cast %cst_v : vector<8x16x32xf16> to vector<4096xf16>
+ %mask = arith.constant dense<true> : vector<4096xi1>
+ %offsets = vector.step : vector<4096xindex>
+ xegpu.store %0, %arg0[%offsets], %mask <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [1]>}> : vector<4096xf16>, memref<4096xf16>, vector<4096xindex>, vector<4096xi1>
+ return
+ }
+}
+
+// -----
+// shape_cast collapse where the outermost src dim is too small to absorb the
+// full lane_layout, so it spills inward across multiple src dims.
+// srcShape=[2, 8, 32], resShape=[512], consumer inst_data=[64],
+// lane_layout=[16], lane_data=[2]
+// lane_layout outer-to-inner: dim0 take=min(16,2)=2 (rem=8);
+// dim1 take=min(8,8)=8 (rem=1) -> [2, 8, 1].
+// lane_data innermost-first: total=2; dim2 cap=32/1=32, take=2 -> [1, 1, 2].
+// inst_data: laneAtom=16*2=32, seed [2*1, 8*1, 1*2] = [2, 8, 2];
+// remaining=64/32=2 spreads innermost-first: dim2 cap=32/2=16, take=2
+// -> inst_data[2] *= 2 -> [2, 8, 4].
+gpu.module @test {
+// CHECK-LABEL: func.func @vector_shape_cast_collapse_lane_layout_spill_inward(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [2, 8, 4], lane_layout = [2, 8, 1], lane_data = [1, 1, 2]>} dense<0.000000e+00> : vector<2x8x32xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<inst_data = [64], lane_layout = [16], lane_data = [2]>} : vector<2x8x32xf16> to vector<512xf16>
+func.func @vector_shape_cast_collapse_lane_layout_spill_inward(%arg0: memref<512xf16>) {
+ %cst_v = arith.constant dense<0.000000e+00> : vector<2x8x32xf16>
+ %0 = vector.shape_cast %cst_v : vector<2x8x32xf16> to vector<512xf16>
+ %mask = arith.constant dense<true> : vector<512xi1>
+ %offsets = vector.step : vector<512xindex>
+ xegpu.store %0, %arg0[%offsets], %mask <{layout = #xegpu.layout<inst_data = [64], lane_layout = [16], lane_data = [2]>}> : vector<512xf16>, memref<512xf16>, vector<512xindex>, vector<512xi1>
+ return
+ }
+}
+
+// -----
+// shape_cast collapse with multiple non-trivial groups: every dst dim
+// collapses >=2 src dims, exercising the general matchDimCollapse path.
+// srcShape=[2, 4, 8, 16], resShape=[8, 128], consumer inst_data=[1, 16],
+// lane_layout=[1, 16], lane_data=[1, 1]
+// - dst[0]=8 collapses src[0, 1]: lane_layout=1, lane_data=1, inst_data=1
+// -> [1, 1, _, _] for all three.
+// - dst[1]=128 collapses src[2, 3]: lane_layout outer-to-inner over [2, 3]:
+// dim2 take=min(16, 8)=8 (rem=2); dim3 take=min(2, 16)=2 (rem=1)
+// -> [_, _, 8, 2]; lane_data=1 -> [_, _, 1, 1]; inst_data laneAtom=16,
+// seed [_, _, 8*1, 2*1] = [_, _, 8, 2]; remaining=16/16=1 -> done
+// -> inst_data=[_, _, 8, 2].
+gpu.module @test {
+// CHECK-LABEL: func.func @vector_shape_cast_collapse_multi_groups(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 8, 2], lane_layout = [1, 1, 8, 2], lane_data = [1, 1, 1, 1]>} dense<0.000000e+00> : vector<2x4x8x16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<2x4x8x16xf16> to vector<8x128xf16>
+func.func @vector_shape_cast_collapse_multi_groups(%arg0: memref<8x128xf16>) {
+ %cst_v = arith.constant dense<0.000000e+00> : vector<2x4x8x16xf16>
+ %0 = vector.shape_cast %cst_v : vector<2x4x8x16xf16> to vector<8x128xf16>
+ %tdesc = xegpu.create_nd_tdesc %arg0 : memref<8x128xf16> -> !xegpu.tensor_desc<8x128xf16>
+ xegpu.store_nd %0, %tdesc[0, 0] <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x128xf16>, !xegpu.tensor_desc<8x128xf16>
+ return
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 98d3ce9d9c4cd..5021c8a746045 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -506,3 +506,72 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+// shape_cast collapse: all src dims fold into a single innermost dst dim.
+// srcShape=[8, 16, 32], resShape=[4096], consumer sg_layout=[8],
+// sg_data=[512]
+// sg_layout outer-to-inner: dim0 take=min(8, 8)=8 (rem=1) -> [8, 1, 1].
+// sg_data innermost-first, capped per dim by srcShape[d]/sgLayout[d]:
+// dim2 cap=32/1=32, take=32 (rem=16); dim1 cap=16/1=16, take=16 (rem=1)
+// -> [1, 16, 32].
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_collapse_innermost(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8, 1, 1], sg_data = [1, 16, 32]>} dense<0.000000e+00> : vector<8x16x32xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [512]>} : vector<8x16x32xf16> to vector<4096xf16>
+ gpu.func @shape_cast_collapse_innermost(%dst: memref<4096xf16>) kernel {
+ %cst = arith.constant dense<0.000000e+00> : vector<8x16x32xf16>
+ %0 = vector.shape_cast %cst : vector<8x16x32xf16> to vector<4096xf16>
+ %mask = arith.constant dense<true> : vector<4096xi1>
+ %offsets = vector.step : vector<4096xindex>
+ xegpu.store %0, %dst[%offsets], %mask <{layout = #xegpu.layout<sg_layout = [8], sg_data = [512]>}> : vector<4096xf16>, memref<4096xf16>, vector<4096xindex>, vector<4096xi1>
+ gpu.return
+ }
+}
+
+// -----
+// shape_cast collapse where sg_layout exceeds the size of the outermost src
+// dim, so it spills inward across multiple src dims.
+// srcShape=[2, 8, 32], resShape=[512], consumer sg_layout=[16], sg_data=[32]
+// sg_layout outer-to-inner: dim0 take=min(16, 2)=2 (rem=8);
+// dim1 take=min(8, 8)=8 (rem=1) -> [2, 8, 1].
+// sg_data innermost-first, capped per dim by srcShape[d]/sgLayout[d]:
+// dim2 cap=32/1=32, take=32 (rem=1) -> [1, 1, 32].
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_collapse_sg_layout_spill_inward(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [2, 8, 1], sg_data = [1, 1, 32]>} dense<0.000000e+00> : vector<2x8x32xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [32]>} : vector<2x8x32xf16> to vector<512xf16>
+ gpu.func @shape_cast_collapse_sg_layout_spill_inward(%dst: memref<512xf16>) kernel {
+ %cst = arith.constant dense<0.000000e+00> : vector<2x8x32xf16>
+ %0 = vector.shape_cast %cst : vector<2x8x32xf16> to vector<512xf16>
+ %mask = arith.constant dense<true> : vector<512xi1>
+ %offsets = vector.step : vector<512xindex>
+ xegpu.store %0, %dst[%offsets], %mask <{layout = #xegpu.layout<sg_layout = [16], sg_data = [32]>}> : vector<512xf16>, memref<512xf16>, vector<512xindex>, vector<512xi1>
+ gpu.return
+ }
+}
+
+// -----
+// shape_cast collapse with multiple non-trivial groups: every dst dim
+// collapses >=2 src dims, exercising the general matchDimCollapse path.
+// srcShape=[2, 4, 8, 16], resShape=[8, 128], consumer sg_layout=[2, 4],
+// sg_data=[4, 32]
+// - dst[0]=8 collapses src[0, 1]: sg_layout outer-to-inner over [0, 1]:
+// dim0 take=min(2, 2)=2 (rem=1) -> [2, 1, _, _]; sg_data innermost-first:
+// dim1 cap=4/1=4, take=4 (rem=1) -> [1, 4, _, _].
+// - dst[1]=128 collapses src[2, 3]: sg_layout outer-to-inner: dim2 take=
+// min(4, 8)=4 (rem=1) -> [_, _, 4, 1]; sg_data innermost-first: dim3
+// cap=16/1=16, take=16 (rem=2); dim2 cap=8/4=2, take=2 (rem=1)
+// -> [_, _, 2, 16].
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_collapse_multi_groups(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [2, 1, 4, 1], sg_data = [1, 4, 2, 16]>} dense<0.000000e+00> : vector<2x4x8x16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [4, 32]>} : vector<2x4x8x16xf16> to vector<8x128xf16>
+ gpu.func @shape_cast_collapse_multi_groups(%dst: memref<8x128xf16>) kernel {
+ %cst = arith.constant dense<0.000000e+00> : vector<2x4x8x16xf16>
+ %0 = vector.shape_cast %cst : vector<2x4x8x16xf16> to vector<8x128xf16>
+ %tdesc = xegpu.create_nd_tdesc %dst : memref<8x128xf16> -> !xegpu.tensor_desc<8x128xf16>
+ xegpu.store_nd %0, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [4, 32]>}> : vector<8x128xf16>, !xegpu.tensor_desc<8x128xf16>
+ gpu.return
+ }
+}
>From b94cd0794b66bc76cbc66d0ef53bf5f9513f6167 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 4 Jun 2026 03:00:36 +0000
Subject: [PATCH 4/5] [MLIR][XeGPU] Refactor shape_cast collapse via expandDims
interface
Add `expandDims(int64_t dim, ArrayRef<int64_t> targetShape)` to the
DistributeLayoutAttr interface, with implementations on both LayoutAttr
and SliceAttr. The method is the rank-increasing dual of `collapseDims`
and bakes in the distribution policy required for a no-data-movement
collapse: sg_layout / lane_layout spread outer-to-inner; sg_data /
lane_data / inst_data fill innermost-first; inst_data is seeded from
lane_layout * lane_data per new dim; `order` is rewritten so the
expanded dims occupy the original dim's slot in innermost-fastest
order.
With this primitive in hand, use case 3 in inferShapeCastSourceLayout
collapses to a use-case-2-style loop: for each dst-side group (in
reverse), call expandDims when the group is non-trivial or dropDims
when the dst dim has no backing src dim. The previous ~190-line inline
distribution body is replaced by a handful of lines.
Also fix LayoutAttr::dropDims so it preserves "no order" when the input
had no order attribute; previously it synthesized the default
[rank-1,...,0] order, which then tripped collapseDims's adjacency
check on a downstream pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/IR/XeGPUAttrs.td | 24 ++
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 237 +++++++++++++++-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 262 ++----------------
3 files changed, 280 insertions(+), 243 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 40edce8a60429..474226b526b37 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -214,6 +214,18 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
"xegpu::DistributeLayoutAttr",
"collapseDims",
(ins "SmallVector<int64_t>": $dimGroup)>,
+ InterfaceMethod<[{Derive a new layout by expanding a single dimension into
+ multiple adjacent dimensions whose extents are given by `targetShape`
+ (`product(targetShape) == originalSize` on `dim`). The original layout's
+ per-dim values on `dim` are distributed across the new dims:
+ `sg_layout` and `lane_layout` spread outer-to-inner (capped per dim by
+ `targetShape[i]`); `sg_data`, `lane_data`, and `inst_data` fill innermost-
+ first; `inst_data` is seeded from `lane_layout * lane_data` per new dim.
+ `order` is rewritten so the expanded dims appear innermost-fastest in the
+ original dim's slot.}],
+ "xegpu::DistributeLayoutAttr",
+ "expandDims",
+ (ins "int64_t": $dim, "ArrayRef<int64_t>": $targetShape)>,
InterfaceMethod<[{Derive a new layout by trasnposing it using `permutation`.}],
"xegpu::DistributeLayoutAttr",
"transposeDims",
@@ -582,6 +594,12 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
// that are collapsed into a single dimension in the derived layout.
DistributeLayoutAttr collapseDims(SmallVector<int64_t> dimGroup);
+ // Derive a new layout by expanding a single dimension `dim` into
+ // multiple adjacent dimensions whose shape is given by `targetShape`.
+ // See the interface method documentation for the per-field distribution
+ // policy.
+ DistributeLayoutAttr expandDims(int64_t dim, ArrayRef<int64_t> targetShape);
+
// Derive a new layout by transposing the layout using `permutation`.
DistributeLayoutAttr transposeDims(ArrayRef<int64_t> permutation);
@@ -814,6 +832,12 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
// that are collapsed into a single dimension in the derived layout.
DistributeLayoutAttr collapseDims(SmallVector<int64_t> dimGroup);
+ // Derive a new layout by expanding a single dimension `dim` into
+ // multiple adjacent dimensions whose shape is given by `targetShape`.
+ // See the interface method documentation for the per-field distribution
+ // policy.
+ DistributeLayoutAttr expandDims(int64_t dim, ArrayRef<int64_t> targetShape);
+
// Derive a new layout by transposing the layout using `permutation`.
DistributeLayoutAttr transposeDims(ArrayRef<int64_t> permutation);
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 25b100d706e5f..64c04437cd723 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -486,7 +486,7 @@ DistributeLayoutAttr LayoutAttr::dropDims(SmallVector<int64_t> dimGroup) {
SmallVector<int64_t> instData = getEffectiveInstDataAsInt();
SmallVector<int64_t> laneLayout = getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> laneData = getEffectiveLaneDataAsInt();
- SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
+ DenseI32ArrayAttr origOrderAttr = getOrder();
SmallVector<int64_t> sortedDimGroup = dimGroup;
llvm::sort(sortedDimGroup);
@@ -504,15 +504,22 @@ DistributeLayoutAttr LayoutAttr::dropDims(SmallVector<int64_t> dimGroup) {
}
}
+ // Only emit a new order attribute when the input had one, so that "no
+ // order" inputs do not gain a synthetic default that would later trip
+ // adjacency checks in collapseDims.
SmallVector<int64_t> newOrder;
- for (int64_t d : origOrder) {
- if (llvm::is_contained(dimGroup, d))
- continue;
- int64_t offset = llvm::count_if(dimGroup, [&](int64_t s) { return s < d; });
- newOrder.push_back(d - offset);
+ if (origOrderAttr && !origOrderAttr.empty()) {
+ SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
+ for (int64_t d : origOrder) {
+ if (llvm::is_contained(dimGroup, d))
+ continue;
+ int64_t offset =
+ llvm::count_if(dimGroup, [&](int64_t s) { return s < d; });
+ newOrder.push_back(d - offset);
+ }
+ if ((sgLayout.empty() && laneLayout.empty()) || newOrder.size() == 1)
+ newOrder.clear();
}
- if ((sgLayout.empty() && laneLayout.empty()) || newOrder.size() == 1)
- newOrder.clear();
auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
if (v.empty())
@@ -637,6 +644,193 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
return collapsedLayout;
}
+// Derive a new layout by expanding a single dimension `dim` into multiple
+// adjacent dimensions whose extents are given by `targetShape`.
+//
+// Distribution policy on the expanded src dims (replacing `dim`):
+// - sg_layout / lane_layout: spread outer-to-inner; each dim takes
+// min(remaining, targetShape[i]); leftover spills into the next inner
+// dim.
+// - sg_data: fill innermost-first, capped per dim by
+// targetShape[i] / sgLayout[i] (the per-sg share of the extent).
+// - lane_data: fill innermost-first, capped per dim by
+// (targetShape[i] / sgLayout[i]) / laneLayout[i] (the per-lane share of
+// the per-sg extent).
+// - inst_data: seeded from laneLayout[i] * laneData[i] per dim, then the
+// remaining factor is distributed innermost-first (capped per dim by
+// the per-sg extent).
+// - order: the original dim index is replaced by the expanded dim indices
+// in innermost-fastest order; entries past `dim` shift up by
+// `targetShape.size() - 1`.
+DistributeLayoutAttr
+LayoutAttr::expandDims(int64_t dim, ArrayRef<int64_t> targetShape) {
+ SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
+ SmallVector<int64_t> sgData = getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneLayout = getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = getEffectiveLaneDataAsInt();
+
+ int64_t origRank = getRank();
+ int64_t expCount = static_cast<int64_t>(targetShape.size());
+ assert(dim >= 0 && dim < origRank && "dim out of range");
+ assert(expCount >= 1 && "targetShape must have at least one dim");
+ int64_t newRank = origRank + expCount - 1;
+
+ // Snapshot the per-dim values we need before any field is mutated by
+ // splice() below; without this, computations that read e.g. `laneLayout[dim]`
+ // after laneLayout has already been expanded would see the wrong value.
+ int64_t origSgLayoutDim = sgLayout.empty() ? 1 : sgLayout[dim];
+ int64_t origSgDataDim = sgData.empty() ? 1 : sgData[dim];
+ int64_t origLaneLayoutDim = laneLayout.empty() ? 1 : laneLayout[dim];
+ int64_t origLaneDataDim = laneData.empty() ? 1 : laneData[dim];
+ int64_t origInstDataDim = instData.empty() ? 1 : instData[dim];
+ (void)origSgLayoutDim;
+ (void)origSgDataDim;
+
+ // Spread `total` across `targetShape` (length expCount), capped per dim by
+ // perDimCap[i]. `outerToInner` selects iteration direction (true = i=0..n-1).
+ auto spread = [&](int64_t total, ArrayRef<int64_t> perDimCap,
+ bool outerToInner) -> SmallVector<int64_t> {
+ SmallVector<int64_t> out(expCount, 1);
+ int64_t remaining = total;
+ auto step = [&](int64_t i) {
+ if (remaining == 1)
+ return;
+ int64_t take = std::min(remaining, perDimCap[i]);
+ assert(take > 0 && "expandDims distribution must not be zero");
+ assert(remaining % take == 0 &&
+ "expandDims must divide evenly across dims");
+ out[i] = take;
+ remaining /= take;
+ };
+ if (outerToInner)
+ for (int64_t i = 0; i < expCount; ++i)
+ step(i);
+ else
+ for (int64_t i = expCount - 1; i >= 0; --i)
+ step(i);
+ assert(remaining == 1 && "expandDims total must fit within target shape");
+ return out;
+ };
+
+ // Splice `expanded` (length expCount) into `vec` at position `dim`,
+ // replacing the single entry at `dim`.
+ auto splice = [&](SmallVector<int64_t> &vec,
+ ArrayRef<int64_t> expanded) {
+ if (vec.empty())
+ return;
+ vec.erase(vec.begin() + dim);
+ vec.insert(vec.begin() + dim, expanded.begin(), expanded.end());
+ };
+
+ bool hasSgLayout = !sgLayout.empty();
+ bool hasSgData = !sgData.empty();
+ bool hasLaneLayout = !laneLayout.empty();
+ bool hasLaneData = !laneData.empty();
+ bool hasInstData = !instData.empty();
+
+ // sg_layout / sg_data
+ SmallVector<int64_t> expSgLayout(expCount, 1);
+ if (hasSgLayout) {
+ expSgLayout = spread(origSgLayoutDim, targetShape, /*outerToInner=*/true);
+ splice(sgLayout, expSgLayout);
+ }
+ if (hasSgData) {
+ SmallVector<int64_t> cap(targetShape.begin(), targetShape.end());
+ if (hasSgLayout)
+ for (int64_t i = 0; i < expCount; ++i)
+ cap[i] /= expSgLayout[i];
+ SmallVector<int64_t> expSgData =
+ spread(origSgDataDim, cap, /*outerToInner=*/false);
+ splice(sgData, expSgData);
+ }
+
+ // Per-sg view used as the base for lane_layout / lane_data / inst_data:
+ // targetShape[i] / sg_layout[i] when sg_layout is present, else
+ // targetShape itself.
+ SmallVector<int64_t> perSgShape(targetShape.begin(), targetShape.end());
+ if (hasSgLayout)
+ for (int64_t i = 0; i < expCount; ++i)
+ perSgShape[i] /= expSgLayout[i];
+
+ // lane_layout / lane_data
+ SmallVector<int64_t> expLaneLayout(expCount, 1);
+ SmallVector<int64_t> expLaneData(expCount, 1);
+ if (hasLaneLayout) {
+ expLaneLayout = spread(origLaneLayoutDim, perSgShape,
+ /*outerToInner=*/true);
+ splice(laneLayout, expLaneLayout);
+ }
+ if (hasLaneData) {
+ SmallVector<int64_t> cap(perSgShape.begin(), perSgShape.end());
+ if (hasLaneLayout)
+ for (int64_t i = 0; i < expCount; ++i)
+ cap[i] /= expLaneLayout[i];
+ expLaneData = spread(origLaneDataDim, cap, /*outerToInner=*/false);
+ splice(laneData, expLaneData);
+ }
+
+ // inst_data: seed each new dim with laneLayout[i] * laneData[i]; spread the
+ // remaining factor innermost-first capped per dim by perSgShape[i] / seed.
+ if (hasInstData) {
+ SmallVector<int64_t> expInstData(expCount, 1);
+ if (!hasLaneLayout || !hasLaneData) {
+ expInstData = spread(origInstDataDim, perSgShape, /*outerToInner=*/false);
+ } else {
+ int64_t laneAtom = origLaneLayoutDim * origLaneDataDim;
+ for (int64_t i = 0; i < expCount; ++i)
+ expInstData[i] = expLaneLayout[i] * expLaneData[i];
+ SmallVector<int64_t> cap(expCount, 0);
+ for (int64_t i = 0; i < expCount; ++i)
+ cap[i] = perSgShape[i] / expInstData[i];
+ int64_t remaining = origInstDataDim / laneAtom;
+ for (int64_t i = expCount - 1; i >= 0; --i) {
+ if (remaining == 1)
+ break;
+ int64_t take = std::min(remaining, cap[i]);
+ assert(take > 0 && "inst_data distribution must not be zero");
+ assert(remaining % take == 0 &&
+ "inst_data must divide evenly across dims");
+ expInstData[i] *= take;
+ remaining /= take;
+ }
+ assert(remaining == 1 && "inst_data must fit within target shape");
+ }
+ splice(instData, expInstData);
+ }
+
+ // order: replace `dim`'s entry with the expanded dim indices in
+ // innermost-fastest order; shift every other entry past `dim` up by
+ // (expCount - 1).
+ SmallVector<int64_t> newOrder;
+ DenseI32ArrayAttr orderAttr = getOrder();
+ if (orderAttr && !orderAttr.empty()) {
+ SmallVector<int64_t> origOrder = getEffectiveOrderAsInt();
+ newOrder.reserve(newRank);
+ for (int64_t o : origOrder) {
+ if (o == dim) {
+ // Innermost dim of the expanded group is fastest-varying.
+ for (int64_t i = expCount - 1; i >= 0; --i)
+ newOrder.push_back(dim + i);
+ } else if (o > dim) {
+ newOrder.push_back(o + expCount - 1);
+ } else {
+ newOrder.push_back(o);
+ }
+ }
+ }
+
+ auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
+ if (v.empty())
+ return DenseI32ArrayAttr();
+ SmallVector<int32_t> v32(v.begin(), v.end());
+ return DenseI32ArrayAttr::get(getContext(), v32);
+ };
+ return xegpu::LayoutAttr::get(
+ getContext(), toAttr(sgLayout), toAttr(sgData), toAttr(instData),
+ toAttr(laneLayout), toAttr(laneData), toAttr(newOrder));
+}
+
// Derive a new layout by transpose the layout using `permutation`.
DistributeLayoutAttr LayoutAttr::transposeDims(ArrayRef<int64_t> permutation) {
@@ -1145,6 +1339,33 @@ DistributeLayoutAttr SliceAttr::collapseDims(SmallVector<int64_t> dimGroup) {
DenseI64ArrayAttr::get(getContext(), sliceDims));
}
+// Derive a new layout by expanding a single sliced-space dim into multiple
+// adjacent dims. The dim is mapped to parent space, the parent layout is
+// expanded there, and the slice dims that lie past the expanded position
+// are shifted up by `targetShape.size() - 1`.
+DistributeLayoutAttr SliceAttr::expandDims(int64_t dim,
+ ArrayRef<int64_t> targetShape) {
+ // `dim` is in slice space; map it to parent space (parent dims listed in
+ // `sliceDims` are removed by the slice, so the mapping always lands on a
+ // non-sliced parent dim).
+ SmallVector<int64_t> sliceDims = llvm::to_vector(getDims().asArrayRef());
+ SmallVector<int64_t> dimSet = {dim};
+ SmallVector<int64_t> dimsInParentSpace =
+ mapSlicedDimsToParentSpace(dimSet, sliceDims);
+ int64_t parentDim = dimsInParentSpace[0];
+
+ auto expandedParent = getParent().expandDims(parentDim, targetShape);
+
+ int64_t shift = static_cast<int64_t>(targetShape.size()) - 1;
+ SmallVector<int64_t> newSliceDims;
+ newSliceDims.reserve(sliceDims.size());
+ for (int64_t s : sliceDims)
+ newSliceDims.push_back(s > parentDim ? s + shift : s);
+
+ return SliceAttr::get(getContext(), expandedParent,
+ DenseI64ArrayAttr::get(getContext(), newSliceDims));
+}
+
SmallVector<int64_t> getPermForParentLayout(ArrayRef<int64_t> sliceDims,
ArrayRef<int64_t> permutation) {
SmallVector<int64_t> sortedSliceDims = llvm::to_vector(sliceDims);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index b008a88ca9e8f..42064b309281f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -685,37 +685,6 @@ xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return resLayout;
}
-/// Walk `srcDims` in the requested direction and assign each src dim
-/// `min(remaining, perDimCap[d])`, spilling the leftover into the next dim.
-/// `srcDims` is always given outer-to-inner; `innerToOuter` selects the
-/// iteration direction. `perDimCap[d]` is the maximum value that may be
-/// placed on src dim `d`; callers compute it as either the full available
-/// extent (e.g. `srcShape[d]`) or the per-sg / per-lane share of that
-/// extent (e.g. `srcShape[d] / inferredSgLayout[d]`) when a layout has
-/// already been placed on that dim.
-static void distributeAcrossSrcDims(int64_t total, ArrayRef<int64_t> srcDims,
- bool innerToOuter,
- ArrayRef<int64_t> perDimCap,
- SmallVectorImpl<int64_t> &out) {
- int64_t remaining = total;
- auto step = [&](int64_t d) {
- if (remaining == 1)
- return;
- int64_t take = std::min(remaining, perDimCap[d]);
- assert(take > 0 && "distribution must not be zero");
- assert(remaining % take == 0 && "must divide evenly across dims");
- out[d] = take;
- remaining /= take;
- };
- if (innerToOuter)
- for (int64_t d : llvm::reverse(srcDims))
- step(d);
- else
- for (int64_t d : srcDims)
- step(d);
- assert(remaining == 1 && "total must fit within collapsed src dims");
-}
-
/// Infers the source layout attribute for a shape cast operation given the
/// result layout attribute, result shape, and source shape.
xegpu::DistributeLayoutAttr
@@ -758,216 +727,39 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
// Use case 3: General dim collapse, for cross-sg reduction to SLM and other
// shape casts where consecutive src dims fold into a single dst dim.
- // For each dst dim that is produced by collapsing >=2 src dims, distribute
- // the consumer layout/data of that dst dim across the src dim group so that
- // each sg/lane owns a contiguous run in the collapsed dst dim:
- // - sg_layout: distribute starting from the OUTERMOST src dim of the
- // group, spreading INWARD. Each dim takes min(remaining, srcShape[d]);
- // leftover spills into the next inner dim. Subgroup partitioning along
- // the slowest-varying axis keeps each sg's tile contiguous in the dst.
- // - lane_layout: distribute starting from the OUTERMOST src dim of the
- // group, spreading INWARD. Each dim takes min(remaining, srcShape[d]);
- // leftover spills into the next inner dim. (Same direction as
- // sg_layout; lane_data is the one that fills innermost-first to
- // match the dst's row-major linearization.)
- // - sg_data / inst_data / lane_data: fill from the innermost src dim
- // outward, capped per dim by srcShape[d] (or srcShape[d]/sg_layout[d] /
- // srcShape[d]/lane_layout[d] when a layout is placed on that dim).
- // Examples:
- // srcShape=[8, 16, 32], resShape=[1, 4096], inst_data=[1, 16]
- // -> inferredInstData=[1, 1, 16]
- // srcShape=[4, 8, 64], resShape=[2048], lane_layout=[16], lane_data=[2]
- // -> inferredLaneLayout=[16, 1, 1], inferredLaneData=[1, 1, 2]
- // srcShape=[8, 16, 32], resShape=[4096], sg_layout=[8], sg_data=[512]
- // -> inferredSgLayout=[8, 1, 1] (outermost holds 8), inferredSgData=
- // [1, 16, 32]
- // srcShape=[2, 8, 32], resShape=[512], sg_layout=[16], sg_data=[32]
- // -> outer dim 0 holds min(16, 2)=2, leftover 8 spills inward to dim 1:
- // inferredSgLayout=[2, 8, 1]; inferredSgData fills innermost first:
- // [1, 1, 32]
- // srcShape=[64, 4, 2], resShape=[512], lane_layout=[16], lane_data=[2]
- // -> outer dim 0 holds min(16, 64)=16: inferredLaneLayout=[16, 1, 1];
- // inferredLaneData=[1, 1, 2] (innermost dim 2 fills first).
+ //
+ // Mirrors use case 2's elegant shape: walk the dst-side groups and call
+ // a single layout-attribute primitive per group. Here the primitive is
+ // `expandDims(dim, targetShape)`, the inverse of `collapseDims`. It applies
+ // the per-field distribution policy required for a no-data-movement collapse
+ // (sg_layout/lane_layout spread outer-to-inner; sg_data/lane_data/inst_data
+ // fill innermost-first; inst_data is seeded from lane_layout * lane_data).
+ // See LayoutAttr::expandDims for the full policy.
+ //
+ // Iteration goes innermost-first (reverse dst order) so that each
+ // expandDims/dropDims call only mutates dst positions whose indices are
+ // unaffected by earlier calls.
SmallVector<SmallVector<int64_t>> collapseDims;
if (xegpu::matchDimCollapse(srcShape, resShape, collapseDims)) {
- int srcShapeSize = srcShape.size();
- auto context = resLayout.getContext();
- auto resSgLayout = resLayout.getEffectiveSgLayoutAsInt();
- auto resSgData = resLayout.getEffectiveSgDataAsInt();
- auto resInstData = resLayout.getEffectiveInstDataAsInt();
- auto resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
- auto resLaneData = resLayout.getEffectiveLaneDataAsInt();
-
- SmallVector<int64_t> inferredSgLayout(
- resSgLayout.empty() ? 0 : srcShapeSize, 1);
- SmallVector<int64_t> inferredSgData(resSgData.empty() ? 0 : srcShapeSize,
- 1);
- SmallVector<int64_t> inferredInstData(
- resInstData.empty() ? 0 : srcShapeSize, 1);
- SmallVector<int64_t> inferredLaneLayout(
- resLaneLayout.empty() ? 0 : srcShapeSize, 1);
- SmallVector<int64_t> inferredLaneData(
- resLaneData.empty() ? 0 : srcShapeSize, 1);
-
- for (size_t dstIdx = 0; dstIdx < collapseDims.size(); ++dstIdx) {
+ auto srcLayout = resLayout;
+ for (int64_t dstIdx = static_cast<int64_t>(collapseDims.size()) - 1;
+ dstIdx >= 0; --dstIdx) {
ArrayRef<int64_t> srcDims = collapseDims[dstIdx];
- if (srcDims.empty())
+ if (srcDims.empty()) {
+ // Unit dst dim with no backing src dim: drop it.
+ srcLayout = srcLayout.dropDims({dstIdx});
continue;
-
- // Order matters: each *_data step depends on the matching *_layout
- // having been computed first (the layout values are used as per-dim
- // divisor caps when distributing the data). So we interleave:
- // sg_layout -> sg_data -> lane_layout -> lane_data -> inst_data
- // sg_data / lane_data / inst_data all fill innermost-first and spill
- // outward; sg_data is capped per dim by srcShape[d]/sg_layout[d],
- // and inst_data is seeded from lane_layout*lane_data (see below).
- //
- // After sg_data is distributed, `srcShape` is rebound to point at
- // `inferredSgData` so that the subsequent lane_layout / lane_data /
- // inst_data steps see the per-subgroup tile rather than the full
- // source. This is what implicitly gives lane_data a cap of
- // sgData[d]/lane_layout[d] when sg_data is available (and the
- // original srcShape[d]/lane_layout[d] otherwise).
- //
- // When the consumer layout replicates the dst dim across
- // subgroups/lanes (i.e. sg_layout[dstIdx] * sg_data[dstIdx] >
- // resShape[dstIdx]), each subgroup/lane owns the full extent of the
- // collapsed src dims, so the per-dim cap drops the layout divisor.
-
- // Helper: build a per-src-dim cap by dividing each entry of `base` by
- // the corresponding entry of `layout` (the consumer's already-placed
- // partitioning on that dim). `layout` may be empty, meaning "no
- // partition on this dim, so the cap is just the base extent".
- auto perDimCap = [&](ArrayRef<int64_t> base,
- ArrayRef<int64_t> layout) -> SmallVector<int64_t> {
- SmallVector<int64_t> cap(base.begin(), base.end());
- if (!layout.empty())
- for (int64_t d : srcDims)
- cap[d] /= layout[d];
- return cap;
- };
-
- // sg_layout: outer-to-inner so the outermost src dim of the group fills
- // first; leftover spreads inward when a single dim can't hold the value.
- if (!resSgLayout.empty())
- distributeAcrossSrcDims(resSgLayout[dstIdx], srcDims,
- /*innerToOuter=*/false,
- /*perDimCap=*/srcShape, inferredSgLayout);
-
- // sg_data: innermost-first, capped per dim by srcShape[d]/sgLayout[d]
- // unless the dst dim is sg-replicated (then each sg owns the full
- // src extent).
- if (!resSgData.empty()) {
- bool sgReplicated =
- !resSgLayout.empty() &&
- resSgLayout[dstIdx] * resSgData[dstIdx] > resShape[dstIdx];
- SmallVector<int64_t> cap =
- sgReplicated
- ? SmallVector<int64_t>(srcShape.begin(), srcShape.end())
- : perDimCap(srcShape, inferredSgLayout);
- distributeAcrossSrcDims(resSgData[dstIdx], srcDims,
- /*innerToOuter=*/true, cap, inferredSgData);
- }
-
- // Use a per-subgroup view for the remaining steps (lane_layout /
- // lane_data / inst_data): they describe how a single subgroup's tile is
- // partitioned across lanes, so their caps must be relative to
- // inferredSgData rather than the full source. This is a local view; the
- // function-scope `srcShape` must stay pointing at the full source so the
- // next dst dim's sg_layout/sg_data steps are computed correctly.
- ArrayRef<int64_t> laneSrcShape = resSgData.empty()
- ? ArrayRef<int64_t>(srcShape)
- : ArrayRef<int64_t>(inferredSgData);
-
- // lane_layout: outer-to-inner so the outermost src dim of the group
- // fills first; leftover spreads inward when a single dim is too small.
- // Computed AFTER sg_data and BEFORE lane_data/inst_data so that the
- // following data steps can use inferredLaneLayout as their cap/seed.
- if (!resLaneLayout.empty())
- distributeAcrossSrcDims(resLaneLayout[dstIdx], srcDims,
- /*innerToOuter=*/false,
- /*perDimCap=*/laneSrcShape, inferredLaneLayout);
-
- // lane_data: innermost-first, capped per dim by
- // (per-sg) srcShape[d] / inferredLaneLayout[d], unless the dst dim is
- // lane-replicated (then each lane owns the full per-sg extent).
- if (!resLaneData.empty()) {
- bool laneReplicated =
- !resLaneLayout.empty() &&
- resLaneLayout[dstIdx] * resLaneData[dstIdx] > resShape[dstIdx];
- SmallVector<int64_t> cap =
- laneReplicated
- ? SmallVector<int64_t>(laneSrcShape.begin(), laneSrcShape.end())
- : perDimCap(laneSrcShape, inferredLaneLayout);
- distributeAcrossSrcDims(resLaneData[dstIdx], srcDims,
- /*innerToOuter=*/true, cap, inferredLaneData);
- }
-
- // inst_data[d] must be a multiple of lane_layout[d] * lane_data[d] on
- // each src dim of the collapsed group; otherwise the consumer's lane
- // partitioning cannot be evenly mapped onto the source. Seed each src
- // dim with the per-lane atomic unit (inferredLaneLayout[d] *
- // inferredLaneData[d]) -- both already computed above -- then
- // distribute the remaining factor innermost-first, capped per dim by
- // srcShape[d] / inferredInstData[d].
- if (!resInstData.empty()) {
- // When the consumer layout has no lane_layout/lane_data, there is no
- // per-lane atomic unit to seed; distribute inst_data the same way as
- // sg_data/lane_data (innermost-first, no divisor cap).
- if (resLaneLayout.empty() || resLaneData.empty()) {
- distributeAcrossSrcDims(resInstData[dstIdx], srcDims,
- /*innerToOuter=*/true,
- /*perDimCap=*/laneSrcShape, inferredInstData);
- } else {
- int64_t laneAtom = resLaneLayout[dstIdx] * resLaneData[dstIdx];
- for (int64_t d : srcDims)
- inferredInstData[d] = inferredLaneLayout[d] * inferredLaneData[d];
- int64_t remaining = resInstData[dstIdx] / laneAtom;
- for (int64_t d : llvm::reverse(srcDims)) {
- if (remaining == 1)
- break;
- int64_t cap = laneSrcShape[d] / inferredInstData[d];
- int64_t take = std::min(remaining, cap);
- assert(take > 0 && "inst_data distribution must not be zero");
- assert(remaining % take == 0 &&
- "inst_data must divide evenly across dims");
- inferredInstData[d] *= take;
- remaining /= take;
- }
- assert(remaining == 1 &&
- "inst_data must fit within collapsed src dims");
- }
}
+ if (srcDims.size() == 1)
+ // 1:1 mapping, nothing to do for this dim.
+ continue;
+ SmallVector<int64_t> targetShape;
+ targetShape.reserve(srcDims.size());
+ for (int64_t d : srcDims)
+ targetShape.push_back(srcShape[d]);
+ srcLayout = srcLayout.expandDims(dstIdx, targetShape);
}
-
- auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
- if (v.empty())
- return DenseI32ArrayAttr();
- SmallVector<int32_t> v32(v.begin(), v.end());
- return DenseI32ArrayAttr::get(context, v32);
- };
-
- // Propagate order: for each dst dim taken in dst-order (fastest first),
- // emit its collapsed src dims from innermost to outermost. Unit dst dims
- // with no backing src (empty groups) contribute nothing.
- // Example: src=[n1,n2,n3,n4,n5], dst=[m1,n3,m2] with collapse groups
- // [[0,1],[2],[3,4]] and dst order=[1,2,0] -> src order=[2,4,3,1,0].
- DenseI32ArrayAttr srcOrderAttr;
- if (DenseI32ArrayAttr resOrder = resLayout.getOrder();
- resOrder && !resOrder.empty()) {
- SmallVector<int64_t> resOrderVec = resLayout.getEffectiveOrderAsInt();
- SmallVector<int64_t> srcOrder;
- srcOrder.reserve(srcShapeSize);
- for (int64_t dstIdx : resOrderVec)
- for (int64_t d : llvm::reverse(collapseDims[dstIdx]))
- srcOrder.push_back(d);
- srcOrderAttr = toAttr(srcOrder);
- }
-
- return xegpu::LayoutAttr::get(
- context, toAttr(inferredSgLayout), toAttr(inferredSgData),
- toAttr(inferredInstData), toAttr(inferredLaneLayout),
- toAttr(inferredLaneData), srcOrderAttr);
+ return srcLayout;
}
llvm_unreachable("running into unsupported shape cast scenarios");
return nullptr;
>From 878ab6316d0a2e98f43f3e6711ed5970f89a1706 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 4 Jun 2026 04:48:01 +0000
Subject: [PATCH 5/5] git clang format
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 15 +++++++--------
1 file changed, 7 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 64c04437cd723..42d370fcf6294 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -662,8 +662,8 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
// - order: the original dim index is replaced by the expanded dim indices
// in innermost-fastest order; entries past `dim` shift up by
// `targetShape.size() - 1`.
-DistributeLayoutAttr
-LayoutAttr::expandDims(int64_t dim, ArrayRef<int64_t> targetShape) {
+DistributeLayoutAttr LayoutAttr::expandDims(int64_t dim,
+ ArrayRef<int64_t> targetShape) {
SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
SmallVector<int64_t> sgData = getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = getEffectiveInstDataAsInt();
@@ -715,8 +715,7 @@ LayoutAttr::expandDims(int64_t dim, ArrayRef<int64_t> targetShape) {
// Splice `expanded` (length expCount) into `vec` at position `dim`,
// replacing the single entry at `dim`.
- auto splice = [&](SmallVector<int64_t> &vec,
- ArrayRef<int64_t> expanded) {
+ auto splice = [&](SmallVector<int64_t> &vec, ArrayRef<int64_t> expanded) {
if (vec.empty())
return;
vec.erase(vec.begin() + dim);
@@ -826,9 +825,9 @@ LayoutAttr::expandDims(int64_t dim, ArrayRef<int64_t> targetShape) {
SmallVector<int32_t> v32(v.begin(), v.end());
return DenseI32ArrayAttr::get(getContext(), v32);
};
- return xegpu::LayoutAttr::get(
- getContext(), toAttr(sgLayout), toAttr(sgData), toAttr(instData),
- toAttr(laneLayout), toAttr(laneData), toAttr(newOrder));
+ return xegpu::LayoutAttr::get(getContext(), toAttr(sgLayout), toAttr(sgData),
+ toAttr(instData), toAttr(laneLayout),
+ toAttr(laneData), toAttr(newOrder));
}
// Derive a new layout by transpose the layout using `permutation`.
@@ -1344,7 +1343,7 @@ DistributeLayoutAttr SliceAttr::collapseDims(SmallVector<int64_t> dimGroup) {
// expanded there, and the slice dims that lie past the expanded position
// are shifted up by `targetShape.size() - 1`.
DistributeLayoutAttr SliceAttr::expandDims(int64_t dim,
- ArrayRef<int64_t> targetShape) {
+ ArrayRef<int64_t> targetShape) {
// `dim` is in slice space; map it to parent space (parent dims listed in
// `sliceDims` are removed by the slice, so the mapping always lands on a
// non-sliced parent dim).
More information about the Mlir-commits
mailing list