[Mlir-commits] [mlir] [mlir][xegpu] Check uArch block shapes in VectorToXeGPU transfer lowering (PR #217179)
Igor Zamyatin
llvmlistbot at llvm.org
Tue Aug 18 17:50:57 PDT 2026
https://github.com/Garra1980 created https://github.com/llvm/llvm-project/pull/217179
vector.transfer_read/transfer_write were lowered to xegpu.load_nd / xegpu.store_nd whenever the target chip was pvc, bmg or cri and the transfer had rank >= 2, without asking whether the target's subgroup 2D block instructions can access the requested tile at all.
Changing it to query the uArch 2D block load/store description up front and keep the block path only for tiles whose two innermost dims are each a multiple of a supported block size- the property the later layout propagation and blocking passes rely on when they split a tile into hardware-sized blocks. Everything else takes the scattered load_gather/store_scatter fallback that both patterns already have.
Deriving 2D block support from the uArch instruction registry also replaces the hardcoded chip list, resolving the TODO left there; the three chips that registry covers are the same three that were listed.
assisted by claude
>From 8f63f37f94704c5a7741c26d30511a7154e4aefe Mon Sep 17 00:00:00 2001
From: Garra1980 <igor.zamyatin at intel.com>
Date: Fri, 14 Aug 2026 04:16:56 +0000
Subject: [PATCH] [mlir][xegpu] Check uArch block shapes in VectorToXeGPU
transfer lowering
vector.transfer_read/transfer_write were lowered to xegpu.load_nd /
xegpu.store_nd whenever the target chip was pvc, bmg or cri and the transfer
had rank >= 2, without asking whether the target's subgroup 2D block
instructions can access the requested tile at all. Reading e.g.
vector<1x8x16x1xf16> out of memref<1x24x1024x1xf16> leaves a 16x1 f16 tile,
which no 2D block instruction can provide because its block width must be a
multiple of 16 f16 elements. The lowering still emitted
%0 = xegpu.create_nd_tdesc %arg0 : memref<1x24x1024x1xf16>
-> !xegpu.tensor_desc<1x8x16x1xf16, ...>
%1 = xegpu.load_nd %0[0, %off0, %off1, 0] : ... -> vector<1x8x16x1xf16>
and the pipeline then crashed downstream: xegpu-propagate-layout finds no
supported block width for the innermost dimension and assigns
inst_data = [1, 1, 16, -1], after which xegpu-blocking asserts while building
vector<1x1x16x-1xf16>.
Query the uArch 2D block load/store description up front instead, and keep the
block path only for tiles whose two innermost extents are each a multiple of a
supported block extent - the property the later layout propagation and blocking
passes rely on when they split a tile into hardware-sized blocks. Everything
else takes the scattered load_gather/store_scatter fallback that both patterns
already have. Deriving 2D block support from the uArch instruction registry
also replaces the hardcoded chip list, resolving the TODO left there; the three
chips that registry covers are the same three that were listed.
Four cases in the tests move to the scattered path because no 2D block
instruction can access their tiles, all of which crashed or failed the
gpu-lower-to-xevm pipeline before: the 8D f32 read and write (innermost extent
2), the transposed f16 read (a transposing 16-bit block load needs a
16-element-wide block, so its 16x8 tile has no block shape) and the 16x8 f16
write in transpose_1x1024x24x64. Transposed f32 reads keep using the block
path, since the transposing 32-bit block load does support an 8-element width.
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 120 ++++++++++++-----
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 127 ++++++++++--------
.../transfer-write-to-xegpu.mlir | 64 ++++++---
3 files changed, 199 insertions(+), 112 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 8a45836426931..a821f031233cd 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -21,6 +21,8 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -82,6 +84,51 @@ static bool isInnermostTwoDimsTransposed(AffineMap map) {
map.getResult(numResults - 1) == getAffineDimExpr(numInputs - 2, ctx);
}
+// Return true if `uArch` can transfer a `shape`-shaped tile of `elemTy`
+// elements with the subgroup 2D block instruction `instKind`.
+//
+// A 2D block instruction accesses the two innermost dimensions of the tile;
+// leading dimensions are unrolled into a sequence of 2D accesses. The
+// innermost 2D tile does not have to match a hardware block exactly, because
+// the later XeGPU passes (layout propagation and blocking) split it into
+// hardware-sized blocks - but that split only exists when each of the two
+// innermost extents is a multiple of a supported block extent. Transfers
+// without such a split cannot be expressed with block instructions at all and
+// must use a different lowering.
+static bool isSupportedBlockShape(const xegpu::uArch::uArch *uArch,
+ xegpu::uArch::InstructionKind instKind,
+ ArrayRef<int64_t> shape, Type elemTy,
+ bool hasTranspose = false) {
+ if (shape.size() < 2)
+ return false;
+ if (!uArch || !uArch->isSupportedInstruction(instKind))
+ return false;
+
+ const xegpu::uArch::Instruction *inst = uArch->getInstruction(instKind);
+ const xegpu::uArch::BlockIOInstructionInterface *blockInst = nullptr;
+ if (const auto *load =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(inst))
+ blockInst = load;
+ else if (const auto *store =
+ dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(inst))
+ blockInst = store;
+ if (!blockInst)
+ return false;
+
+ // A missing entry means the element type itself is not supported.
+ std::optional<xegpu::uArch::BlockIOInstructionInterface::BlockShapes>
+ blockShapes = blockInst->getBlockWidthHeightCount(
+ elemTy, /*hasTransform=*/false, hasTranspose);
+ if (!blockShapes)
+ return false;
+
+ auto [widths, heights, counts] = *blockShapes;
+ int width = static_cast<int>(shape.back());
+ int height = static_cast<int>(shape[shape.size() - 2]);
+ return xegpu::getLargestDivisor(width, widths) != -1 &&
+ xegpu::getLargestDivisor(height, heights) != -1;
+}
+
static LogicalResult transferPreconditions(PatternRewriter &rewriter,
VectorTransferOpInterface xferOp) {
if (xferOp.getMask())
@@ -603,10 +650,8 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
return success();
}
- // TODO: This check needs to be replaced with proper uArch capability check.
- auto chip = xegpu::getChipStr(readOp);
- bool hasBlockLoadSupport =
- (chip == "pvc" || chip == "bmg" || chip == "cri");
+ const xegpu::uArch::uArch *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(readOp).value_or(""));
// An nd block load can realize a minor-identity map directly, or an
// innermost-two-dims transpose via a trailing vector.transpose. Any other
@@ -615,29 +660,36 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
AffineMap readMap = readOp.getPermutationMap();
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
- // Prefer an nd block load. It requires HW block-load support, a vector of
- // rank >= 2 backed by a scalar-element memref, and a map the block load can
- // realize. 1D vectors use the scattered xegpu.load path instead, which has
- // a richer interface (e.g. layout capabilities). Out-of-bounds reads are
- // allowed as long as the padding matches load_nd's implicit zero padding.
+ // A block load transfers the tile as it is laid out in memory, so a
+ // transposing read describes a tile whose innermost two dims are swapped.
+ Type elementType = loadedVecTy.getElementType();
+ SmallVector<int64_t> descShape(loadedVecTy.getShape());
+ if (isTransposeLoad) {
+ size_t rank = descShape.size();
+ assert(rank >= 2 && "Transpose requires at least 2 dimensions");
+ std::swap(descShape[rank - 1], descShape[rank - 2]);
+ }
+
+ // Prefer an nd block load. It requires a vector of rank >= 2 backed by a
+ // scalar-element memref, a map the block load can realize, and a tile shape
+ // the target's 2D block load can access. 1D vectors use the scattered
+ // xegpu.load path instead, which has a richer interface (e.g. layout
+ // capabilities). Out-of-bounds reads are allowed as long as the padding
+ // matches load_nd's implicit zero padding.
bool canLowerToLoadNd =
- hasBlockLoadSupport && loadedVecTy.getRank() > 1 &&
+ loadedVecTy.getRank() > 1 &&
(readMap.isMinorIdentity() || isTransposeLoad) &&
readMemTy.getElementType().isIntOrFloat() &&
- (!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
+ (!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding())) &&
+ isSupportedBlockShape(
+ uArch, xegpu::uArch::InstructionKind::Subgroup2DBlockLoad,
+ descShape, elementType, /*hasTranspose=*/isTransposeLoad);
if (canLowerToLoadNd) {
- auto elementType = loadedVecTy.getElementType();
-
- SmallVector<int64_t> descShape(loadedVecTy.getShape());
- if (isTransposeLoad) {
- // If load is transposed, simply swap the last two dimensions of the
- // loaded vector type to get the descriptor shape.
- size_t rank = descShape.size();
- assert(rank >= 2 && "Transpose requires at least 2 dimensions");
- std::swap(descShape[rank - 1], descShape[rank - 2]);
+ // The load produces the memory-ordered tile; the transpose below restores
+ // the shape the transfer_read asked for.
+ if (isTransposeLoad)
loadedVecTy = VectorType::get(descShape, elementType);
- }
auto descType = xegpu::TensorDescType::get(
descShape, elementType, /*array_length=*/1,
/*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
@@ -722,20 +774,22 @@ struct TransferWriteLowering
return success();
}
- // TODO: This check needs to be replaced with proper uArch capability check.
- auto chip = xegpu::getChipStr(writeOp);
- bool hasBlockStoreSupport =
- (chip == "pvc" || chip == "bmg" || chip == "cri");
+ const xegpu::uArch::uArch *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(writeOp).value_or(""));
- // Prefer an nd block store. It requires HW block-store support, a vector of
- // rank >= 2 backed by a scalar-element memref, and a minor-identity map
- // (block stores have no transpose support). 1D vectors use the scattered
- // xegpu.store path instead, which has a richer interface. Out-of-bounds
- // writes are handled by the descriptor's boundary check.
+ // Prefer an nd block store. It requires a vector of rank >= 2 backed by a
+ // scalar-element memref, a minor-identity map (block stores have no
+ // transpose support), and a tile shape the target's 2D block store can
+ // access. 1D vectors use the scattered xegpu.store path instead, which has
+ // a richer interface. Out-of-bounds writes are handled by the descriptor's
+ // boundary check.
AffineMap map = writeOp.getPermutationMap();
- bool canLowerToStoreNd = hasBlockStoreSupport && vecTy.getRank() > 1 &&
- map.isMinorIdentity() &&
- writeMemTy.getElementType().isIntOrFloat();
+ bool canLowerToStoreNd =
+ vecTy.getRank() > 1 && map.isMinorIdentity() &&
+ writeMemTy.getElementType().isIntOrFloat() &&
+ isSupportedBlockShape(
+ uArch, xegpu::uArch::InstructionKind::Subgroup2DBlockStore,
+ vecTy.getShape(), vecTy.getElementType());
if (canLowerToStoreNd) {
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 4cc4a0db5b63c..a54ec48bdec48 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -340,6 +340,31 @@ gpu.func @load_high_dim_vector(%source: memref<16x32x64xf32>,
// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16x32xindex>, vector<8x16x32xi1> -> vector<8x16x32xf32>
}
+// -----
+// The unit innermost dimension leaves a 16x1 f16 tile to transfer, and a 2D
+// block load needs a block width that is a multiple of 16 f16 elements. The
+// read must use the scattered path rather than a block load whose shape no
+// hardware instruction can provide.
+gpu.module @xevm_module {
+gpu.func @load_high_dim_unsupported_block_width(%source: memref<1x24x1024x1xf16>,
+ %offset: index) -> vector<1x8x16x1xf16> {
+ %pad = ub.poison : f16
+ %c0 = arith.constant 0 : index
+ %0 = vector.transfer_read %source[%c0, %offset, %offset, %c0], %pad
+ {in_bounds = [true, true, true, true]}
+ : memref<1x24x1024x1xf16>, vector<1x8x16x1xf16>
+ gpu.return %0 : vector<1x8x16x1xf16>
+}
+
+// CHECK-LABEL: @load_high_dim_unsupported_block_width(
+// CHECK-SAME: %[[SRC:.+]]: memref<1x24x1024x1xf16>,
+// CHECK-NOT: xegpu.create_nd_tdesc
+// CHECK-NOT: xegpu.load_nd
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<1x24x1024x1xf16> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%{{.+}}], %{{.+}} : i64, vector<1x8x16x1xindex>, vector<1x8x16x1xi1> -> vector<1x8x16x1xf16>
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_8D_vector(%source: memref<2x2x2x2x2x2x2x2xf32>,
@@ -350,25 +375,22 @@ gpu.func @load_8D_vector(%source: memref<2x2x2x2x2x2x2x2xf32>,
gpu.return %0 : vector<2x2x2x2x2x2x2x2xf32>
}
-// LOAD-ND-LABEL: @load_8D_vector(
-// LOAD-ND-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>,
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32>
-// LOAD-ND-SAME: -> !xegpu.tensor_desc<2x2x2x2x2x2x2x2xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
-// LOAD-ND-SAME: -> vector<2x2x2x2x2x2x2x2xf32>
-
-// LOAD-GATHER-LABEL: @load_8D_vector(
-// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>,
-// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<2x2x2x2x2x2x2x2xi1>
-// LOAD-GATHER-COUNT8: vector.step
-// LOAD-GATHER-COUNT7: vector.shape_cast
-// LOAD-GATHER-COUNT8: vector.broadcast {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// LOAD-GATHER-COUNT7: arith.addi {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x2x2x2x2x2x2x2xindex>
-// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32> -> index
-// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<2x2x2x2x2x2x2x2xindex>, vector<2x2x2x2x2x2x2x2xi1> -> vector<2x2x2x2x2x2x2x2xf32>
+// The innermost extent of 2 f32 elements is not a multiple of any block width
+// the 2D block load supports (16 f32 elements), so the read uses the scattered
+// path on every target.
+// CHECK-LABEL: @load_8D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>,
+// CHECK-NOT: xegpu.load_nd
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<2x2x2x2x2x2x2x2xi1>
+// CHECK-COUNT-8: vector.step
+// CHECK-COUNT-7: vector.shape_cast
+// CHECK-COUNT-8: vector.broadcast {{.*}} to vector<2x2x2x2x2x2x2x2xindex>
+// CHECK-COUNT-7: arith.addi {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x2x2x2x2x2x2x2xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<2x2x2x2x2x2x2x2xindex>, vector<2x2x2x2x2x2x2x2xi1> -> vector<2x2x2x2x2x2x2x2xf32>
}
// -----
@@ -382,23 +404,21 @@ gpu.func @load_transpose_f16(%source: memref<32x64xf16>,
gpu.return %0 : vector<8x16xf16>
}
-// LOAD-ND-LABEL: @load_transpose_f16(
-// LOAD-ND: %[[LOAD:.*]] = xegpu.load_nd
-// LOAD-ND: vector.transpose %[[LOAD]], [1, 0] : vector<16x8xf16> to vector<8x16xf16>
-
-// LOAD-GATHER-LABEL: @load_transpose_f16(
-// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<32x64xf16>,
-// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
-// LOAD-GATHER-COUNT2: vector.step
-// LOAD-GATHER-COUNT2: vector.shape_cast
-// LOAD-GATHER-COUNT2: vector.broadcast
-// LOAD-GATHER-COUNT2: arith.muli {{.*}} : index
-// LOAD-GATHER-COUNT2: arith.addi {{.*}} : index
-// LOAD-GATHER: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>
-// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}}: vector<8x16xindex>
-// LOAD-GATHER: %[[COLLAPSE:.*]] = memref.extract_aligned_pointer_as_index %arg0 : memref<32x64xf16> -> index
-// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// LOAD-GATHER: %[[LOAD:.*]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf16>
+// Unlike the f32 case above, a transposing block load of 16-bit elements needs
+// a block 16 elements wide, so the 16x8 tile this read would have to fetch has
+// no hardware block shape and the read uses the scattered path instead.
+// CHECK-LABEL: @load_transpose_f16(
+// CHECK-SAME: %[[SRC:.+]]: memref<32x64xf16>,
+// CHECK-NOT: xegpu.load_nd
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
+// CHECK-COUNT-2: vector.step
+// CHECK: vector.shape_cast
+// CHECK-COUNT-2: vector.broadcast {{.*}} to vector<8x16xindex>
+// CHECK: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}}: vector<8x16xindex>
+// CHECK: %[[COLLAPSE:.*]] = memref.extract_aligned_pointer_as_index %arg0 : memref<32x64xf16> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: %[[LOAD:.*]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf16>
}
// -----
@@ -688,34 +708,25 @@ gpu.func @transpose_1x1024x24x64(
// The vector.transpose is folded into the transfer_read via
// CombineTransferReadOpTranspose, giving the read a mid-vector permutation map
// (d0, d1, d2, d3) -> (d0, d2, d1, d3). An nd block load can only realize an
-// innermost-two-dims transpose, so the read falls back to the scattered path
-// while the identity-map write still lowers to store_nd.
-// LOAD-ND-LABEL: @transpose_1x1024x24x64
-// LOAD-ND-DAG: %[[C1536:.+]] = arith.constant 1536 : index
-// LOAD-ND-DAG: %[[C64:.+]] = arith.constant 64 : index
-// LOAD-ND: arith.muli %{{.+}}, %[[C1536]] : index
-// LOAD-ND: arith.muli %block_id_x, %[[C64]] : index
-// LOAD-ND: %[[VEC:.+]] = xegpu.load {{.*}} -> vector<1x1x16x8xf16>
-// LOAD-ND: %[[WDESC:.+]] = xegpu.create_nd_tdesc %arg1 : memref<1x24x1024x64xf16>
-// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x1x16x8xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
-// LOAD-ND: xegpu.store_nd %[[VEC]], %[[WDESC]]
-
-// LOAD-GATHER-LABEL: @transpose_1x1024x24x64
-// LOAD-GATHER-DAG: %[[C1536:.+]] = arith.constant 1536 : index
-// LOAD-GATHER-DAG: %[[C64:.+]] = arith.constant 64 : index
-// LOAD-GATHER-DAG: %[[C65536:.+]] = arith.constant 65536 : index
+// innermost-two-dims transpose, so the read falls back to the scattered path.
+// The identity-map write cannot use a block store either: its 16x8 f16 tile is
+// narrower than the 16-element block width the 2D block store requires.
+// CHECK-LABEL: @transpose_1x1024x24x64
+// CHECK-DAG: %[[C1536:.+]] = arith.constant 1536 : index
+// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
+// CHECK-DAG: %[[C65536:.+]] = arith.constant 65536 : index
// Read from memref<1x1024x24x64xf16>, strides [1572864, 1536, 64, 1].
// Scalar base offset: seq_off * 1536 (original dim1 stride),
// block_id_x * 64 (original dim2 stride).
-// LOAD-GATHER: arith.muli %{{.+}}, %[[C1536]] : index
-// LOAD-GATHER: arith.muli %block_id_x, %[[C64]] : index
-// LOAD-GATHER: xegpu.load {{.*}} -> vector<1x1x16x8xf16>
+// CHECK: arith.muli %{{.+}}, %[[C1536]] : index
+// CHECK: arith.muli %block_id_x, %[[C64]] : index
+// CHECK: xegpu.load {{.*}} -> vector<1x1x16x8xf16>
// Write to memref<1x24x1024x64xf16>, strides [1572864, 65536, 64, 1].
// Scalar base offset: block_id_x * 65536 (original dim1 stride),
// seq_off * 64 (original dim2 stride).
-// LOAD-GATHER: arith.muli %block_id_x, %[[C65536]] : index
-// LOAD-GATHER: arith.muli %{{.+}}, %[[C64]] : index
-// LOAD-GATHER: xegpu.store {{.*}}
+// CHECK: arith.muli %block_id_x, %[[C65536]] : index
+// CHECK: arith.muli %{{.+}}, %[[C64]] : index
+// CHECK: xegpu.store {{.*}}
}
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index dad2740f5c0ea..1b71e533dc64b 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -211,6 +211,32 @@ gpu.func @store_high_dim_vector(%vec: vector<8x16x32xf32>,
// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : vector<8x16x32xf32>, i64, vector<8x16x32xindex>, vector<8x16x32xi1>
}
+// -----
+// The unit innermost dimension leaves a 16x1 f16 tile to transfer, and a 2D
+// block store needs a block width that is a multiple of 16 f16 elements. The
+// write must use the scattered path rather than a block store whose shape no
+// hardware instruction can provide.
+gpu.module @xevm_module {
+gpu.func @store_high_dim_unsupported_block_width(%vec: vector<1x8x16x1xf16>,
+ %source: memref<1x24x1024x1xf16>, %offset: index) {
+ %c0 = arith.constant 0 : index
+ vector.transfer_write %vec, %source[%c0, %offset, %offset, %c0]
+ {in_bounds = [true, true, true, true]}
+ : vector<1x8x16x1xf16>, memref<1x24x1024x1xf16>
+ gpu.return
+}
+
+// CHECK-LABEL: @store_high_dim_unsupported_block_width(
+// CHECK-SAME: %[[VEC:.+]]: vector<1x8x16x1xf16>,
+// CHECK-SAME: %[[SRC:.+]]: memref<1x24x1024x1xf16>
+// CHECK-NOT: xegpu.create_nd_tdesc
+// CHECK-NOT: xegpu.store_nd
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<1x24x1024x1xf16> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%{{.+}}], %{{.+}} : vector<1x8x16x1xf16>, i64, vector<1x8x16x1xindex>, vector<1x8x16x1xi1>
+
+}
+
// -----
gpu.module @xevm_module {
gpu.func @store_8D_vector(%vec: vector<2x2x2x2x2x2x2x2xf32>,
@@ -221,27 +247,23 @@ gpu.func @store_8D_vector(%vec: vector<2x2x2x2x2x2x2x2xf32>,
gpu.return
}
-// STORE-ND-LABEL: @store_8D_vector(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<2x2x2x2x2x2x2x2xf32>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32>
-// STORE-ND-SAME: -> !xegpu.tensor_desc<2x2x2x2x2x2x2x2xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
-// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]]
-// STORE-ND-SAME: : vector<2x2x2x2x2x2x2x2xf32>, !xegpu.tensor_desc<2x2x2x2x2x2x2x2xf32
-
-// STORE-SCATTER-LABEL: @store_8D_vector(
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<2x2x2x2x2x2x2x2xf32>,
-// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>
-// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<2x2x2x2x2x2x2x2xi1>
-// STORE-SCATTER-COUNT8: vector.step
-// STORE-SCATTER-COUNT7: vector.shape_cast
-// STORE-SCATTER-COUNT8: vector.broadcast {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// STORE-SCATTER-COUNT7: arith.addi {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// STORE-SCATTER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x2x2x2x2x2x2x2xindex>
-// STORE-SCATTER: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// STORE-SCATTER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32> -> index
-// STORE-SCATTER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : vector<2x2x2x2x2x2x2x2xf32>, i64, vector<2x2x2x2x2x2x2x2xindex>, vector<2x2x2x2x2x2x2x2xi1>
+// The innermost extent of 2 f32 elements is not a multiple of any block width
+// the 2D block store supports (16 f32 elements), so the write uses the
+// scattered path on every target.
+// CHECK-LABEL: @store_8D_vector(
+// CHECK-SAME: %[[VEC:.+]]: vector<2x2x2x2x2x2x2x2xf32>,
+// CHECK-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>
+// CHECK-NOT: xegpu.store_nd
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<2x2x2x2x2x2x2x2xi1>
+// CHECK-COUNT-8: vector.step
+// CHECK-COUNT-7: vector.shape_cast
+// CHECK-COUNT-8: vector.broadcast {{.*}} to vector<2x2x2x2x2x2x2x2xindex>
+// CHECK-COUNT-7: arith.addi {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x2x2x2x2x2x2x2xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : vector<2x2x2x2x2x2x2x2xf32>, i64, vector<2x2x2x2x2x2x2x2xindex>, vector<2x2x2x2x2x2x2x2xi1>
}
// -----
More information about the Mlir-commits
mailing list