[Mlir-commits] [mlir] [mlir][XeGPU] Distribute create_nd_tdesc/load/store with SliceAttr layout (PR #216104)
Andrey Pavlenko
llvmlistbot at llvm.org
Thu Aug 13 10:54:36 PDT 2026
https://github.com/AndreyPavlenko updated https://github.com/llvm/llvm-project/pull/216104
>From 4b712bcc54903e7955d102876db28f9baee746fc Mon Sep 17 00:00:00 2001
From: Andrey Pavlenko <andrey.a.pavlenko at gmail.com>
Date: Thu, 13 Aug 2026 16:19:01 +0000
Subject: [PATCH 1/2] [mlir][XeGPU] Distribute create_nd_tdesc/load/store with
SliceAttr layout
WgToSgCreateNdOp and the pass's legality check only matched
xegpu::LayoutAttr on a tensor_desc, missing xegpu::SliceAttr (also a
DistributeLayoutAttr). A tensor_desc feeding a unit-dim-expanding
vector.shape_cast carries a SliceAttr, so it was left undistributed
while its consumers were converted to subgroup shape, producing a
shape mismatch on the shape_cast. Match on DistributeLayoutAttr
instead, consistent with the rest of the pass.
---
.../Transforms/XeGPUWgToSgDistribute.cpp | 6 ++--
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 35 +++++++++++++++++++
2 files changed, 38 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index afea358dafe29..a24e9b2fd7e0f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -189,7 +189,7 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
Location loc = op.getLoc();
MLIRContext *ctx = op.getContext();
xegpu::TensorDescType tdescTy = op.getType();
- auto layout = dyn_cast<xegpu::LayoutAttr>(tdescTy.getLayout());
+ auto layout = dyn_cast<xegpu::DistributeLayoutAttr>(tdescTy.getLayout());
if (!layout || !layout.isForWorkgroup())
return failure();
@@ -1599,8 +1599,8 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
xegpu::StoreNdOp, xegpu::PrefetchNdOp>(
[=](Operation *op) -> bool {
auto tdescTy = getTensorDescType(op);
- auto layout =
- dyn_cast_if_present<xegpu::LayoutAttr>(tdescTy.getLayout());
+ auto layout = dyn_cast_if_present<xegpu::DistributeLayoutAttr>(
+ tdescTy.getLayout());
return isLegal(layout);
});
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index cdc9392f02e04..3944c6c9f23c5 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -1370,3 +1370,38 @@ func.func @no_crash_on_dynamic_tensor(%arg0: tensor<?xi32>, %arg1: index) -> ten
}
return %result : tensor<?xi32>
}
+
+// -----
+
+// Regression test: a tensor_desc whose layout is a SliceAttr (produced e.g.
+// by unit-dim-expanding a vector.shape_cast) was not recognized as a
+// workgroup-level layout by WgToSgCreateNdOp, since it only matched
+// LayoutAttr. This left the tensor_desc and its load_nd undistributed,
+// causing a shape mismatch on the consuming shape_cast.
+gpu.module @test_slice_layout {
+ // CHECK-LABEL: slice_layout_feeding_shape_cast
+ gpu.func @slice_layout_feeding_shape_cast(%arg0: memref<1024x1536xf16>, %arg1: memref<1024x1536xf16>) {
+ %cst = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>} dense<3.000000e+00> : vector<1x32x64xf16>
+ // CHECK: xegpu.create_nd_tdesc %{{.*}} : memref<1024x1536xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %tdesc0 = xegpu.create_nd_tdesc %arg0 : memref<1024x1536xf16>
+ -> !xegpu.tensor_desc<32x64xf16, #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>>
+ // CHECK: xegpu.load_nd %{{.*}} : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %load = xegpu.load_nd %tdesc0[0, 0] <{layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>}>
+ : !xegpu.tensor_desc<32x64xf16, #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>> -> vector<32x64xf16>
+ // CHECK: vector.shape_cast {{.*}} : vector<8x16xf16> to vector<1x8x16xf16>
+ %expand = vector.shape_cast %load {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>}
+ : vector<32x64xf16> to vector<1x32x64xf16>
+ // CHECK: arith.addf {{.*}} : vector<1x8x16xf16>
+ %add = arith.addf %expand, %cst {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>}
+ : vector<1x32x64xf16>
+ // CHECK: vector.shape_cast {{.*}} : vector<1x8x16xf16> to vector<8x16xf16>
+ %collapse = vector.shape_cast %add {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>}
+ : vector<1x32x64xf16> to vector<32x64xf16>
+ // CHECK: xegpu.store_nd {{.*}} : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
+ %tdesc1 = xegpu.create_nd_tdesc %arg1 : memref<1024x1536xf16>
+ -> !xegpu.tensor_desc<32x64xf16, #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>>
+ xegpu.store_nd %collapse, %tdesc1[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>}>
+ : vector<32x64xf16>, !xegpu.tensor_desc<32x64xf16, #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>>
+ gpu.return
+ }
+}
>From 8137347902568652e37dc2444810a2644eccc500 Mon Sep 17 00:00:00 2001
From: Andrey Pavlenko <andrey.a.pavlenko at gmail.com>
Date: Thu, 13 Aug 2026 17:37:55 +0000
Subject: [PATCH 2/2] [mlir][XeGPU] Trim regression test to just
create_nd_tdesc/load_nd
Address review: drop the redundant shape_cast/arith/store_nd chain and
the second (plain-LayoutAttr) tensor_desc, and shorten the comment.
The SliceAttr-on-tensor_desc distribution is exercised by
create_nd_tdesc/load_nd alone.
---
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 29 ++++-----------------
1 file changed, 5 insertions(+), 24 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 3944c6c9f23c5..2140bd65c95dc 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -1373,35 +1373,16 @@ func.func @no_crash_on_dynamic_tensor(%arg0: tensor<?xi32>, %arg1: index) -> ten
// -----
-// Regression test: a tensor_desc whose layout is a SliceAttr (produced e.g.
-// by unit-dim-expanding a vector.shape_cast) was not recognized as a
-// workgroup-level layout by WgToSgCreateNdOp, since it only matched
-// LayoutAttr. This left the tensor_desc and its load_nd undistributed,
-// causing a shape mismatch on the consuming shape_cast.
+// Regression test: a tensor_desc with a SliceAttr layout was not distributed.
gpu.module @test_slice_layout {
- // CHECK-LABEL: slice_layout_feeding_shape_cast
- gpu.func @slice_layout_feeding_shape_cast(%arg0: memref<1024x1536xf16>, %arg1: memref<1024x1536xf16>) {
- %cst = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>} dense<3.000000e+00> : vector<1x32x64xf16>
+ // CHECK-LABEL: slice_layout
+ gpu.func @slice_layout(%arg0: memref<1024x1536xf16>) {
// CHECK: xegpu.create_nd_tdesc %{{.*}} : memref<1024x1536xf16> -> !xegpu.tensor_desc<8x16xf16>
- %tdesc0 = xegpu.create_nd_tdesc %arg0 : memref<1024x1536xf16>
+ %tdesc = xegpu.create_nd_tdesc %arg0 : memref<1024x1536xf16>
-> !xegpu.tensor_desc<32x64xf16, #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>>
// CHECK: xegpu.load_nd %{{.*}} : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %load = xegpu.load_nd %tdesc0[0, 0] <{layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>}>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>}>
: !xegpu.tensor_desc<32x64xf16, #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>, dims = [0]>> -> vector<32x64xf16>
- // CHECK: vector.shape_cast {{.*}} : vector<8x16xf16> to vector<1x8x16xf16>
- %expand = vector.shape_cast %load {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>}
- : vector<32x64xf16> to vector<1x32x64xf16>
- // CHECK: arith.addf {{.*}} : vector<1x8x16xf16>
- %add = arith.addf %expand, %cst {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 4], sg_data = [1, 8, 16]>}
- : vector<1x32x64xf16>
- // CHECK: vector.shape_cast {{.*}} : vector<1x8x16xf16> to vector<8x16xf16>
- %collapse = vector.shape_cast %add {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>}
- : vector<1x32x64xf16> to vector<32x64xf16>
- // CHECK: xegpu.store_nd {{.*}} : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
- %tdesc1 = xegpu.create_nd_tdesc %arg1 : memref<1024x1536xf16>
- -> !xegpu.tensor_desc<32x64xf16, #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>>
- xegpu.store_nd %collapse, %tdesc1[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>}>
- : vector<32x64xf16>, !xegpu.tensor_desc<32x64xf16, #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 16]>>
gpu.return
}
}
More information about the Mlir-commits
mailing list