[Mlir-commits] [mlir] [mlir][xegpu] Derive load_nd lane_data from the hardware block variant (PR #223164)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 12 11:08:40 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
setupLoadNdAnchorLayout copied lane_data from the consumer layout. But lane_data follows from the element width and the block variant the hardware will use, not from what the consumer wants.
This PR adds get2DBlockLoadLaneData, deriving lane_data from the packing size the uArch already reports: the load's packed format size along the innermost dim for a regular 2d block load, the general packed format size along the second-innermost dim for transform (VNNI), or along the innermost dim for transpose.
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: width 16 for an 8-bit scale cannot spread 16 lanes of 2 elements over 32 columns, so it is loaded transformed, packing down the column instead.
scale_a inst_data = [16, 32], lane_layout = [16, 1], lane_data = [1, 4], order = [0, 1]
scale_b inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]
Also correct the uArch table: an 8-bit regular block load packs two elements per lane, so 16 lanes need 32 columns. Its supported width is 32, not {16, 32}.
assisted-by-claude
---
Full diff: https://github.com/llvm/llvm-project/pull/223164.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h (+1-2)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+66-10)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+18-18)
``````````diff
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>
``````````
</details>
https://github.com/llvm/llvm-project/pull/223164
More information about the Mlir-commits
mailing list