[Mlir-commits] [mlir] [mlir][xegpu] Derive load_nd lane_data from the hardware block variant (PR #223164)
Jianhui Li
llvmlistbot at llvm.org
Thu Sep 17 19:46:11 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/223164
>From aa09fce6ba1744e9fb1a01b24cf406d732a47bfe Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 12 Sep 2026 17:36:54 +0000
Subject: [PATCH 1/4] [mlir][xegpu] Derive load_nd lane_data from the hardware
block variant
setupLoadNdAnchorLayout took lane_data straight from the consumer layout. That
is only correct by accident: lane_data counts the elements a lane holds in one
32-bit register, so it follows from the element width and the block variant the
hardware will use, not from what the consumer would like. A consumer asking for
more elements per lane than a block load can pack produced an unrealizable
layout.
Add get2DBlockLoadLaneData, which derives lane_data from the packing size the
uArch already reports:
* regular: the load's packed format size (16 bits) along the innermost dim
* transform: the general packed format size (32 bits) along the
second-innermost dim, i.e. VNNI
* transpose: the general packed format size along the innermost dim
An element at least as wide as its packing size gets unit lane_data, so f32 is
always [1, 1].
On the fallback path, where the consumer's inst_data is not a usable block,
also pick the variant that can host that lane_data. DPAS_MX scales land here:
their inst_data is smaller than any block load, and a scale whose lanes pack
along the innermost dim cannot be loaded regularly, because that dim is
narrower than lane_layout * lane_data needs (width 16 for an 8-bit scale, which
spreads 16 lanes of 2 elements over 32 columns). Loading it transformed packs
down the column instead, which fits the narrow width. A scale asking for a
vertical lane_layout is already treated as a transpose and needs no adjustment.
Scale layouts now come out as:
scale_a inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4]
scale_b inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]
Also correct the uArch block table: an 8-bit regular block load packs two
elements per lane, so spreading a subgroup of 16 lanes over it needs 32
columns. Its supported width is 32, not {16, 32}. This is what makes both
scales fall through to the block selection above instead of being handed a
block shape they cannot distribute.
The ui8 bitcast test moves with the rule: its 8-bit load now packs two elements
per lane, so inst_data widens to 32 and the f4 consumer layout it is derived
from becomes inst_data = [32, 64], lane_data = [1, 4].
Full MLIR check: 4613 tests, only the 5 pre-existing failures.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/uArch/uArchBase.h | 3 +-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 76 ++++++++++++++++---
.../XeGPU/propagate-layout-inst-data.mlir | 36 ++++-----
3 files changed, 85 insertions(+), 30 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index e1879bb3ffe1c..d4ba2391729e1 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -285,7 +285,6 @@ struct Subgroup2DBlockLoadInstruction : public Instruction,
static const int kWidth64[] = {64};
static const int kWidth32[] = {32};
static const int kWidth16[] = {16};
- static const int kWidthAtLeast16[] = {16, 32};
static const int kWidthAtLeast32[] = {32, 64};
static const int kWidth8[] = {8};
@@ -302,7 +301,7 @@ struct Subgroup2DBlockLoadInstruction : public Instruction,
// expressed directly. 4-bit elements are packed two-per-byte, so their
// widths (or heights, when transformed) are double the 8-bit rows.
static const llvm::DenseMap<Key, Value> kMap = {
- {{8, false, false, false}, {kWidthAtLeast16, kHeightAtLeast1, kCount2}},
+ {{8, false, false, false}, {kWidth32, kHeightAtLeast1, kCount2}},
{{8, false, false, true}, {kWidth16, kHeightAtLeast8, kCount4Only}},
{{16, false, false, false}, {kWidth16, kHeightAtLeast1, kCount2}},
{{32, false, false, false}, {kWidth16, kHeightAtLeast1, kCount1}},
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index f696020b2de12..b33470db6b902 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1759,11 +1759,39 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
}
+/// Returns the lane_data a 2D block load can deliver for `elemTy`. lane_data
+/// counts the elements a lane holds in one register, so it follows the packing
+/// granularity of the block variant:
+///
+/// * regular: packs along the innermost dim, to the load's packed format
+/// size.
+/// * transform: packs along the second-innermost dim, to the uArch's general
+/// packed format size. This is VNNI.
+/// * transpose: packs along the innermost dim, to the general packed format
+/// size.
+///
+/// An element at least as wide as its packing size gets unit lane_data, so f32
+/// is always [1, 1]. Leading dims are unit.
+static SmallVector<int64_t> get2DBlockLoadLaneData(
+ Type elemTy, int rank, bool hasTransform, bool hasTranspose,
+ const xegpu::uArch::Subgroup2DBlockLoadInstruction *uArchInstruction,
+ const xegpu::uArch::uArch *uArch) {
+ SmallVector<int64_t> laneData(rank, 1);
+ unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
+ unsigned packingSize = hasTransform || hasTranspose
+ ? uArch->getGeneralPackedFormatBitSize()
+ : uArchInstruction->getPackedFormatBitSize();
+ laneData[hasTransform ? rank - 2 : rank - 1] =
+ llvm::divideCeil(packingSize, bitwidth);
+ return laneData;
+}
+
/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
/// consumer layout (from its result's downstream uses) and validates it
/// against uArch constraints; if valid, the consumer's `inst_data` /
/// `sg_layout` are honored. Otherwise the helper falls back to defaults
-/// derived from uArch block parameters.
+/// derived from uArch block parameters. `lane_data` never comes from the
+/// consumer; see get2DBlockLoadLaneData.
xegpu::DistributeLayoutAttr
xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
VectorType resVecTy,
@@ -1795,7 +1823,6 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
consumerLayout.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> consumerLaneData =
consumerLayout.getEffectiveLaneDataAsInt();
- auto consumerOrderAttr = consumerLayout.getOrder();
assert(!consumerLaneLayout.empty() && !consumerLaneData.empty() &&
"Expected consumer layout to have lane_layout and lane_data");
@@ -1830,6 +1857,22 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
laneLayout.push_back(1);
}
+ // A consumer may ask for more elements per lane than the block load can
+ // pack, e.g. an f32 multi_reduction wanting 16. A downstream convert_layout
+ // reconciles the difference.
+ SmallVector<int64_t> laneData = get2DBlockLoadLaneData(
+ elemTy, rank, hasTransform, hasTranspose, uArchInstruction, uArch);
+
+ // Only a transposed block departs from the default of innermost-fastest.
+ DenseI32ArrayAttr orderAttr = nullptr;
+ if (hasTranspose) {
+ SmallVector<int32_t> order(rank);
+ for (int i = 0; i < rank; i++)
+ order[i] = rank - 1 - i;
+ std::swap(order[0], order[1]);
+ orderAttr = DenseI32ArrayAttr::get(context, order);
+ }
+
// See whether the consumer's inst_data satisfies the block constraints.
int64_t height = consumerInstData[rank - 2];
int64_t width = consumerInstData[rank - 1];
@@ -1838,23 +1881,36 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
if (llvm::is_contained(bWidths, static_cast<int>(width)) ||
(width % maxWidth == 0 && width / maxWidth < maxBlockCount)) {
if (llvm::is_contained(bHeights, static_cast<int>(height))) {
- return buildInstDataLayoutWithLane(context, consumerInstData,
- laneLayout, consumerLaneData,
- consumerOrderAttr);
+ assert(isValidLaneLayout(consumerInstData, laneLayout, laneData) &&
+ "Expected the load layout to satisfy uArch block constraints");
+ return buildInstDataLayoutWithLane(context, consumerInstData, laneLayout,
+ laneData, orderAttr);
}
}
- // if consumer instData size too small, try the larger one. like DPAS_MX's
- // scale is smaller than block load
+ // The consumer's inst_data is not a usable block, so take the uArch's.
+ // DPAS_MX scales land here. A scale whose innermost dim is too narrow to
+ // spread packed lanes over (width 16 for an 8-bit scale needing 16 lanes of
+ // 2 elements) cannot be loaded regularly; transformed, it packs down the
+ // column instead, which fits. A vertical lane_layout is already a transpose.
+ if (!hasTranspose && !hasTransform &&
+ consumerInstData[rank - 1] %
+ (laneLayout[rank - 1] * laneData[rank - 1]) !=
+ 0) {
+ hasTransform = true;
+ laneData = get2DBlockLoadLaneData(elemTy, rank, hasTransform, hasTranspose,
+ uArchInstruction, uArch);
+ }
+
auto instData = get2DBlockIOInstDataLayout(
dataShape, elemTy, uArchInstruction, hasTransform, hasTranspose);
// Shape not realizable as a 2D-block instruction; let the caller report it.
if (!instData)
return nullptr;
- assert(isValidLaneLayout(*instData, laneLayout, consumerLaneData) &&
+ assert(isValidLaneLayout(*instData, laneLayout, laneData) &&
"Expected the load layout to satisfy uArch block constraints");
- return buildInstDataLayoutWithLane(context, *instData, laneLayout,
- consumerLaneData, consumerOrderAttr);
+ return buildInstDataLayoutWithLane(context, *instData, laneLayout, laneData,
+ orderAttr);
}
if (layoutKind == xegpu::LayoutKind::Lane) {
assert(isValidLaneLayout(dataShape, consumerLaneLayout, consumerLaneData) &&
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index df5d95f1683fa..841f06545eb19 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -354,19 +354,19 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
gpu.module @test {
// CHECK-LABEL: func.func @bitcast_ui8_to_f4(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256x32xui8>) {
-// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<256x32xui8>
-// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>} : vector<256x32xui8> to vector<256x64xf4E2M1FN>
+// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<256x32xui8, #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<256x32xui8>
+// CHECK: %[[BC:.*]] = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>} : vector<256x32xui8> to vector<256x64xf4E2M1FN>
// CHECK: xegpu.convert_layout %[[BC]]
-// CHECK-SAME: <{target_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}>
+// CHECK-SAME: <{target_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
// CHECK-SAME: : vector<256x64xf4E2M1FN>
func.func @bitcast_ui8_to_f4(%arg0: memref<256x32xui8>) {
%0 = xegpu.create_nd_tdesc %arg0 : memref<256x32xui8> -> !xegpu.tensor_desc<256x32xui8>
%1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<256x32xui8> -> vector<256x32xui8>
%2 = vector.bitcast %1 : vector<256x32xui8> to vector<256x64xf4E2M1FN>
%3 = xegpu.convert_layout %2
- <{target_layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>}>
+ <{target_layout = #xegpu.layout<inst_data = [32, 64], lane_layout = [1, 16], lane_data = [1, 4]>}>
: vector<256x64xf4E2M1FN>
return
}
@@ -406,12 +406,12 @@ func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
// CHECK-SAME: !xegpu.tensor_desc<16x1024xf8E5M2, #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>> -> vector<16x1024xf8E5M2>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<1024x32xf8E5M2, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<1024x32xf8E5M2>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x32xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: <{layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
@@ -448,12 +448,12 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x1024xf8E5M2>, %arg1: memref<1024x32xf
// CHECK-SAME: !xegpu.tensor_desc<16x1024xf4E2M1FN, #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>> -> vector<16x1024xf4E2M1FN>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<1024x32xf4E2M1FN, #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>> -> vector<1024x32xf4E2M1FN>
-// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>>
-// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x32xf8E8M0FNU>
-// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
+// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %[[ARG3]] : memref<16x32xf8E8M0FNU> -> !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]>>
+// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x32xf8E8M0FNU, #xegpu.layout<inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]>> -> vector<16x32xf8E8M0FNU>
+// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %[[ARG4]] : memref<32x32xf8E8M0FNU> -> !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>>
+// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: <{layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
>From ea1d086ab80c128205ea10d6ee6366b30d541dd4 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 12 Sep 2026 19:02:18 +0000
Subject: [PATCH 2/4] [mlir][xegpu] Let the load lane-layout helper own
lane_layout and order
Fold the lane_layout and order derivation into the helper next to lane_data,
so the whole per-lane layout of a block load is decided in one place. The
consumer's lane_layout is no longer used as a base: like lane_data, it follows
from the block variant, spreading the subgroup along the innermost dim, or the
second-innermost one when transposed, which also fixes the order.
Drop the claim that lane_data counts the elements a lane holds in one register.
That is not necessarily how the data lands; lane_data follows the element width
and the packing granularity of the variant, which is all this needs to say.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 93 ++++++++++---------
1 file changed, 47 insertions(+), 46 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index b33470db6b902..d4924a05ad6b7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1759,31 +1759,52 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
}
-/// Returns the lane_data a 2D block load can deliver for `elemTy`. lane_data
-/// counts the elements a lane holds in one register, so it follows the packing
-/// granularity of the block variant:
+/// Returns the (lane_layout, lane_data, order) a 2D block load can deliver for
+/// `elemTy`. All three follow the element width and the block variant, not the
+/// consumer.
///
-/// * regular: packs along the innermost dim, to the load's packed format
-/// size.
-/// * transform: packs along the second-innermost dim, to the uArch's general
-/// packed format size. This is VNNI.
-/// * transpose: packs along the innermost dim, to the general packed format
-/// size.
+/// * regular: lane_layout [1, 16], lane_data packs elements narrower than
+/// the load's packed format size along the innermost dim.
+/// * transform: lane_layout [1, 16], lane_data packs along the
+/// second-innermost dim to the general packed format size. This
+/// is VNNI, so 16-bit and narrower only.
+/// * transpose: lane_layout [16, 1] with order [0, 1], lane_data packs along
+/// the innermost dim to the general packed format size.
///
-/// An element at least as wide as its packing size gets unit lane_data, so f32
-/// is always [1, 1]. Leading dims are unit.
-static SmallVector<int64_t> get2DBlockLoadLaneData(
- Type elemTy, int rank, bool hasTransform, bool hasTranspose,
+/// Leading dims are unit. `order` is left null for the two variants that keep
+/// the default innermost-fastest linearization.
+static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, DenseI32ArrayAttr>
+get2DBlockLoadLaneLayout(
+ mlir::MLIRContext *context, Type elemTy, int rank, int64_t subgroupSize,
+ bool hasTransform, bool hasTranspose,
const xegpu::uArch::Subgroup2DBlockLoadInstruction *uArchInstruction,
const xegpu::uArch::uArch *uArch) {
- SmallVector<int64_t> laneData(rank, 1);
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
unsigned packingSize = hasTransform || hasTranspose
? uArch->getGeneralPackedFormatBitSize()
: uArchInstruction->getPackedFormatBitSize();
- laneData[hasTransform ? rank - 2 : rank - 1] =
- llvm::divideCeil(packingSize, bitwidth);
- return laneData;
+
+ if (hasTranspose) {
+ laneLayout[rank - 2] = subgroupSize;
+ laneData[rank - 1] = llvm::divideCeil(packingSize, bitwidth);
+ SmallVector<int32_t> order(rank);
+ for (int i = 0; i < rank; i++)
+ order[i] = rank - 1 - i;
+ std::swap(order[0], order[1]);
+ return {laneLayout, laneData, DenseI32ArrayAttr::get(context, order)};
+ }
+
+ if (hasTransform) {
+ assert(bitwidth <= packingSize / 2 && "VNNI needs a sub-dword element");
+ laneLayout[rank - 1] = subgroupSize;
+ laneData[rank - 2] = llvm::divideCeil(packingSize, bitwidth);
+ return {laneLayout, laneData, nullptr};
+ }
+
+ laneLayout[rank - 1] = subgroupSize;
+ laneData[rank - 1] = llvm::divideCeil(packingSize, bitwidth);
+ return {laneLayout, laneData, nullptr};
}
/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
@@ -1845,33 +1866,12 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
auto [bWidths, bHeights, bCounts] = blockWHC.value();
- SmallVector<int64_t> laneLayout;
- // set the laneLayout to use consumer's LaneLayout as base, but adjust its
- // size to match the subgroupsize in case its original value is larger than
- // 1
- for (int i = 0; i < rank; i++) {
- if (consumerLaneLayout[i] > 1)
- laneLayout.push_back(std::max(static_cast<int64_t>(subgroupSize),
- consumerLaneLayout[i]));
- else
- laneLayout.push_back(1);
- }
-
// A consumer may ask for more elements per lane than the block load can
// pack, e.g. an f32 multi_reduction wanting 16. A downstream convert_layout
// reconciles the difference.
- SmallVector<int64_t> laneData = get2DBlockLoadLaneData(
- elemTy, rank, hasTransform, hasTranspose, uArchInstruction, uArch);
-
- // Only a transposed block departs from the default of innermost-fastest.
- DenseI32ArrayAttr orderAttr = nullptr;
- if (hasTranspose) {
- SmallVector<int32_t> order(rank);
- for (int i = 0; i < rank; i++)
- order[i] = rank - 1 - i;
- std::swap(order[0], order[1]);
- orderAttr = DenseI32ArrayAttr::get(context, order);
- }
+ auto [laneLayout, laneData, orderAttr] = get2DBlockLoadLaneLayout(
+ context, elemTy, rank, subgroupSize, hasTransform, hasTranspose,
+ uArchInstruction, uArch);
// See whether the consumer's inst_data satisfies the block constraints.
int64_t height = consumerInstData[rank - 2];
@@ -1883,8 +1883,8 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
if (llvm::is_contained(bHeights, static_cast<int>(height))) {
assert(isValidLaneLayout(consumerInstData, laneLayout, laneData) &&
"Expected the load layout to satisfy uArch block constraints");
- return buildInstDataLayoutWithLane(context, consumerInstData, laneLayout,
- laneData, orderAttr);
+ return buildInstDataLayoutWithLane(context, consumerInstData,
+ laneLayout, laneData, orderAttr);
}
}
@@ -1892,14 +1892,15 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
// DPAS_MX scales land here. A scale whose innermost dim is too narrow to
// spread packed lanes over (width 16 for an 8-bit scale needing 16 lanes of
// 2 elements) cannot be loaded regularly; transformed, it packs down the
- // column instead, which fits. A vertical lane_layout is already a transpose.
+ // column instead. A vertical lane_layout is already transposed.
if (!hasTranspose && !hasTransform &&
consumerInstData[rank - 1] %
(laneLayout[rank - 1] * laneData[rank - 1]) !=
0) {
hasTransform = true;
- laneData = get2DBlockLoadLaneData(elemTy, rank, hasTransform, hasTranspose,
- uArchInstruction, uArch);
+ std::tie(laneLayout, laneData, orderAttr) = get2DBlockLoadLaneLayout(
+ context, elemTy, rank, subgroupSize, hasTransform, hasTranspose,
+ uArchInstruction, uArch);
}
auto instData = get2DBlockIOInstDataLayout(
>From 5743c0d8343163fe7d956e927dcf2d1655da16a8 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 12 Sep 2026 19:05:30 +0000
Subject: [PATCH 3/4] [mlir][xegpu] Consolidate the block-IO and load
lane-layout helpers
The load path and the store / prefetch / DPAS paths were deriving per-lane
layouts with two separate helpers that differed only in which knobs they
exposed. compute2DBlockIOLaneLayoutAndData already took a shape, a subgroup
size, an element width and a packing size, and its `vnni` flag is exactly the
transform variant, so one function covers every caller.
Fold the load helper into it as compute2DBlockIOLaneLayout, which now also
returns the order and handles the transposed variant. Callers that cannot
transpose ignore the order, which always comes back empty for them.
The load path gains the shape clamp the block-IO helper already had, so a block
narrower than the subgroup no longer claims more lanes than it can supply.
No functional change: 4613 MLIR tests, only the 5 pre-existing failures.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 180 ++++++++----------
1 file changed, 79 insertions(+), 101 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d4924a05ad6b7..411cd86b92a01 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1240,20 +1240,35 @@ computeScatterIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
return {laneLayout, laneData};
}
-// Computes the per-lane layout and data for a 2D block load/store/prefetch:
-// lanes are spread across the subgroup along the last dim (or rank-2 if
-// transposed), and laneData packs sub-bitwidth elements along the packing dim.
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
- int64_t subgroupSize, int64_t bitwidth,
- int64_t packingSize, bool transform = false) {
+// Computes the (lane_layout, lane_data, order) a 2D block instruction can
+// deliver over `instShape`, from the element `bitwidth` and the variant's
+// `packingSize`. lane_data packs elements narrower than the packing size along
+// the packing dim: the innermost one, or the one above it when transformed
+// (VNNI). Lanes are spread along the innermost dim, or the one above it when
+// transposed, which makes that dim fastest-changing and so returns a reversed
+// `order`. `order` is empty for the variants that keep the default. Leading
+// dims are unit.
+static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>
+compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
+ int64_t bitwidth, int64_t packingSize,
+ bool transform = false, bool transpose = false) {
int64_t rank = instShape.size();
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- int kDim = transform ? rank - 2 : rank - 1;
- unsigned vnniFactor = packingSize / bitwidth;
- laneData[kDim] = bitwidth < packingSize ? vnniFactor : 1;
- laneLayout.back() =
- std::min(subgroupSize, instShape.back() / laneData.back());
+ assert(rank >= 2 && "Expected at least a 2D shape for a 2D block op");
+ SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1), order;
+
+ int64_t packDim = transform ? rank - 2 : rank - 1;
+ laneData[packDim] = llvm::divideCeil(packingSize, bitwidth);
+
+ int64_t laneDim = transpose ? rank - 2 : rank - 1;
+ laneLayout[laneDim] =
+ std::min(subgroupSize, instShape[laneDim] / laneData[laneDim]);
+
+ if (transpose) {
+ for (int64_t i = 0; i < rank; ++i)
+ order.push_back(rank - 1 - i);
+ std::swap(order[0], order[1]);
+ }
// assert that the lane layout and data fit in the inst shape
for (int64_t i = 0; i < rank; ++i) {
@@ -1262,7 +1277,7 @@ compute2DBlockIOLaneLayoutAndData(ArrayRef<int64_t> instShape,
"lane_layout * lane_data must evenly divide the inst shape");
(void)laneProduct;
}
- return {laneLayout, laneData};
+ return {laneLayout, laneData, order};
}
/// Computes the (lane_layout, lane_data) for a multi-reduction's source layout.
@@ -1439,18 +1454,18 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
auto subgroupSize = uArch->getSubgroupSize();
- auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
- aTy.getShape(), subgroupSize,
- aTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeA());
- auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ auto [laneLayoutA, laneDataA, orderA] =
+ compute2DBlockIOLaneLayout(aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB, orderB] = compute2DBlockIOLaneLayout(
bTy.getShape(), subgroupSize,
bTy.getElementType().getIntOrFloatBitWidth(),
uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
- auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
- cdTy.getShape(), subgroupSize,
- cdTy.getElementType().getIntOrFloatBitWidth(),
- cdTy.getElementType().getIntOrFloatBitWidth());
+ auto [laneLayoutCD, laneDataCD, orderCD] =
+ compute2DBlockIOLaneLayout(cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction);
if (!instDataVecs)
@@ -1579,18 +1594,18 @@ xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
auto subgroupSize = uArch->getSubgroupSize();
- auto [laneLayoutA, laneDataA] = compute2DBlockIOLaneLayoutAndData(
- aTy.getShape(), subgroupSize,
- aTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeA());
- auto [laneLayoutB, laneDataB] = compute2DBlockIOLaneLayoutAndData(
+ auto [laneLayoutA, laneDataA, orderA] =
+ compute2DBlockIOLaneLayout(aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ auto [laneLayoutB, laneDataB, orderB] = compute2DBlockIOLaneLayout(
bTy.getShape(), subgroupSize,
bTy.getElementType().getIntOrFloatBitWidth(),
uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
- auto [laneLayoutCD, laneDataCD] = compute2DBlockIOLaneLayoutAndData(
- cdTy.getShape(), subgroupSize,
- cdTy.getElementType().getIntOrFloatBitWidth(),
- cdTy.getElementType().getIntOrFloatBitWidth());
+ auto [laneLayoutCD, laneDataCD, orderCD] =
+ compute2DBlockIOLaneLayout(cdTy.getShape(), subgroupSize,
+ cdTy.getElementType().getIntOrFloatBitWidth(),
+ cdTy.getElementType().getIntOrFloatBitWidth());
auto instDataVecs = getDpasInstDataLayouts(aTy, bTy, cdTy, uArchInstruction);
if (!instDataVecs)
return std::nullopt;
@@ -1671,9 +1686,9 @@ xegpu::setupStoreNdAnchorLayout(xegpu::LayoutKind layoutKind,
// Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- dataShape, subgroupSize, bitwidth,
- uArchInstruction->getPackedFormatBitSize());
+ auto [laneLayout, laneData, order] =
+ compute2DBlockIOLaneLayout(dataShape, subgroupSize, bitwidth,
+ uArchInstruction->getPackedFormatBitSize());
if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
@@ -1727,9 +1742,9 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
// Compute the default 2D block IO lane layout / lane data.
unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
- dataShape, subgroupSize, bitwidth,
- uArchInstruction->getPackedFormatBitSize());
+ auto [laneLayout, laneData, order] =
+ compute2DBlockIOLaneLayout(dataShape, subgroupSize, bitwidth,
+ uArchInstruction->getPackedFormatBitSize());
if (layoutKind == xegpu::LayoutKind::Lane)
return buildLaneLayout(context, laneLayout, laneData);
@@ -1759,54 +1774,6 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
}
-/// Returns the (lane_layout, lane_data, order) a 2D block load can deliver for
-/// `elemTy`. All three follow the element width and the block variant, not the
-/// consumer.
-///
-/// * regular: lane_layout [1, 16], lane_data packs elements narrower than
-/// the load's packed format size along the innermost dim.
-/// * transform: lane_layout [1, 16], lane_data packs along the
-/// second-innermost dim to the general packed format size. This
-/// is VNNI, so 16-bit and narrower only.
-/// * transpose: lane_layout [16, 1] with order [0, 1], lane_data packs along
-/// the innermost dim to the general packed format size.
-///
-/// Leading dims are unit. `order` is left null for the two variants that keep
-/// the default innermost-fastest linearization.
-static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>, DenseI32ArrayAttr>
-get2DBlockLoadLaneLayout(
- mlir::MLIRContext *context, Type elemTy, int rank, int64_t subgroupSize,
- bool hasTransform, bool hasTranspose,
- const xegpu::uArch::Subgroup2DBlockLoadInstruction *uArchInstruction,
- const xegpu::uArch::uArch *uArch) {
- SmallVector<int64_t> laneLayout(rank, 1), laneData(rank, 1);
- unsigned bitwidth = elemTy.getIntOrFloatBitWidth();
- unsigned packingSize = hasTransform || hasTranspose
- ? uArch->getGeneralPackedFormatBitSize()
- : uArchInstruction->getPackedFormatBitSize();
-
- if (hasTranspose) {
- laneLayout[rank - 2] = subgroupSize;
- laneData[rank - 1] = llvm::divideCeil(packingSize, bitwidth);
- SmallVector<int32_t> order(rank);
- for (int i = 0; i < rank; i++)
- order[i] = rank - 1 - i;
- std::swap(order[0], order[1]);
- return {laneLayout, laneData, DenseI32ArrayAttr::get(context, order)};
- }
-
- if (hasTransform) {
- assert(bitwidth <= packingSize / 2 && "VNNI needs a sub-dword element");
- laneLayout[rank - 1] = subgroupSize;
- laneData[rank - 2] = llvm::divideCeil(packingSize, bitwidth);
- return {laneLayout, laneData, nullptr};
- }
-
- laneLayout[rank - 1] = subgroupSize;
- laneData[rank - 1] = llvm::divideCeil(packingSize, bitwidth);
- return {laneLayout, laneData, nullptr};
-}
-
/// Sets up the anchor layout for a load_nd operation. LoadNd takes a
/// consumer layout (from its result's downstream uses) and validates it
/// against uArch constraints; if valid, the consumer's `inst_data` /
@@ -1866,12 +1833,21 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
return nullptr;
auto [bWidths, bHeights, bCounts] = blockWHC.value();
- // A consumer may ask for more elements per lane than the block load can
- // pack, e.g. an f32 multi_reduction wanting 16. A downstream convert_layout
+ // lane_layout and lane_data are the block load's own, not the consumer's: a
+ // consumer may ask for more elements per lane than the hardware can pack,
+ // e.g. an f32 multi_reduction wanting 16. A downstream convert_layout
// reconciles the difference.
- auto [laneLayout, laneData, orderAttr] = get2DBlockLoadLaneLayout(
- context, elemTy, rank, subgroupSize, hasTransform, hasTranspose,
- uArchInstruction, uArch);
+ unsigned packingSize = hasTransform || hasTranspose
+ ? uArch->getGeneralPackedFormatBitSize()
+ : uArchInstruction->getPackedFormatBitSize();
+ auto [laneLayout, laneData, order] = compute2DBlockIOLaneLayout(
+ dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(), packingSize,
+ hasTransform, hasTranspose);
+ DenseI32ArrayAttr orderAttr =
+ order.empty()
+ ? nullptr
+ : DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
// See whether the consumer's inst_data satisfies the block constraints.
int64_t height = consumerInstData[rank - 2];
@@ -1898,9 +1874,9 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
(laneLayout[rank - 1] * laneData[rank - 1]) !=
0) {
hasTransform = true;
- std::tie(laneLayout, laneData, orderAttr) = get2DBlockLoadLaneLayout(
- context, elemTy, rank, subgroupSize, hasTransform, hasTranspose,
- uArchInstruction, uArch);
+ std::tie(laneLayout, laneData, order) = compute2DBlockIOLaneLayout(
+ dataShape, subgroupSize, elemTy.getIntOrFloatBitWidth(),
+ uArch->getGeneralPackedFormatBitSize(), hasTransform, hasTranspose);
}
auto instData = get2DBlockIOInstDataLayout(
@@ -2185,7 +2161,7 @@ xegpu::completeBlockStoreLaneLayoutFromInstData(
return specifiedLayout;
auto *context = specifiedLayout.getContext();
- auto [laneLayout, laneData] = compute2DBlockIOLaneLayoutAndData(
+ auto [laneLayout, laneData, order] = compute2DBlockIOLaneLayout(
specifiedInstData, subgroupSize, elemTy.getIntOrFloatBitWidth(),
uArchInstruction->getPackedFormatBitSize());
if (!isValidLaneLayout(specifiedInstData, laneLayout, laneData))
@@ -2266,20 +2242,22 @@ xegpu::completeDpasLaneLayoutFromInstData(xegpu::DistributeLayoutAttr aLayout,
auto subgroupSize = uArch->getSubgroupSize();
llvm::SmallVector<int64_t> laneLayoutA, laneDataA, laneLayoutB, laneDataB,
laneLayoutCD, laneDataCD;
+ // DPAS operands are never transposed, so the order comes back empty.
+ llvm::SmallVector<int64_t> orderA, orderB, orderCD;
SmallVector<int64_t> instDataA = aLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> instDataB = bLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> instDataCD = cdLayout.getEffectiveInstDataAsInt();
if (isa<xegpu::uArch::Xe2, xegpu::uArch::Xe3>(uArch)) {
- std::tie(laneLayoutA, laneDataA) = compute2DBlockIOLaneLayoutAndData(
- aTy.getShape(), subgroupSize,
- aTy.getElementType().getIntOrFloatBitWidth(),
- uArchInstruction->getPackedFormatBitSizeA());
- std::tie(laneLayoutB, laneDataB) = compute2DBlockIOLaneLayoutAndData(
+ std::tie(laneLayoutA, laneDataA, orderA) =
+ compute2DBlockIOLaneLayout(aTy.getShape(), subgroupSize,
+ aTy.getElementType().getIntOrFloatBitWidth(),
+ uArchInstruction->getPackedFormatBitSizeA());
+ std::tie(laneLayoutB, laneDataB, orderB) = compute2DBlockIOLaneLayout(
bTy.getShape(), subgroupSize,
bTy.getElementType().getIntOrFloatBitWidth(),
uArchInstruction->getPackedFormatBitSizeB(), /*vnni=*/true);
- std::tie(laneLayoutCD, laneDataCD) = compute2DBlockIOLaneLayoutAndData(
+ std::tie(laneLayoutCD, laneDataCD, orderCD) = compute2DBlockIOLaneLayout(
cdTy.getShape(), subgroupSize,
cdTy.getElementType().getIntOrFloatBitWidth(),
cdTy.getElementType().getIntOrFloatBitWidth());
>From 471ed3d5f8425c12d7522224327baf261225357a Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 18 Sep 2026 02:45:47 +0000
Subject: [PATCH 4/4] remove comments, correct uArch table
---
mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h | 2 +-
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 6 ++----
2 files changed, 3 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index d4ba2391729e1..cad9c07f283bc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -313,7 +313,7 @@ struct Subgroup2DBlockLoadInstruction : public Instruction,
{{16, false, true, false}, {kWidth16, kHeightAtLeast16, kCount1}},
{{32, false, true, false}, {kWidth8, kHeightAtLeast16, kCount1}},
// 4-bit elements (sub-byte):
- {{4, false, false, false}, {kWidthAtLeast32, kHeightAtLeast1, kCount2}},
+ {{4, false, false, false}, {kWidth64, kHeightAtLeast1, kCount2}},
{{4, false, false, true}, {kWidth32, kHeightAtLeast8, kCount4Only}},
{{4, true, false, false}, {kWidth16, kHeight64, kCount4}},
{{4, false, true, false}, {kWidth64, kHeightAtLeast16, kCount1}}};
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 411cd86b92a01..d719559925932 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1778,8 +1778,7 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
/// consumer layout (from its result's downstream uses) and validates it
/// against uArch constraints; if valid, the consumer's `inst_data` /
/// `sg_layout` are honored. Otherwise the helper falls back to defaults
-/// derived from uArch block parameters. `lane_data` never comes from the
-/// consumer; see get2DBlockLoadLaneData.
+/// derived from uArch block parameters.
xegpu::DistributeLayoutAttr
xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
VectorType resVecTy,
@@ -1835,8 +1834,7 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
// lane_layout and lane_data are the block load's own, not the consumer's: a
// consumer may ask for more elements per lane than the hardware can pack,
- // e.g. an f32 multi_reduction wanting 16. A downstream convert_layout
- // reconciles the difference.
+ // e.g. an f32 multi_reduction wanting 16.
unsigned packingSize = hasTransform || hasTranspose
? uArch->getGeneralPackedFormatBitSize()
: uArchInstruction->getPackedFormatBitSize();
More information about the Mlir-commits
mailing list