[Mlir-commits] [mlir] [MLIR][XeGPU] Enable WG-level mxfp GEMM via generalized shape_cast collapse inference (PR #201496)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 3 20:08:04 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
Summary
Bringing up two WG-level mxfp GEMM integration tests — simple_mxfp_gemm_quantizeA_F4 and
simple_mxfp_gemm_dequantizeB_F4 — exposed several gaps in the XeGPU layout-propagation and unroll paths
that previously kept them from compiling end-to-end. This PR lands those two tests as the motivating
workloads, plus the supporting changes:
1. A generalized shape_cast collapse layout inference — required because the mxfp lowering inserts
vector.shape_cast ops that collapse multiple src dims into a single dst dim with non-trivial sg / lane
layouts spanning across them. The previous matchCollapseToInnermostDim only covered the narrow […] →
[N] / [1, N] shape and could not infer correct source layouts for these patterns.
2. A small primitive (expandDims) on the layout attribute so the new code stays as elegant as the use
case of collapseDims.
3. Bug fixes uncovered while running these workloads end-to-end (transpose layout check, layout-attr
unroll cast crash, drop-dims order pollution).
What's in this PR
Motivating integration tests (the driving force)
- mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
- mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir
These exercise WG-level GEMM with mxfp quantization (BF16 × F4 paths). They depend on every other
change in the PR; without them, layout propagation crashes or yields conflicting layouts on the
inserted vector.shape_cast and xegpu.load_matrix / xegpu.store_matrix ops.
Generalized shape_cast collapse inference
- New utility xegpu::matchDimCollapse(srcShape, resShape, collapseDims) in XeGPUUtils.{h,cpp} — the
dual of matchSplitDimExpansion, returning per-dst-dim groups of src indices.
- inferShapeCastSourceLayout use case 3 now handles arbitrary collapse patterns:
- sg_layout / lane_layout spread outer-to-inner, so each subgroup / lane owns a contiguous run in the
collapsed dst dim's row-major linearization.
- sg_data / lane_data / inst_data fill innermost-first, with per-dim caps from any layout already placed.
- inst_data is seeded from lane_layout * lane_data per dim; the remaining factor spreads innermost-first.
- order is rewritten by walking dst order fastest-first and emitting each group's src dims innermost-fastest.
- Net effect for the mxfp tests: no data movement across sg / lane boundaries when shape_cast collapses dims.
Refactor: expandDims interface method
- Added expandDims(int64_t dim, ArrayRef<int64_t> targetShape) to the DistributeLayoutAttr interface,
with implementations on both LayoutAttr and SliceAttr. It's the rank-increasing dual of collapseDims
and bakes in the distribution policy above.
- inferShapeCastSourceLayout use case 3 now mirrors use case 2's per-group loop:
auto srcLayout = resLayout;
for (dst dim in reverse) {
if (group.empty()) srcLayout = srcLayout.dropDims({dstIdx});
else if (group.size() > 1) srcLayout = srcLayout.expandDims(dstIdx, targetShape);
}
return srcLayout;
- Replaces ~190 lines of inlined per-field distribution logic with a handful of lines.
Bug fixes uncovered while bringing up the integration tests
- LayoutAttr::isTransposeOf: corrected the per-dim check to match vector.transpose semantics (dst[i] =
src[perm[i]]); the old comparison indexed src and dst inversely.
- LayoutAttr::dropDims: stop synthesizing a default [rank-1,...,0] order when the input had none — that
synthesized order tripped collapseDims's adjacency check downstream.
- UnrollLoadMatrixOp / UnrollStoreMatrixOp: stop assuming the op's layout is always a LayoutAttr. Use
DistributeLayoutAttr and guard dropInstData() so SliceAttr / missing-layout inputs no longer crash
unrolling.
Unit-test coverage
- New shape_cast collapse coverage in both propagate-layout-subgroup.mlir and
propagate-layout-inst-data.mlir for: plain innermost collapse, layout spill across multiple src dims,
and multi-group collapse.
- Updated one lane_layout expectation in propagate-layout.mlir to reflect the generalized distribution.
Files changed
- mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td — interface + class declarations for expandDims
- mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h — declaration for matchDimCollapse
- mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp — expandDims impls, dropDims order fix, isTransposeOf fix
- mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp — refactored use case 3
- mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp — load/store_matrix unroll hardening
- mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp — matchDimCollapse impl
- mlir/test/Dialect/XeGPU/propagate-layout-{subgroup,inst-data,}.mlir — new tests / updated expectation
- mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_{dequantizeB_F4,quantizeA_F4}.mlir — new
motivating integration tests
---
Patch is 49.70 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/201496.diff
11 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td (+24)
- (modified) mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h (+9)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp (+233-9)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+40-61)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp (+7-4)
- (modified) mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp (+53)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+75)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir (+69)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+1-1)
- (added) mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir (+71)
- (added) mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir (+110)
``````````diff
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/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/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index e92b109c2223e..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) {
@@ -689,10 +883,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;
@@ -1142,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 7384f6be8d051..42064b309281f 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
@@ -719,68 +725,41 @@ 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)) {
- int srcShapeSize = srcShape.size();
- int resShapeSize = resShape.size();
- auto context = resLayout.getContext();
- 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<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 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.
+ //
+ // M...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/201496
More information about the Mlir-commits
mailing list