[Mlir-commits] [mlir] [MLIR][XeGPU] Enhance unrolling of convert layout (PR #209822)
Nishant Patel
llvmlistbot at llvm.org
Sun Jul 19 21:47:06 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/209822
>From e9401f094f7163276fd348056d310a895236c428 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 14 Jul 2026 21:43:37 +0000
Subject: [PATCH 1/2] Enhance unrolling of convert layout
---
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 124 ++++++++++++++++--
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 32 ++++-
2 files changed, 137 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 74c358cef90df..bcb9946373298 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -957,8 +957,92 @@ struct UnrollStoreMatrixOp : public UnrollPattern<xegpu::StoreMatrixOp> {
/// after inst_data stripped. If it does, it will unroll the vector into
/// multiple smaller vectors according to the target shape, and create multiple
/// ConvertLayoutOp with the unrolled vectors and the stripped layouts.
+///
+/// When the input and target layouts have different inst_data, the source is
+/// extracted at the input inst_data granularity and the result is inserted at
+/// the target inst_data granularity, enabling slice cancellation during
+/// canonicalization.
struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
using UnrollPattern<xegpu::ConvertLayoutOp>::UnrollPattern;
+
+ /// Extracts source in `inTile` slices, regroups into `convTile`-sized
+ /// ConvertLayoutOps, and inserts the result in `outTile` slices.
+ /// Returns failure if the tiles do not evenly divide.
+ LogicalResult
+ rewriteWithRegrouping(xegpu::ConvertLayoutOp op, VectorType valueTy,
+ ArrayRef<int64_t> convTile, ArrayRef<int64_t> inTile,
+ ArrayRef<int64_t> outTile,
+ xegpu::DistributeLayoutAttr inputLayout,
+ xegpu::DistributeLayoutAttr targetLayout, Location loc,
+ PatternRewriter &rewriter) const {
+ ArrayRef<int64_t> vecShape = valueTy.getShape();
+ if (!computeShapeRatio(vecShape, convTile) ||
+ !computeShapeRatio(convTile, inTile) ||
+ !computeShapeRatio(convTile, outTile))
+ return failure();
+
+ Type elemTy = valueTy.getElementType();
+ int64_t rank = valueTy.getRank();
+ VectorType convTy = VectorType::get(convTile, elemTy);
+ SmallVector<int64_t> strides(rank, 1);
+
+ Value source = op.getSource();
+ auto zeroOf = [&](VectorType ty) -> Value {
+ return arith::ConstantOp::create(
+ rewriter, loc, ty,
+ DenseElementsAttr::get(ty, rewriter.getZeroAttr(elemTy)));
+ };
+ auto addOffsets = [](ArrayRef<int64_t> a,
+ ArrayRef<int64_t> b) -> SmallVector<int64_t> {
+ SmallVector<int64_t> res(a);
+ for (auto [r, v] : llvm::zip_equal(res, b))
+ r += v;
+ return res;
+ };
+
+ Value result = zeroOf(valueTy);
+ for (SmallVector<int64_t> convOff :
+ StaticTileOffsetRange(vecShape, convTile)) {
+ // Build the convert tile from inTile-sized slices of the source.
+ Value conv;
+ if (convTile == inTile) {
+ conv = vector::ExtractStridedSliceOp::create(
+ rewriter, loc, source, convOff, convTile, strides);
+ } else {
+ conv = zeroOf(convTy);
+ for (SmallVector<int64_t> inLocal :
+ StaticTileOffsetRange(convTile, inTile)) {
+ Value piece = vector::ExtractStridedSliceOp::create(
+ rewriter, loc, source, addOffsets(convOff, inLocal), inTile,
+ strides);
+ conv = vector::InsertStridedSliceOp::create(rewriter, loc, piece,
+ conv, inLocal, strides);
+ }
+ }
+
+ conv = xegpu::ConvertLayoutOp::create(rewriter, loc, convTy, conv,
+ inputLayout, targetLayout);
+
+ // Write the converted tile into the result as outTile-sized slices.
+ if (convTile == outTile) {
+ result = vector::InsertStridedSliceOp::create(rewriter, loc, conv,
+ result, convOff, strides);
+ } else {
+ for (SmallVector<int64_t> outLocal :
+ StaticTileOffsetRange(convTile, outTile)) {
+ Value piece = vector::ExtractStridedSliceOp::create(
+ rewriter, loc, conv, outLocal, outTile, strides);
+ result = vector::InsertStridedSliceOp::create(
+ rewriter, loc, piece, result, addOffsets(convOff, outLocal),
+ strides);
+ }
+ }
+ }
+
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+
LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
@@ -974,8 +1058,10 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
return success();
}
- if (inputLayout.getEffectiveInstDataAsInt().empty() ||
- targetLayout.getEffectiveInstDataAsInt().empty())
+ // Capture inst_data granularities before stripping them.
+ SmallVector<int64_t> inTile = inputLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> outTile = targetLayout.getEffectiveInstDataAsInt();
+ if (inTile.empty() || outTile.empty())
return rewriter.notifyMatchFailure(op, "Not a target ConvertLayoutOp.");
inputLayout = inputLayout.dropInstData();
@@ -988,20 +1074,30 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
if (!targetShape || targetShape->size() != (size_t)valueTy.getRank())
return failure();
- Value newSource = op.getSource();
+ // Nothing to convert if layouts match after stripping inst_data.
+ if (!inputLayout || !targetLayout || inputLayout.isEqualTo(targetLayout)) {
+ rewriter.replaceOp(op, op.getSource());
+ return success();
+ }
+
+ // Try regrouping: extract at inTile, convert, insert at outTile.
+ if (succeeded(rewriteWithRegrouping(op, valueTy, *targetShape, inTile,
+ outTile, inputLayout, targetLayout, loc,
+ rewriter)))
+ return success();
+
+ // Fallback: pack/unpack at the convert tile granularity.
+ SmallVector<Type> convertedValTypes =
+ getUnrolledTypes(valueTy, *targetShape);
+ SmallVector<Value> convertedValues =
+ pack(op.getOperand(), convertedValTypes, *targetShape, loc, rewriter);
SmallVector<Value> newOps;
- if (inputLayout && targetLayout && !inputLayout.isEqualTo(targetLayout)) {
- SmallVector<Type> convertedValTypes =
- getUnrolledTypes(valueTy, *targetShape);
- SmallVector<Value> convertedValues =
- pack(op.getOperand(), convertedValTypes, *targetShape, loc, rewriter);
- for (auto [v, t] : llvm::zip(convertedValues, convertedValTypes)) {
- auto newOp = xegpu::ConvertLayoutOp::create(rewriter, loc, t, v,
- inputLayout, targetLayout);
- newOps.push_back(newOp);
- }
- newSource = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
+ for (auto [v, t] : llvm::zip(convertedValues, convertedValTypes)) {
+ auto newOp = xegpu::ConvertLayoutOp::create(rewriter, loc, t, v,
+ inputLayout, targetLayout);
+ newOps.push_back(newOp);
}
+ Value newSource = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, newSource);
return success();
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index b0b16c2adba6e..27a88efcf443c 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -463,19 +463,19 @@ gpu.module @test_kernel {
//CHECK: gpu.func @convert_layout([[arg0:%.+]]: vector<8x32x2xf16>) -> vector<8x32x2xf16> {
//CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<8x32x2xf16>
//CHECK: [[e0:%.+]] = vector.extract_strided_slice [[arg0]] {offsets = [0, 0, 0], sizes = [4, 32, 2], strides = [1, 1, 1]} : vector<8x32x2xf16> to vector<4x32x2xf16>
- //CHECK: [[e1:%.+]] = vector.extract_strided_slice [[arg0]] {offsets = [4, 0, 0], sizes = [4, 32, 2], strides = [1, 1, 1]} : vector<8x32x2xf16> to vector<4x32x2xf16>
//CHECK: [[c0:%.+]] = xegpu.convert_layout [[e0]] <{input_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 2]>, target_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>}> : vector<4x32x2xf16>
- //CHECK: [[c1:%.+]] = xegpu.convert_layout [[e1]] <{input_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 2]>, target_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>}> : vector<4x32x2xf16>
//CHECK: [[e2:%.+]] = vector.extract_strided_slice [[c0]] {offsets = [0, 0, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
+ //CHECK: [[e3:%.+]] = vector.extract_strided_slice [[c0]] {offsets = [0, 16, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
+ //CHECK: [[e1:%.+]] = vector.extract_strided_slice [[arg0]] {offsets = [4, 0, 0], sizes = [4, 32, 2], strides = [1, 1, 1]} : vector<8x32x2xf16> to vector<4x32x2xf16>
+ //CHECK: [[c1:%.+]] = xegpu.convert_layout [[e1]] <{input_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 2]>, target_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>}> : vector<4x32x2xf16>
+ //CHECK: [[e4:%.+]] = vector.extract_strided_slice [[c1]] {offsets = [0, 0, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
+ //CHECK: [[e5:%.+]] = vector.extract_strided_slice [[c1]] {offsets = [0, 16, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
//CHECK: [[m0:%.+]] = math.exp [[e2]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
//CHECK: [[i0:%.+]] = vector.insert_strided_slice [[m0]], [[cst]] {offsets = [0, 0, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
- //CHECK: [[e3:%.+]] = vector.extract_strided_slice [[c0]] {offsets = [0, 16, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
//CHECK: [[m1:%.+]] = math.exp [[e3]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
//CHECK: [[i1:%.+]] = vector.insert_strided_slice [[m1]], [[i0]] {offsets = [0, 16, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
- //CHECK: [[e4:%.+]] = vector.extract_strided_slice [[c1]] {offsets = [0, 0, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
//CHECK: [[m2:%.+]] = math.exp [[e4]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
//CHECK: [[i2:%.+]] = vector.insert_strided_slice [[m2]], [[i1]] {offsets = [4, 0, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
- //CHECK: [[e5:%.+]] = vector.extract_strided_slice [[c1]] {offsets = [0, 16, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
//CHECK: [[m3:%.+]] = math.exp [[e5]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
//CHECK: [[i3:%.+]] = vector.insert_strided_slice [[m3]], [[i2]] {offsets = [4, 16, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
//CHECK: gpu.return [[i3]] : vector<8x32x2xf16>
@@ -488,6 +488,28 @@ gpu.module @test_kernel {
}
}
+// -----
+
+// Test regrouping when input and target inst_data differ ([8, 16] vs [8, 64]).
+#in = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#tgt = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
+
+gpu.module @test_kernel {
+ //CHECK-LABEL: gpu.func @convert_layout_regroup
+ //CHECK-SAME: ([[arg0:%.+]]: vector<16x64xbf16>)
+ //CHECK: vector.extract_strided_slice [[arg0]] {offsets = [0, 0], sizes = [8, 16]
+ //CHECK: math.exp {{.*}} : vector<8x16xbf16>
+ //CHECK-NOT: vector.extract_strided_slice [[arg0]] {{.*}}sizes = [8, 64]
+ //CHECK: vector.insert_strided_slice {{.*}} : vector<8x16xbf16> into vector<8x64xbf16>
+ //CHECK: xegpu.convert_layout {{.*}} : vector<8x64xbf16>
+ //CHECK: vector.insert_strided_slice {{.*}} : vector<8x64xbf16> into vector<16x64xbf16>
+ gpu.func @convert_layout_regroup(%a: vector<16x64xbf16>) -> vector<16x64xbf16> {
+ %p = math.exp %a {layout_result_0 = #in} : vector<16x64xbf16>
+ %0 = xegpu.convert_layout %p <{input_layout = #in, target_layout = #tgt}> : vector<16x64xbf16>
+ gpu.return %0 : vector<16x64xbf16>
+ }
+}
+
// -----
gpu.module @test_kernel {
//CHECK-LABEL: unroll_load_matrix
>From 41c298a78f70f3c2727bae352876565063aa5fde Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 20 Jul 2026 04:46:26 +0000
Subject: [PATCH 2/2] Add test
---
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 21 +++++++++++++++++++++
1 file changed, 21 insertions(+)
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index 27a88efcf443c..bc50bad2b2302 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -510,6 +510,27 @@ gpu.module @test_kernel {
}
}
+// -----
+
+// Test regrouping when input and target inst_data differ ([8, 64] vs [8, 16]).
+#in = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
+#tgt = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+
+gpu.module @test_kernel {
+ //CHECK-LABEL: gpu.func @convert_layout_regroup_swapped
+ //CHECK-SAME: ([[arg0:%.+]]: vector<16x64xbf16>)
+ //CHECK: vector.extract_strided_slice [[arg0]] {offsets = [0, 0], sizes = [8, 64], strides = [1, 1]} : vector<16x64xbf16> to vector<8x64xbf16>
+ //CHECK-NOT: vector.extract_strided_slice [[arg0]] {{.*}}sizes = [8, 16]
+ //CHECK: [[cvt:%.+]] = xegpu.convert_layout {{.*}} : vector<8x64xbf16>
+ //CHECK: vector.extract_strided_slice [[cvt]] {{.*}}sizes = [8, 16]{{.*}} : vector<8x64xbf16> to vector<8x16xbf16>
+ //CHECK: vector.insert_strided_slice {{.*}} : vector<8x16xbf16> into vector<16x64xbf16>
+ gpu.func @convert_layout_regroup_swapped(%a: vector<16x64xbf16>) -> vector<16x64xbf16> {
+ %p = math.exp %a {layout_result_0 = #in} : vector<16x64xbf16>
+ %0 = xegpu.convert_layout %p <{input_layout = #in, target_layout = #tgt}> : vector<16x64xbf16>
+ gpu.return %0 : vector<16x64xbf16>
+ }
+}
+
// -----
gpu.module @test_kernel {
//CHECK-LABEL: unroll_load_matrix
More information about the Mlir-commits
mailing list