[Mlir-commits] [mlir] 7bcdec0 - [MLIR][XeGPU] Enable WG-level mxfp GEMM via generalized shape_cast collapse inference (#201496)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 10 12:17:43 PDT 2026
Author: Jianhui Li
Date: 2026-06-10T12:17:38-07:00
New Revision: 7bcdec0b48ee8f64c16d1c13d7940073c3cb03a7
URL: https://github.com/llvm/llvm-project/commit/7bcdec0b48ee8f64c16d1c13d7940073c3cb03a7
DIFF: https://github.com/llvm/llvm-project/commit/7bcdec0b48ee8f64c16d1c13d7940073c3cb03a7.diff
LOG: [MLIR][XeGPU] Enable WG-level mxfp GEMM via generalized shape_cast collapse inference (#201496)
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
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
Added:
mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir
mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
mlir/test/Dialect/XeGPU/propagate-layout.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
index 40edce8a60429..dabbfd4c79de2 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td
@@ -214,6 +214,23 @@ 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`, when present, is rewritten so the expanded dims appear
+ innermost-fastest in the original dim's slot; if the input has no
+ `order`, the result has no `order` either (the default FCD-first
+ order `[rank-1, ..., 0]` is assumed).
+ Example: layout<sg_layout=[2, 16], sg_data=[1, 1]> with
+ expandDim(1, [2, 8]) -> layout<sg_layout=[2, 2, 8], sg_data=[1, 1, 1]>
+ (no order in, no order out).}],
+ "xegpu::DistributeLayoutAttr",
+ "expandDim",
+ (ins "int64_t": $dim, "ArrayRef<int64_t>": $targetShape)>,
InterfaceMethod<[{Derive a new layout by trasnposing it using `permutation`.}],
"xegpu::DistributeLayoutAttr",
"transposeDims",
@@ -582,6 +599,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 expandDim(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 +837,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 expandDim(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..b780c66594eb0 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,19 +504,26 @@ 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())
- return DenseI32ArrayAttr();
+ return nullptr;
SmallVector<int32_t> v32(v.begin(), v.end());
return DenseI32ArrayAttr::get(getContext(), v32);
};
@@ -627,7 +634,7 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
auto toAttr = [&](ArrayRef<int64_t> v) -> DenseI32ArrayAttr {
if (v.empty())
- return DenseI32ArrayAttr();
+ return nullptr;
SmallVector<int32_t> v32(v.begin(), v.end());
return DenseI32ArrayAttr::get(getContext(), v32);
};
@@ -637,6 +644,207 @@ 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 is always outer-to-inner to make sure larger contiguous
+// chunks are given to each compute unit first. Concretely: sg_layout and
+// lane_layout walk the new dims outer-to-inner so each compute unit owns a
+// contiguous run after collapse; sg_data / lane_data / inst_data fill
+// innermost-first so the per-unit data tile is contiguous in the
+// fastest-varying expanded dim. The expanded dims are assumed to be in
+// row-major (FCD-first) order, which is intrinsic to `vector.shape_cast`'s
+// linear-order-preserving semantics.
+//
+// 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`.
+//
+// Examples (only the affected dim shown; assume rank-1 input):
+// layout<sg_layout=[8], sg_data=[512]>, expandDim(0, [8, 16, 32])
+// -> sg_layout=[8, 1, 1], sg_data=[1, 16, 32]
+// layout<sg_layout=[16], sg_data=[32]>, expandDim(0, [2, 8, 32])
+// -> sg_layout=[2, 8, 1], sg_data=[1, 1, 32] (sg_layout spills inward)
+// layout<inst_data=[32], lane_layout=[16], lane_data=[1]>,
+// expandDim(0, [8, 16, 32])
+// -> inst_data=[8, 2, 2], lane_layout=[8, 2, 1], lane_data=[1, 1, 1]
+DistributeLayoutAttr LayoutAttr::expandDim(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];
+
+ // Spread `total` across the new dims (length expCount), capped per dim by
+ // `dimSizeCap[i]`. `outerToInner` selects iteration direction
+ // (true = i=0..n-1).
+ auto spread = [&](int64_t total, ArrayRef<int64_t> dimSizeCap,
+ 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, dimSizeCap[i]);
+ assert(take > 0 && "expandDim 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> dimSizeCap(targetShape.begin(), targetShape.end());
+ if (hasSgLayout)
+ for (int64_t i = 0; i < expCount; ++i)
+ dimSizeCap[i] /= expSgLayout[i];
+ SmallVector<int64_t> expSgData =
+ spread(origSgDataDim, dimSizeCap, /*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> dimSizeCap(perSgShape.begin(), perSgShape.end());
+ if (hasLaneLayout)
+ for (int64_t i = 0; i < expCount; ++i)
+ dimSizeCap[i] /= expLaneLayout[i];
+ expLaneData = spread(origLaneDataDim, dimSizeCap, /*outerToInner=*/false);
+ splice(laneData, expLaneData);
+ }
+
+ // inst_data: when lane info is present, the per-lane atom
+ // `laneLayout[i] * laneData[i]` is the minimum granularity each new dim must
+ // hold to keep the lane-level distribution unit aligned with inst_data; the
+ // remaining factor `inst_data / laneAtom` is then spread innermost-first
+ // (capped by `perSgShape[i] / atom[i]`) and multiplied back onto the atom.
+ // Without lane info, fall back to a plain innermost-first spread over
+ // perSgShape.
+ if (hasInstData) {
+ SmallVector<int64_t> expInstData;
+ if (!hasLaneLayout || !hasLaneData) {
+ expInstData = spread(origInstDataDim, perSgShape, /*outerToInner=*/false);
+ } else {
+ int64_t laneAtom = origLaneLayoutDim * origLaneDataDim;
+ SmallVector<int64_t> atom(expCount, 1);
+ SmallVector<int64_t> dimSizeCap(expCount, 1);
+ for (int64_t i = 0; i < expCount; ++i) {
+ atom[i] = expLaneLayout[i] * expLaneData[i];
+ dimSizeCap[i] = perSgShape[i] / atom[i];
+ }
+ expInstData = spread(origInstDataDim / laneAtom, dimSizeCap,
+ /*outerToInner=*/false);
+ for (int64_t i = 0; i < expCount; ++i)
+ expInstData[i] *= atom[i];
+ }
+ 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 nullptr;
+ 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 +897,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 +1353,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::expandDim(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).
+ ArrayRef<int64_t> sliceDims = getDims().asArrayRef();
+ SmallVector<int64_t> dimSet = {dim};
+ SmallVector<int64_t> dimsInParentSpace =
+ mapSlicedDimsToParentSpace(dimSet, sliceDims);
+ int64_t parentDim = dimsInParentSpace[0];
+
+ auto expandedParent = getParent().expandDim(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 11b36f56efa30..4b821fce3e40c 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
@@ -611,7 +617,7 @@ xegpu::inferInsertSourceLayout(xegpu::DistributeLayoutAttr resLayout,
/// Infers the source layout attribute for extract operation
/// given the result layout attribute, result shape, and source shape. Adds
/// leading dimensions to the source layout to match the source shape size.
-// TODO: add layout attribute interface: expandDims() and use it here.
+// TODO: add layout attribute interface: expandDim() and use it here.
// TODO: add propagation support for extract op
xegpu::DistributeLayoutAttr
xegpu::inferExtractSourceLayout(xegpu::DistributeLayoutAttr resLayout,
@@ -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.
+ //
+ // Mirrors use case 2's elegant shape: walk the dst-side groups and call
+ // a single layout-attribute primitive per group. Here the primitive is
+ // `expandDim(dim, targetShape)`, the inverse of `collapseDims`. It applies
+ // the per-field distribution policy required for a no-data-movement collapse
+ // (sg_layout/lane_layout spread outer-to-inner; sg_data/lane_data/inst_data
+ // fill innermost-first; inst_data is seeded from lane_layout * lane_data).
+ // See LayoutAttr::expandDim for the full policy.
+ //
+ // Iteration goes innermost-first (reverse dst order) so that each
+ // expandDim/dropDims call only mutates dst positions whose indices are
+ // unaffected by earlier calls.
+ SmallVector<SmallVector<int64_t>> collapseDims;
+ if (xegpu::matchDimCollapse(srcShape, resShape, collapseDims)) {
+ auto srcLayout = resLayout;
+ for (int64_t dstIdx = static_cast<int64_t>(collapseDims.size()) - 1;
+ dstIdx >= 0; --dstIdx) {
+ ArrayRef<int64_t> srcDims = collapseDims[dstIdx];
+ if (srcDims.empty()) {
+ // Unit dst dim with no backing src dim: drop it.
+ srcLayout = srcLayout.dropDims({dstIdx});
+ continue;
}
- 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);
+ 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.expandDim(dstIdx, targetShape);
}
+ return srcLayout;
}
llvm_unreachable("running into unsupported shape cast scenarios");
return nullptr;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index f84d29aa51164..fc7c3b170dd3b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -882,7 +882,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);
@@ -897,7 +897,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);
@@ -921,7 +922,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);
@@ -940,7 +943,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/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index e4a085bdde6d3..4b9d0ceb5350a 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -1002,3 +1002,85 @@ 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, in-between, or trailing) get empty groups; src unit dims that
+// fall past the last consumed dst dim are absorbed into the most-recent
+// non-empty group.
+// Examples:
+// src=[8,16,32], dst=[1,4096] -> true, collapseDims=[[],[0,1,2]]
+// src=[8,16,32], dst=[4096,1] -> 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());
+
+ // Cheap precondition: src and dst must describe the same number of
+ // elements. Bails out early on mismatched shapes without walking the dims.
+ int64_t srcProd = std::accumulate(src.begin(), src.end(), int64_t{1},
+ std::multiplies<int64_t>());
+ int64_t dstProd = std::accumulate(dst.begin(), dst.end(), int64_t{1},
+ std::multiplies<int64_t>());
+ if (srcProd != dstProd)
+ return false;
+
+ // Step 1: validate the partition on the unit-dim-stripped (compact) shapes.
+ // Unit dims play no role in the matching decision — they only need to be
+ // placed somewhere in the final groups (handled in step 2).
+ SmallVector<int64_t> srcCompact, dstCompact;
+ for (int64_t s : src)
+ if (s != 1)
+ srcCompact.push_back(s);
+ for (int64_t d : dst)
+ if (d != 1)
+ dstCompact.push_back(d);
+
+ size_t s = 0;
+ for (int64_t need : dstCompact) {
+ int64_t acc = 1;
+ while (s < srcCompact.size() && acc < need)
+ acc *= srcCompact[s++];
+ if (acc != need)
+ return false;
+ }
+ if (s != srcCompact.size())
+ return false;
+
+ // Step 2: assign each original src index to the correct original dst group.
+ // Walk dst in original order, advancing past unit dst dims (they keep their
+ // pre-initialized empty group). Walk src in original order; non-unit src
+ // dims accumulate into the current dst group, unit src dims attach to the
+ // current group when one is open or to the most-recent non-empty group
+ // after dst is exhausted (leading unit src dims with no group yet are
+ // dropped).
+ size_t dstIdx = 0;
+ while (dstIdx < dst.size() && dst[dstIdx] == 1)
+ dstIdx++;
+
+ int64_t lastNonEmpty = -1;
+ int64_t acc = 1;
+ for (size_t srcIdx = 0; srcIdx < src.size(); ++srcIdx) {
+ if (dstIdx >= dst.size()) {
+ // dst exhausted; remaining src dims are unit (validated above) and
+ // attach to the last non-empty group, if any.
+ if (lastNonEmpty >= 0)
+ collapseDims[lastNonEmpty].push_back(srcIdx);
+ continue;
+ }
+ acc *= src[srcIdx];
+ collapseDims[dstIdx].push_back(srcIdx);
+ lastNonEmpty = dstIdx;
+ if (acc == dst[dstIdx]) {
+ acc = 1;
+ ++dstIdx;
+ while (dstIdx < dst.size() && dst[dstIdx] == 1)
+ ++dstIdx;
+ }
+ }
+ return true;
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 6f587959b697d..5f493c8ca0df6 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -530,3 +530,80 @@ 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).
+// Distribution is always outer-to-inner to make sure larger contiguous chunks
+// are given to each compute unit first.
+// 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
+ }
+}
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..6325aaaf69a4e
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_dequantizeB_F4.mlir
@@ -0,0 +1,203 @@
+// 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: *
+// Note: layouts used by dpas_mx need to match HW constaint. Otherwise dpas_mx is not unrolled.
+#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#b_packed = #xegpu.layout<sg_layout = [2, 2], sg_data = [512, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
+#b = #xegpu.layout<sg_layout = [2, 2], sg_data = [1024, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+#b_f16 = #xegpu.layout<sg_layout = [2, 2], sg_data = [1024, 16], inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>
+#c = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// Note: inst_data is chosen to utilize 2D block load
+#b_scale = #xegpu.layout<sg_layout = [2, 2], sg_data = [32, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// Note: scales for dpas_mx needs separate layouts with inst_data to match HW constraint. Otherwise dpas_mx is not unrolled
+
+
+module @gemm attributes {gpu.container_module} {
+ gpu.module @kernel {
+ // 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<256x4096xbf16>, %arg1: memref<2048x256xi8>, %arg3: memref<128x256xf8E8M0FNU>, %arg4: memref<256x256xf32>) kernel {
+ %c0 = arith.constant 0 : index
+ %mstep = arith.constant 32 : index
+ %nstep = arith.constant 32 : index
+ %kstep = arith.constant 1024 : index
+ %mbound = arith.constant 256 : index
+ %nbound = arith.constant 256 : index
+ %kbound = arith.constant 4096 : index
+ %kbstep = arith.constant 512 : index
+ %kscalestep = arith.constant 32 : index
+ %block_id_x = gpu.block_id x
+ %block_id_y = gpu.block_id y
+ %m = arith.muli %block_id_x, %mstep : index
+ %n = arith.muli %block_id_y, %nstep : index
+
+ %a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<256x4096xbf16> -> !xegpu.tensor_desc<32x1024xbf16>
+ %bp_tdesc = xegpu.create_nd_tdesc %arg1 : memref<2048x256xi8> -> !xegpu.tensor_desc<512x32xi8>
+ %b_scale_tdesc = xegpu.create_nd_tdesc %arg3 : memref<128x256xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU>
+
+ // Load initial C
+ %cd_tdesc = xegpu.create_nd_tdesc %arg4 : memref<256x256xf32> -> !xegpu.tensor_desc<32x32xf32, #c>
+ %c_init = xegpu.load_nd %cd_tdesc[%m, %n] {layout = #c}: !xegpu.tensor_desc<32x32xf32, #c> -> vector<32x32xf32>
+
+ %res:3 = scf.for %k = %c0 to %kbound step %kstep
+ iter_args(%c_partial = %c_init, %kb = %c0, %kscale = %c0) -> (vector<32x32xf32>, index, index) {
+ // -------- Load A (bf16) --------
+ %a = xegpu.load_nd %a_tdesc[%m, %k] {layout = #a}: !xegpu.tensor_desc<32x1024xbf16> -> vector<32x1024xbf16>
+
+ %bp = xegpu.load_nd %bp_tdesc[%kb, %n] {layout = #b_packed}: !xegpu.tensor_desc<512x32xi8> -> vector<512x32xi8>
+
+ // Bitcast to fp4: 512x32 uint8 -> 512x64 fp4 (each uint8 holds 2 fp4 values)
+ %b_bitcast = vector.bitcast %bp : vector<512x32xi8> to vector<512x64xf4E2M1FN>
+
+ // De-interleave: extract even and odd columns
+ // Even columns (indices 0, 2, 4, ..., 62) -> first half
+ // Odd columns (indices 1, 3, 5, ..., 63) -> second half
+ %b_even, %b_odd = vector.deinterleave %b_bitcast : vector<512x64xf4E2M1FN> -> vector<512x32xf4E2M1FN>
+
+ // Reconstruct 1024x32 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<512x32xf4E2M1FN> to vector<32x512xf4E2M1FN>
+ %b_odd_t = vector.transpose %b_odd, [1, 0] : vector<512x32xf4E2M1FN> to vector<32x512xf4E2M1FN>
+ %b_interleaved = vector.interleave %b_even_t, %b_odd_t : vector<32x512xf4E2M1FN> -> vector<32x1024xf4E2M1FN>
+ %b = vector.transpose %b_interleaved, [1, 0] : vector<32x1024xf4E2M1FN> to vector<1024x32xf4E2M1FN>
+
+
+ %scale_b = xegpu.load_nd %b_scale_tdesc[%kscale, %n] {layout = #b_scale}: !xegpu.tensor_desc<32x32xf8E8M0FNU> -> vector<32x32xf8E8M0FNU>
+ // 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<32x32xf8E8M0FNU> to vector<32x32x32xf8E8M0FNU>
+ %scale_b_t = vector.transpose %scale_b_bcast, [1, 0, 2] : vector<32x32x32xf8E8M0FNU> to vector<32x32x32xf8E8M0FNU>
+ %scale_b_full = vector.shape_cast %scale_b_t : vector<32x32x32xf8E8M0FNU> to vector<1024x32xf8E8M0FNU>
+
+ // Dequantize B from f4E2M1FN to bf16 using scale_b.
+ %b_bf16 = arith.scaling_extf %b, %scale_b_full : vector<1024x32xf4E2M1FN>, vector<1024x32xf8E8M0FNU> to vector<1024x32xbf16>
+
+ %new_c_partial = xegpu.dpas %a, %b_bf16, %c_partial
+ {layout_a = #a,
+ layout_b = #b_f16,
+ layout_cd = #c}
+ : vector<32x1024xbf16>, vector<1024x32xbf16>,
+ vector<32x32xf32>
+ -> vector<32x32xf32>
+
+ // b and b_scale take
diff erent steps compared to a
+ // compute adjusted k index for those tiles.
+ %new_kb = arith.addi %kb, %kbstep : index
+ %new_kscale = arith.addi %kscale, %kscalestep : index
+ scf.yield %new_c_partial, %new_kb, %new_kscale : vector<32x32xf32>, index, index
+ }
+
+ // store_nd with offset
+ xegpu.store_nd %res#0, %cd_tdesc[%m, %n] {layout = #c} : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32, #c>
+ gpu.return
+ }
+ }
+
+ func.func @test(%a: memref<256x4096xbf16>, %b: memref<2048x256xi8>, %b_scale: memref<128x256xf8E8M0FNU>, %c: memref<256x256xf32>) -> memref<256x256xf32> attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c64 = arith.constant 64 : index
+
+ %memref_a = gpu.alloc() : memref<256x4096xbf16>
+ gpu.memcpy %memref_a, %a : memref<256x4096xbf16>, memref<256x4096xbf16>
+
+ %memref_b = gpu.alloc() : memref<2048x256xi8>
+ gpu.memcpy %memref_b, %b : memref<2048x256xi8>, memref<2048x256xi8>
+
+ %memref_c = gpu.alloc() : memref<256x256xf32>
+ gpu.memcpy %memref_c, %c : memref<256x256xf32>, memref<256x256xf32>
+
+ %memref_b_scale = gpu.alloc() : memref<128x256xf8E8M0FNU>
+ gpu.memcpy %memref_b_scale, %b_scale : memref<128x256xf8E8M0FNU>, memref<128x256xf8E8M0FNU>
+
+ gpu.launch_func @kernel::@gemm_mxfp blocks in (%c8, %c8, %c1) threads in (%c64, %c1, %c1)
+ args(%memref_a : memref<256x4096xbf16>, %memref_b : memref<2048x256xi8>, %memref_b_scale : memref<128x256xf8E8M0FNU>, %memref_c : memref<256x256xf32>)
+ gpu.dealloc %memref_a : memref<256x4096xbf16>
+ gpu.dealloc %memref_b : memref<2048x256xi8>
+ gpu.dealloc %memref_b_scale : memref<128x256xf8E8M0FNU>
+
+ %res = memref.alloc() : memref<256x256xf32>
+ gpu.memcpy %res, %memref_c : memref<256x256xf32>, memref<256x256xf32>
+ gpu.dealloc %memref_c : memref<256x256xf32>
+ return %res : memref<256x256xf32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %c256 = arith.constant 256 : index
+ %c2K = arith.constant 2048 : index
+ %c4K = arith.constant 4096 : index
+ %c512K = arith.constant 524288 : index
+ %c1bf16 = arith.constant 1.0 : bf16
+ %c1packed_e2m1 = arith.constant 0x22 : i8
+ %c0f32 = arith.constant 0.0 : f32
+ %c1f8E8M0FNU = arith.constant 1.0 : f8E8M0FNU
+
+ %A = memref.alloc() : memref<256x4096xbf16>
+ scf.for %i = %c0 to %c256 step %c1 {
+ scf.for %j = %c0 to %c4K step %c1 {
+ memref.store %c1bf16, %A[%i,%j] : memref<256x4096xbf16>
+ }
+ }
+
+ %B = memref.alloc() : memref<2048x256xi8>
+ scf.for %i = %c0 to %c2K step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c1packed_e2m1, %B[%i, %j] : memref<2048x256xi8>
+ }
+ }
+
+ %C = memref.alloc() : memref<256x256xf32>
+ scf.for %i = %c0 to %c256 step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c0f32, %C[%i, %j] : memref<256x256xf32>
+ }
+ }
+
+ %B_scale = memref.alloc() : memref<128x256xf8E8M0FNU>
+ scf.for %i = %c0 to %c128 step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c1f8E8M0FNU, %B_scale[%i, %j] : memref<128x256xf8E8M0FNU>
+ }
+ }
+
+
+ %c4Kf = arith.constant 4096.0 : f32
+ %C_ref = memref.alloc() : memref<256x256xf32>
+ scf.for %i = %c0 to %c256 step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c4Kf, %C_ref[%i, %j] : memref<256x256xf32>
+ }
+ }
+
+ %C_res = call @test(%A, %B, %B_scale, %C) : (memref<256x4096xbf16>, memref<2048x256xi8>, memref<128x256xf8E8M0FNU>, memref<256x256xf32>) -> memref<256x256xf32>
+ %C_cast = memref.cast %C_res : memref<256x256xf32> to memref<*xf32>
+ %C_ref_cast = memref.cast %C_ref : memref<256x256xf32> to memref<*xf32>
+ %
diff = call @verifyMemRefF32(%C_cast, %C_ref_cast) : (memref<*xf32>, memref<*xf32>) -> i64
+ call @printI64(%
diff ) : (i64) -> ()
+ //call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
+
+ // CHECK: 0
+ memref.dealloc %A : memref<256x4096xbf16>
+ memref.dealloc %B : memref<2048x256xi8>
+ memref.dealloc %B_scale : memref<128x256xf8E8M0FNU>
+ memref.dealloc %C : memref<256x256xf32>
+ memref.dealloc %C_res : memref<256x256xf32>
+ return
+ }
+ func.func private @verifyMemRefF32(%acutal : memref<*xf32>, %expected : memref<*xf32>) -> i64 attributes { llvm.emit_c_interface }
+ func.func private @printI64(%num : i64)
+ //func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }
+
+}
\ No newline at end of file
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..cfaa8077c742c
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
@@ -0,0 +1,235 @@
+// 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: *
+// Note: layouts used by dpas_mx need to match HW constaint. Otherwise dpas_mx is not unrolled.
+#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#a_ld = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#b_packed = #xegpu.layout<sg_layout = [2, 2], sg_data = [512, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
+#b = #xegpu.layout<sg_layout = [2, 2], sg_data = [1024, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+#c = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// Note: inst_data is chosen to utilize 2D block load
+#b_scale = #xegpu.layout<sg_layout = [2, 2], sg_data = [32, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+// Note: scales for dpas_mx needs separate layouts with inst_data to match HW constraint. Otherwise dpas_mx is not unrolled
+#dpas_a_scale = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 32], inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>
+#dpas_b_scale = #xegpu.layout<sg_layout = [2, 2], sg_data = [32, 16], inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+
+
+module @gemm attributes {gpu.container_module} {
+ gpu.module @kernel {
+ // 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<256x4096xbf16>, %arg1: memref<2048x256xi8>, %arg3: memref<128x256xf8E8M0FNU>, %arg4: memref<256x256xf32>) kernel {
+ %c0 = arith.constant 0 : index
+ %mstep = arith.constant 32 : index
+ %nstep = arith.constant 32 : index
+ %kstep = arith.constant 1024 : index
+ %mbound = arith.constant 256 : index
+ %nbound = arith.constant 256 : index
+ %kbound = arith.constant 4096 : index
+ %kbstep = arith.constant 512 : index
+ %kscalestep = arith.constant 32 : index
+ %block_id_x = gpu.block_id x
+ %block_id_y = gpu.block_id y
+ %m = arith.muli %block_id_x, %mstep : index
+ %n = arith.muli %block_id_y, %nstep : index
+
+ %a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<256x4096xbf16> -> !xegpu.tensor_desc<32x1024xbf16>
+ %bp_tdesc = xegpu.create_nd_tdesc %arg1 : memref<2048x256xi8> -> !xegpu.tensor_desc<512x32xi8>
+ %b_scale_tdesc = xegpu.create_nd_tdesc %arg3 : memref<128x256xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU>
+
+ // Load initial C
+ %cd_tdesc = xegpu.create_nd_tdesc %arg4 : memref<256x256xf32> -> !xegpu.tensor_desc<32x32xf32, #c>
+ %c_init = xegpu.load_nd %cd_tdesc[%m, %n] {layout = #c}: !xegpu.tensor_desc<32x32xf32, #c> -> vector<32x32xf32>
+
+ %res:3 = scf.for %k = %c0 to %kbound step %kstep
+ iter_args(%c_partial = %c_init, %kb = %c0, %kscale = %c0) -> (vector<32x32xf32>, index, index) {
+ // -------- Load A (bf16) --------
+ %a_bf16 = xegpu.load_nd %a_tdesc[%m, %k] {layout = #a_ld}: !xegpu.tensor_desc<32x1024xbf16> -> vector<32x1024xbf16>
+
+ // -------- 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<32x1024xbf16>
+ %a_abs_r = vector.shape_cast %a_abs : vector<32x1024xbf16> to vector<32x32x32xbf16>
+ %a_neg_inf_i = arith.constant dense<0xFF80> : vector<32x32xi16>
+ %a_neg_inf = arith.bitcast %a_neg_inf_i : vector<32x32xi16> to vector<32x32xbf16>
+ %a_amax = vector.multi_reduction <maximumf>, %a_abs_r, %a_neg_inf [2]
+ : vector<32x32x32xbf16> to vector<32x32xbf16>
+
+ // 2) Largest power-of-two <= amax: mask out mantissa bits of bf16.
+ %a_amax_i16 = arith.bitcast %a_amax : vector<32x32xbf16> to vector<32x32xi16>
+ %a_exp_mask = arith.constant dense<0x7F80> : vector<32x32xi16>
+ %a_pow2_i16 = arith.andi %a_amax_i16, %a_exp_mask : vector<32x32xi16>
+ %a_pow2 = arith.bitcast %a_pow2_i16 : vector<32x32xi16> to vector<32x32xbf16>
+
+ // 3) Divide by largest power-of-two representable by E2M1 (= 4.0).
+ %a_e2m1_max = arith.constant dense<4.000000e+00> : vector<32x32xbf16>
+ %a_scale_bf16 = arith.divf %a_pow2, %a_e2m1_max : vector<32x32xbf16>
+
+ // 4) Truncate scale to f8E8M0FNU.
+ %a_scale = arith.truncf %a_scale_bf16 : vector<32x32xbf16> to vector<32x32xf8E8M0FNU>
+
+ // 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<32x32xf8E8M0FNU> to vector<32x32x32xf8E8M0FNU>
+ %a_scale_t = vector.transpose %a_scale_lead, [1, 2, 0]
+ : vector<32x32x32xf8E8M0FNU> to vector<32x32x32xf8E8M0FNU>
+ %a_scale_full = vector.shape_cast %a_scale_t
+ : vector<32x32x32xf8E8M0FNU> to vector<32x1024xf8E8M0FNU>
+
+ // 6) Scaled truncf to fp4 (to_nearest_even).
+ %a = arith.scaling_truncf %a_bf16, %a_scale_full
+ : vector<32x1024xbf16>, vector<32x1024xf8E8M0FNU> to vector<32x1024xf4E2M1FN>
+
+ %bp = xegpu.load_nd %bp_tdesc[%kb, %n] {layout = #b_packed}: !xegpu.tensor_desc<512x32xi8> -> vector<512x32xi8>
+
+ // Bitcast to fp4: 512x32 uint8 -> 512x64 fp4 (each uint8 holds 2 fp4 values)
+ %b_bitcast = vector.bitcast %bp : vector<512x32xi8> to vector<512x64xf4E2M1FN>
+
+ // De-interleave: extract even and odd columns
+ // Even columns (indices 0, 2, 4, ..., 62) -> first half
+ // Odd columns (indices 1, 3, 5, ..., 63) -> second half
+ %b_even, %b_odd = vector.deinterleave %b_bitcast : vector<512x64xf4E2M1FN> -> vector<512x32xf4E2M1FN>
+
+ // Reconstruct 1024x32 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<512x32xf4E2M1FN> to vector<32x512xf4E2M1FN>
+ %b_odd_t = vector.transpose %b_odd, [1, 0] : vector<512x32xf4E2M1FN> to vector<32x512xf4E2M1FN>
+ %b_interleaved = vector.interleave %b_even_t, %b_odd_t : vector<32x512xf4E2M1FN> -> vector<32x1024xf4E2M1FN>
+ %b = vector.transpose %b_interleaved, [1, 0] : vector<32x1024xf4E2M1FN> to vector<1024x32xf4E2M1FN>
+
+
+ %scale_b = xegpu.load_nd %b_scale_tdesc[%kscale, %n] {layout = #b_scale}: !xegpu.tensor_desc<32x32xf8E8M0FNU> -> vector<32x32xf8E8M0FNU>
+ %new_c_partial = xegpu.dpas_mx %a, %b, %c_partial scale_a = %a_scale scale_b = %scale_b
+ {layout_a = #a,
+ layout_b = #b,
+ layout_cd = #c,
+ layout_a_scale = #dpas_a_scale,
+ layout_b_scale = #dpas_b_scale}
+ : (vector<32x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>,
+ vector<32x32xf32>,
+ vector<32x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>)
+ -> vector<32x32xf32>
+
+ // b, a_scale and b_scale take
diff erent steps compared to a
+ // compute adjusted k index for those tiles.
+ %new_kb = arith.addi %kb, %kbstep : index
+ %new_kscale = arith.addi %kscale, %kscalestep : index
+ scf.yield %new_c_partial, %new_kb, %new_kscale : vector<32x32xf32>, index, index
+ }
+
+ // store_nd with offset
+ xegpu.store_nd %res#0, %cd_tdesc[%m, %n] {layout = #c} : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32, #c>
+ gpu.return
+ }
+ }
+
+ func.func @test(%a: memref<256x4096xbf16>, %b: memref<2048x256xi8>, %b_scale: memref<128x256xf8E8M0FNU>, %c: memref<256x256xf32>) -> memref<256x256xf32> attributes {llvm.emit_c_interface} {
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %c64 = arith.constant 64 : index
+
+ %memref_a = gpu.alloc() : memref<256x4096xbf16>
+ gpu.memcpy %memref_a, %a : memref<256x4096xbf16>, memref<256x4096xbf16>
+
+ %memref_b = gpu.alloc() : memref<2048x256xi8>
+ gpu.memcpy %memref_b, %b : memref<2048x256xi8>, memref<2048x256xi8>
+
+ %memref_c = gpu.alloc() : memref<256x256xf32>
+ gpu.memcpy %memref_c, %c : memref<256x256xf32>, memref<256x256xf32>
+
+ %memref_b_scale = gpu.alloc() : memref<128x256xf8E8M0FNU>
+ gpu.memcpy %memref_b_scale, %b_scale : memref<128x256xf8E8M0FNU>, memref<128x256xf8E8M0FNU>
+
+ gpu.launch_func @kernel::@gemm_mxfp blocks in (%c8, %c8, %c1) threads in (%c64, %c1, %c1)
+ args(%memref_a : memref<256x4096xbf16>, %memref_b : memref<2048x256xi8>, %memref_b_scale : memref<128x256xf8E8M0FNU>, %memref_c : memref<256x256xf32>)
+ gpu.dealloc %memref_a : memref<256x4096xbf16>
+ gpu.dealloc %memref_b : memref<2048x256xi8>
+ gpu.dealloc %memref_b_scale : memref<128x256xf8E8M0FNU>
+
+ %res = memref.alloc() : memref<256x256xf32>
+ gpu.memcpy %res, %memref_c : memref<256x256xf32>, memref<256x256xf32>
+ gpu.dealloc %memref_c : memref<256x256xf32>
+ return %res : memref<256x256xf32>
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %c256 = arith.constant 256 : index
+ %c2K = arith.constant 2048 : index
+ %c4K = arith.constant 4096 : index
+ %c512K = arith.constant 524288 : index
+ %c1bf16 = arith.constant 1.0 : bf16
+ %c1packed_e2m1 = arith.constant 0x22 : i8
+ %c0f32 = arith.constant 0.0 : f32
+ %c1f8E8M0FNU = arith.constant 1.0 : f8E8M0FNU
+
+ %A = memref.alloc() : memref<256x4096xbf16>
+ scf.for %i = %c0 to %c256 step %c1 {
+ scf.for %j = %c0 to %c4K step %c1 {
+ memref.store %c1bf16, %A[%i,%j] : memref<256x4096xbf16>
+ }
+ }
+
+ %B = memref.alloc() : memref<2048x256xi8>
+ scf.for %i = %c0 to %c2K step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c1packed_e2m1, %B[%i, %j] : memref<2048x256xi8>
+ }
+ }
+
+ %C = memref.alloc() : memref<256x256xf32>
+ scf.for %i = %c0 to %c256 step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c0f32, %C[%i, %j] : memref<256x256xf32>
+ }
+ }
+
+ %B_scale = memref.alloc() : memref<128x256xf8E8M0FNU>
+ scf.for %i = %c0 to %c128 step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c1f8E8M0FNU, %B_scale[%i, %j] : memref<128x256xf8E8M0FNU>
+ }
+ }
+
+
+ %c4Kf = arith.constant 4096.0 : f32
+ %C_ref = memref.alloc() : memref<256x256xf32>
+ scf.for %i = %c0 to %c256 step %c1 {
+ scf.for %j = %c0 to %c256 step %c1 {
+ memref.store %c4Kf, %C_ref[%i, %j] : memref<256x256xf32>
+ }
+ }
+
+ %C_res = call @test(%A, %B, %B_scale, %C) : (memref<256x4096xbf16>, memref<2048x256xi8>, memref<128x256xf8E8M0FNU>, memref<256x256xf32>) -> memref<256x256xf32>
+ %C_cast = memref.cast %C_res : memref<256x256xf32> to memref<*xf32>
+ %C_ref_cast = memref.cast %C_ref : memref<256x256xf32> to memref<*xf32>
+ %
diff = call @verifyMemRefF32(%C_cast, %C_ref_cast) : (memref<*xf32>, memref<*xf32>) -> i64
+ call @printI64(%
diff ) : (i64) -> ()
+ //call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()
+
+ // CHECK: 0
+ memref.dealloc %A : memref<256x4096xbf16>
+ memref.dealloc %B : memref<2048x256xi8>
+ memref.dealloc %B_scale : memref<128x256xf8E8M0FNU>
+ memref.dealloc %C : memref<256x256xf32>
+ memref.dealloc %C_res : memref<256x256xf32>
+ return
+ }
+ func.func private @verifyMemRefF32(%acutal : memref<*xf32>, %expected : memref<*xf32>) -> i64 attributes { llvm.emit_c_interface }
+ func.func private @printI64(%num : i64)
+ //func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }
+
+}
\ No newline at end of file
More information about the Mlir-commits
mailing list