[Mlir-commits] [mlir] [MLIR][XeGPU] Fix layout inference issues blocking MXFP_GEMM test (PR #196243)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 6 23:23:00 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
This branch fixes layout inference issues in XeGPU passes that were blocking MXFP (microscaled floating point) GEMM workloads:
- Fix bitcast/interleave layout adjustment to use result shape instead of source shape. The setupBitCastResultLayout and setupInterleaveResultLayout functions were incorrectly bounding the layout adjustment loop against the source shape.
- Fix blocking pass to drop inst_data from anchor operations. Operations whose shape already matches inst_data don't get unrolled, so their layout attributes retained stale inst_data that broke downstream passes. Now inst_data is unconditionally stripped from all op attributes after blocking.
- Propagate layout to both results of vector.deinterleave. The layout recovery pass was only setting the layout on result 0, leaving result 1 without a layout.
Test plan
- Added mlir/test/Integration/Dialect/XeGPU/WG/simple_ mxfp_gemm.mlir integration test exercising the full
MXFP GEMM pipeline (bitcast, deinterleave, transpose, interleave, dpas_mx).
Assisted by Claude
---
Full diff: https://github.com/llvm/llvm-project/pull/196243.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp (+7)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+14-10)
- (added) mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir (+73)
``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 7db887915b275..8804e5c9919f2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -488,6 +488,13 @@ void XeGPUBlockingPass::runOnOperation() {
}
}
+ // Drop inst_data from operation attributes (e.g., layout, layout_a,
+ // layout_b, etc.) This is necessary for anchor operations that don't get
+ // unrolled because their inst_data already matches their shape.
+ SmallVector<NamedAttribute> newAttrs =
+ xegpu::dropInstDataOnAttrs(op->getAttrs());
+ op->setAttrs(newAttrs);
+
// Resolve unrealized conversion cast ops emulating pack/unpack
if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op))
resolveUnrealizedConversionCastOp(castOp);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 4cab1e24bf9e6..156c10b4118a1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -148,6 +148,9 @@ static void propagateResultsToRegularOperands(Operation *op) {
if (isa<VectorType>(resultType) || isa<vector::MultiDimReductionOp>(op))
xegpu::setTemporaryLayout(result, resLayout);
+ if (isa<vector::DeinterleaveOp>(op))
+ xegpu::setTemporaryLayout(op->getResult(1), resLayout);
+
for (OpOperand &opr : op->getOpOperands()) {
xegpu::DistributeLayoutAttr operandLayout =
xegpu::inferSourceLayoutFromResult(opr, resLayout);
@@ -1003,6 +1006,7 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
int resElemTyBitWidth = resVecTy.getElementType().getIntOrFloatBitWidth();
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
@@ -1013,7 +1017,6 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
int64_t instDataValue = -1;
int64_t laneDataValue = -1;
const int subgroupSize = uArch->getSubgroupSize();
-
if (srcElemTyBitWidth > resElemTyBitWidth) {
// When casting to a smaller bitwidth, multiply the result layout
// accordingly to ensure it can be divided by the ratio back to the
@@ -1026,14 +1029,14 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
instDataValue = instData[dim];
// Adjust instDataValue so it still fits within an instruction after
// dividing by bitWidthRatio
- while ((instDataValue <= srcShape[dim]) &&
+ while ((instDataValue <= resShape[dim]) &&
(instDataValue % (innermostDimLaneLayout * bitWidthRatio) != 0))
instDataValue *= 2;
- assert((srcShape[dim] % instDataValue) == 0 &&
- "srcShape, instData, and lanelayout for innermost must be 2^n !");
+ assert((resShape[dim] % instDataValue) == 0 &&
+ "resShape, instData, and lanelayout for innermost must be 2^n !");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
laneDataValue = laneData[dim];
- while ((laneDataValue <= srcShape[dim]) &&
+ while ((laneDataValue <= resShape[dim]) &&
(laneDataValue % bitWidthRatio != 0))
laneDataValue *= 2;
}
@@ -1066,6 +1069,7 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
DistributeLayoutAttr consumerLayout, const xegpu::uArch::uArch *uArch) {
ArrayRef<int64_t> srcShape = srcVecTy.getShape();
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
@@ -1084,23 +1088,23 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
if (layoutKind == xegpu::LayoutKind::Subgroup) {
sgDataValue = sgData[innerMostDim];
// Ensure sgDataValue is divisible by ratio so source sgData can be inferred
- while ((sgDataValue <= srcShape[innerMostDim]) &&
+ while ((sgDataValue <= resShape[innerMostDim]) &&
(sgDataValue % ratio != 0))
sgDataValue *= ratio;
} else if (layoutKind == xegpu::LayoutKind::InstData) {
instDataValue = instData[innerMostDim];
// Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
// ratio) when inferring the source layout
- while ((instDataValue <= srcShape[innerMostDim]) &&
+ while ((instDataValue <= resShape[innerMostDim]) &&
(instDataValue % (innermostDimLaneLayout * ratio) != 0))
instDataValue *= ratio;
- assert((srcShape[innerMostDim] % instDataValue) == 0 &&
- "srcShape, instData, and laneLayout for innermost must be 2^n!");
+ assert((resShape[innerMostDim] % instDataValue) == 0 &&
+ "resShape, instData, and laneLayout for innermost must be 2^n!");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
laneDataValue = laneData[innerMostDim];
// Ensure laneDataValue is at least 2 and divisible by ratio
// so that source laneData = laneDataValue/2 is valid
- while ((laneDataValue <= srcShape[innerMostDim]) &&
+ while ((laneDataValue <= resShape[innerMostDim]) &&
(laneDataValue % ratio != 0))
laneDataValue *= ratio;
}
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
new file mode 100644
index 0000000000000..d4f9e08019ead
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
@@ -0,0 +1,73 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=lane zebin-chip=cri" \
+// RUN: | mlir-runner \
+// RUN: --shared-libs=%mlir_levelzero_runtime \
+// RUN: --shared-libs=%mlir_runner_utils \
+// RUN: --shared-libs=%mlir_c_runner_utils \
+// RUN: --entry-point-result=void \
+// RUN: | FileCheck %s
+
+// XFAIL: *
+#a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#b_packed = #xegpu.layout<sg_layout = [8, 8], sg_data = [256, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
+#b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
+#c = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>
+#b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+
+gpu.module @test {
+ gpu.func @gemm_mxfp(%arg0: memref<1024x4096xf4E2M1FN>, %arg1: memref<2048x1024xui8>, %arg2: memref<1024x128xf8E8M0FNU>, %arg3: memref<128x1024xf8E8M0FNU>, %arg4: memref<1024x1024xf32>) {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c128 = arith.constant 128 : index
+ %c1024 = arith.constant 1024 : index
+ %block_id_x = gpu.block_id x
+ %block_id_y = gpu.block_id y
+ %0 = arith.muli %block_id_x, %c128 : index
+ %1 = arith.muli %block_id_y, %c128 : index
+
+ %a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<1024x4096xf4E2M1FN> -> !xegpu.tensor_desc<128x512xf4E2M1FN>
+ %bp_tdesc = xegpu.create_nd_tdesc %arg1 : memref<2048x1024xui8> -> !xegpu.tensor_desc<256x128xui8>
+ // load_nd with offset
+ %a = xegpu.load_nd %a_tdesc[%0, %c0] {layout = #a}: !xegpu.tensor_desc<128x512xf4E2M1FN> -> vector<128x512xf4E2M1FN>
+ %bp = xegpu.load_nd %bp_tdesc[%c0, %1] {layout = #b_packed}: !xegpu.tensor_desc<256x128xui8> -> vector<256x128xui8>
+
+ // Bitcast to fp4: 256x128 uint8 -> 256x256 fp4 (each uint8 holds 2 fp4 values)
+ %b_bitcast = vector.bitcast %bp : vector<256x128xui8> to vector<256x256xf4E2M1FN>
+
+ // De-interleave: extract even and odd columns
+ // Even columns (indices 0, 2, 4, ..., 254) -> first half
+ // Odd columns (indices 1, 3, 5, ..., 255) -> second half
+ %b_even, %b_odd = vector.deinterleave %b_bitcast : vector<256x256xf4E2M1FN> -> vector<256x128xf4E2M1FN>
+
+ // Reconstruct 512x128 by interleaving even/odd rows:
+ // Transpose to move the row dim to trailing position, interleave, transpose back.
+ %b_even_t = vector.transpose %b_even, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_odd_t = vector.transpose %b_odd, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_interleaved = vector.interleave %b_even_t, %b_odd_t : vector<128x256xf4E2M1FN> -> vector<128x512xf4E2M1FN>
+ %b = vector.transpose %b_interleaved, [1, 0] : vector<128x512xf4E2M1FN> to vector<512x128xf4E2M1FN>
+
+ %cd_tdesc = xegpu.create_nd_tdesc %arg4 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32, #c>
+ %c = xegpu.load_nd %cd_tdesc[%0, %1] {layout = #c}: !xegpu.tensor_desc<128x128xf32, #c> -> vector<128x128xf32>
+
+ %a_scale_tdesc = xegpu.create_nd_tdesc %arg2 : memref<1024x128xf8E8M0FNU> -> !xegpu.tensor_desc<128x16xf8E8M0FNU>
+ %scale_a = xegpu.load_nd %a_scale_tdesc[%0, %c0] {layout = #a_scale}: !xegpu.tensor_desc<128x16xf8E8M0FNU> -> vector<128x16xf8E8M0FNU>
+
+ %b_scale_tdesc = xegpu.create_nd_tdesc %arg3 : memref<128x1024xf8E8M0FNU> -> !xegpu.tensor_desc<16x128xf8E8M0FNU>
+ %scale_b = xegpu.load_nd %b_scale_tdesc[%c0, %1] {layout = #b_scale}: !xegpu.tensor_desc<16x128xf8E8M0FNU> -> vector<16x128xf8E8M0FNU>
+
+ %d = xegpu.dpas_mx %a, %b, %c scale_a = %scale_a scale_b = %scale_b
+ {layout_a = #a,
+ layout_b = #b,
+ layout_cd = #c,
+ layout_a_scale = #a_scale,
+ layout_b_scale = #b_scale}
+ : vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>,
+ vector<128x128xf32>,
+ vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>
+ -> vector<128x128xf32>
+
+ // store_nd with offset
+ xegpu.store_nd %d, %cd_tdesc[%0, %1] {layout = #c} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #c>
+ gpu.return
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/196243
More information about the Mlir-commits
mailing list