[Mlir-commits] [mlir] [mlir][xegpu] Change inst-data of load/load_matrix to fit lane layout (PR #209661)
Jianhui Li
llvmlistbot at llvm.org
Tue Jul 14 19:11:05 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/209661
The consumer layout may have multiple distribution of lane data, but load_gather/load_matrix only allow one single distribution, change their inst_data according to the single distribution of lane data.
>From 8d4af3f7516a86beb8634bff744bbc020bee84be Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 15 Jul 2026 02:02:16 +0000
Subject: [PATCH] change the inst-data layout of load to fit lane layout
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 15 ++----
.../XeGPU/propagate-layout-inst-data.mlir | 47 +++++++++++++++----
2 files changed, 41 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 8da151a333fb3..d32d53fa39613 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1686,9 +1686,6 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
- // Pick lane_layout / lane_data: prefer consumer's, fall back to the
- // scatter-store default (subgroupSize lanes on innermost dim, per-lane
- // vector capped by maxChunkSize).
SmallVector<int64_t> laneLayout;
SmallVector<int64_t> laneData;
assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
@@ -1697,16 +1694,10 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
if (layoutKind == xegpu::LayoutKind::InstData) {
- // Take consumer's inst_data as-is. If the consumer doesn't have one,
- // fall back to lane_layout * lane_data per dim.
SmallVector<int64_t> instData;
- if (!consumerInstData.empty()) {
- instData.assign(consumerInstData.begin(), consumerInstData.end());
- } else {
- instData.resize(resShape.size());
- for (size_t i = 0; i < resShape.size(); ++i)
- instData[i] = laneLayout[i] * laneData[i];
- }
+ instData.resize(resShape.size());
+ for (size_t i = 0; i < resShape.size(); ++i)
+ instData[i] = laneLayout[i] * laneData[i];
return buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
}
if (layoutKind == xegpu::LayoutKind::Lane)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index c72e917b6fafa..ea2a7f395b764 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -128,37 +128,65 @@ func.func @store_matrix(%arg0: !xegpu.mem_desc<16x64xf16>) {
}
}
+// TODO: colaesce_chunksize tests should be adjusted when we enable colasce optimization
// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @load_matrix_with_coalesce_chunksize(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<2048x512xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<2048x512xf32>)
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<128xf32, 3>
+// CHECK: %[[MDESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<128xf32, 3> -> !xegpu.mem_desc<128xf32>
+// CHECK: %[[LOAD:.*]] = xegpu.load_matrix %[[MDESC]][%[[C0]]] <{layout = #xegpu.layout<inst_data = [1], lane_layout = [1], lane_data = [1]>}>: !xegpu.mem_desc<128xf32>, index -> vector<16xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>} : vector<16xf32> to vector<16x16xf32>
+// CHECK: %[[TRANS:.*]] = vector.transpose %[[BCAST]], [1, 0] {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf32> to vector<16x16xf32>
+// CHECK: %[[EXP:.*]] = math.exp %[[TRANS]] {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf32>
+// CHECK: %[[DIV:.*]] = arith.divf %[[EXP]], %[[TRANS]] {layout_result_0 = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf32>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<2048x512xf32> -> !xegpu.tensor_desc<16x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+// CHECK: xegpu.store_nd %[[DIV]], %[[TDESC]][%[[C0]], %[[C0]]] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+func.func @load_matrix_with_coalesce_chunksize(%arg0: memref<2048x512xf32>, %arg1: memref<2048x512xf32>) {
+ %c0 = arith.constant 0 : index
+ %alloca = memref.alloca() : memref<128xf32, 3>
+ %6 = xegpu.create_mem_desc %alloca : memref<128xf32, 3> -> !xegpu.mem_desc<128xf32>
+ %13 = xegpu.load_matrix %6[%c0] : !xegpu.mem_desc<128xf32>, index -> vector<16xf32>
+ %14 = vector.broadcast %13 : vector<16xf32> to vector<16x16xf32>
+ %15 = vector.transpose %14, [1, 0] : vector<16x16xf32> to vector<16x16xf32>
+ %17 = math.exp %15 : vector<16x16xf32>
+ %18 = arith.divf %17, %15 : vector<16x16xf32>
+ %22 = xegpu.create_nd_tdesc %arg1 : memref<2048x512xf32> -> !xegpu.tensor_desc<16x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+ xegpu.store_nd %18, %22[%c0, %c0] : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+ return
+}
+}
+// -----
gpu.module @test {
// CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<512xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<true> : vector<16x32xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<12> : vector<16x32xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} dense<true> : vector<16x32xi1>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} dense<12> : vector<16x32xindex>
+// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
// CHECK-SAME: memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
-// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
+// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
// CHECK-SAME: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
func.func @scatter_ops_coalesce_chunksize(%src: memref<512xf32>) {
%1 = arith.constant dense<1>: vector<16x32xi1>
%offset = arith.constant dense<12> : vector<16x32xindex>
%3 = xegpu.load %src[%offset], %1
: memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
- xegpu.store %3, %src[%offset], %1 <{layout = #xegpu.layout<inst_data = [16, 16]>}>
+ xegpu.store %3, %src[%offset], %1 <{layout = #xegpu.layout<inst_data = [1, 32]>}>
: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
return
}
}
// -----
-
gpu.module @test {
// CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>}
// CHECK-SAME: dense<0> : vector<16x16xindex>
-// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
+// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
@@ -176,6 +204,7 @@ func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: m
}
// -----
+
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
More information about the Mlir-commits
mailing list