[Mlir-commits] [mlir] [mlir][xegpu] Support N-D block transfers in VectorToXeGPU (PR #210527)
Jianhui Li
llvmlistbot at llvm.org
Sat Jul 18 14:55:25 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/210527
>From 0d7ff5b87542e7ecee78d6d33e6696b9338212b9 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 18 Jul 2026 17:41:56 +0000
Subject: [PATCH 1/2] [mlir][xegpu] Support N-D block transfers in
VectorToXeGPU
Extend the vector.transfer_read/transfer_write lowerings so they can
produce N-D xegpu.load_nd/store_nd, not just 1D/2D, and relax the
out-of-bounds handling to match load_nd's implicit-zero padding.
Changes:
- Allow high-dim (rank > 2) vectors on the nd block path. The rank
restriction in storeLoadPreconditions is gated behind a new
allowHighDimVector flag so plain vector.load/store keep their 1D/2D
limit while transfers accept any non-0D rank.
- Accept an out-of-bounds read whose padding is a zero constant or a
ub.poison value (isZeroOrPoisonPadding): load_nd zero-fills
out-of-bounds elements, and poison means those elements are
don't-care, so both match its semantics.
- Restructure both patterns to attempt the nd block op first via a
single canLowerToLoadNd/canLowerToStoreNd predicate, then fall back
to the scattered path. The block path only realizes a minor-identity
map (reads also handle an innermost-two-dims transpose via a trailing
vector.transpose); any other permutation is routed to the scattered
path, which permutes strides explicitly. This fixes silently dropping
the permutation on high-dim transposed transfers.
Tests updated for the new high-dim block lowering, plus new
load_high_dim_transposed / store_high_dim_transposed coverage.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 268 ++++++++++--------
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 169 +++++++----
.../transfer-write-to-xegpu.mlir | 225 +++++++--------
3 files changed, 380 insertions(+), 282 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9a994e87697f6..306ce6f1e8008 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -49,14 +50,53 @@ static bool isZeroConstant(Value val) {
.Default(false);
}
+// Return true if the transfer padding value is compatible with the implicit
+// padding of an nd block load. LoadNdOp fills out-of-bounds elements with zero,
+// so a zero constant matches its semantics exactly. A poison padding means the
+// out-of-bounds elements are "don't care", so any implicit padding (including
+// zero) is also acceptable.
+static bool isZeroOrPoisonPadding(Value val) {
+ return isZeroConstant(val) || val.getDefiningOp<ub::PoisonOp>();
+}
+
+// Return true if the permutation map keeps every dimension in place except the
+// innermost two, which are swapped, e.g.:
+// (d0, d1) -> (d1, d0)
+// (d0, d1, d2) -> (d2, d1)
+// (d0, d1, d2, d3) -> (d0, d1, d3, d2)
+// This is the only non-identity permutation an nd block load can realize (by
+// loading the untransposed block and applying a trailing vector.transpose).
+static bool isInnermostTwoDimsTransposed(AffineMap map) {
+ unsigned numResults = map.getNumResults();
+ if (numResults < 2)
+ return false;
+ MLIRContext *ctx = map.getContext();
+ unsigned numInputs = map.getNumInputs();
+ // All but the innermost two results must match the minor-identity map.
+ for (unsigned i = 0; i + 2 < numResults; ++i)
+ if (map.getResult(i) != getAffineDimExpr(numInputs - numResults + i, ctx))
+ return false;
+ // The innermost two results must be the last two input dims, swapped.
+ return map.getResult(numResults - 2) == getAffineDimExpr(numInputs - 1, ctx) &&
+ map.getResult(numResults - 1) == getAffineDimExpr(numInputs - 2, ctx);
+}
+
static LogicalResult storeLoadPreconditions(PatternRewriter &rewriter,
Operation *op, VectorType vecTy,
- MemRefType memTy) {
+ MemRefType memTy,
+ bool allowHighDimVector = false) {
// Validate only vector as the basic vector store and load ops guarantee
// XeGPU-compatible memref source.
unsigned vecRank = vecTy.getRank();
- if (!(vecRank == 1 || vecRank == 2))
+ // nd block loads/stores support N-D block transfers, so callers that lower to
+ // them (transfer_read/transfer_write) allow any non-zero rank. The plain
+ // vector.load/store lowering keeps the 1D/2D restriction.
+ if (allowHighDimVector) {
+ if (vecRank == 0)
+ return rewriter.notifyMatchFailure(op, "Expects non-0D vector");
+ } else if (!(vecRank == 1 || vecRank == 2)) {
return rewriter.notifyMatchFailure(op, "Expects 1D or 2D vector");
+ }
if (!vecTy.getElementType().isIntOrFloat())
return rewriter.notifyMatchFailure(
@@ -594,92 +634,78 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
return success();
}
- // TODO:This check needs to be replaced with proper uArch capability check
+ // TODO: This check needs to be replaced with proper uArch capability check.
auto chip = xegpu::getChipStr(readOp);
- // Lower to scattered load Op if the target HW doesn't have 2d block load
- // support and the load is not from shared memory.
- if ((chip != "pvc" && chip != "bmg" && chip != "cri") ||
- readOp.getVectorType().getRank() > 2) {
-
- // TODO: add support for OutOfBound access
- if (isOutOfBounds)
- return failure();
- return lowerToScatteredLoadOp(readOp, rewriter);
- }
-
- // Handle the 1D non-SLM case using load.gather.
- if (loadedVecTy.getRank() == 1 && !isOutOfBounds)
- return lowerToScatteredLoadOp(readOp, rewriter);
-
- // Perform common data transfer checks.
- // TODO: Maybe too strict for SLM case.
- if (failed(
- storeLoadPreconditions(rewriter, readOp, loadedVecTy, readMemTy)))
- return failure();
-
- if (isOutOfBounds && !isZeroConstant(readOp.getPadding()))
- return rewriter.notifyMatchFailure(
- readOp, "Unsupported non-zero padded out-of-bounds read");
+ bool hasBlockLoadSupport =
+ (chip == "pvc" || chip == "bmg" || chip == "cri");
+ // An nd block load can realize a minor-identity map directly, or an
+ // innermost-two-dims transpose via a trailing vector.transpose. Any other
+ // permutation (e.g. a mid-vector transpose of a high-dim load) is left to
+ // the scattered path, which permutes strides explicitly.
AffineMap readMap = readOp.getPermutationMap();
- // Check if this is a transpose: the map must have exactly 2 results,
- // and those 2 results must be the last 2 input dimensions interchanged.
- // Examples:
- // (d0, d1) -> (d1, d0) // transpose
- // (d0, d1) -> (d0, d1) // not a transpose
- // (d0, d1, d2) -> (d2, d1) // transpose (last 2 dims swapped)
- bool isTransposeLoad = false;
- if (readMap.getNumResults() == 2) {
- auto results = readMap.getResults();
- unsigned numInputs = readMap.getNumInputs();
- if (numInputs >= 2) {
- auto lastDim = getAffineDimExpr(numInputs - 1, readMap.getContext());
- auto secondLastDim =
- getAffineDimExpr(numInputs - 2, readMap.getContext());
- isTransposeLoad =
- (results[0] == lastDim && results[1] == secondLastDim);
+ bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
+
+ // Prefer an nd block load. It requires HW block-load support, a >1D vector
+ // of a scalar element type backed by a scalar-element memref, and a map the
+ // block load can realize. Out-of-bounds reads are allowed as long as the
+ // padding matches load_nd's implicit zero padding.
+ bool canLowerToLoadNd =
+ hasBlockLoadSupport && loadedVecTy.getRank() > 1 &&
+ (readMap.isMinorIdentity() || isTransposeLoad) &&
+ loadedVecTy.getElementType().isIntOrFloat() &&
+ readMemTy.getElementType().isIntOrFloat() &&
+ (!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
+
+ 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]);
+ loadedVecTy = VectorType::get(descShape, elementType);
}
+ auto descType = xegpu::TensorDescType::get(
+ descShape, elementType, /*array_length=*/1,
+ /*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
+ auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
+ rewriter, loc, readOp.getBase(),
+ getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
+ // By default, no specific caching policy is assigned.
+ xegpu::CachePolicyAttr hint = nullptr;
+ xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
+
+ Operation *loadedOp =
+ xegpu::LoadNdOp::create(rewriter, loc, loadedVecTy, ndDesc, indices,
+ /*packed=*/nullptr, /*transpose=*/nullptr,
+ /*l1_hint=*/hint,
+ /*l2_hint=*/hint, /*l3_hint=*/hint,
+ /*layout=*/nullptr);
+ if (isTransposeLoad) {
+ // Undo the innermost-two-dims swap with a trailing vector.transpose:
+ // keep the leading dimensions in place and interchange only the last
+ // two.
+ int64_t rank = loadedVecTy.getRank();
+ SmallVector<int64_t> perm(llvm::to_vector(llvm::seq<int64_t>(0, rank)));
+ std::swap(perm[rank - 1], perm[rank - 2]);
+ loadedOp = vector::TransposeOp::create(rewriter, loc,
+ loadedOp->getResult(0), perm);
+ }
+ rewriter.replaceOp(readOp, loadedOp);
+ return success();
}
- 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]);
- loadedVecTy = VectorType::get(descShape, elementType);
- }
- auto descType = xegpu::TensorDescType::get(
- descShape, elementType, /*array_length=*/1,
- /*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
- auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
- rewriter, loc, readOp.getBase(), getAsOpFoldResult(readOp.getIndices()),
- loadedVecTy.getRank());
- // By default, no specific caching policy is assigned.
- xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
- rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
-
- Operation *loadedOp =
- xegpu::LoadNdOp::create(rewriter, loc, loadedVecTy, ndDesc, indices,
- /*packed=*/nullptr, /*transpose=*/nullptr,
- /*l1_hint=*/hint,
- /*l2_hint=*/hint, /*l3_hint=*/hint,
- /*layout=*/nullptr);
- if (isTransposeLoad) {
- // Transposing the loaded vector with a separate vector.transpose
- // operation
- auto range = llvm::seq<int64_t>(0, readMap.getResults().size());
- SmallVector<int64_t> perm(
- range.rbegin(), range.rend()); // reverse the range for transpose
- loadedOp = vector::TransposeOp::create(rewriter, loc,
- loadedOp->getResult(0), perm);
- }
- rewriter.replaceOp(readOp, loadedOp);
- return success();
+ // Fall back to a scattered load. It supports arbitrary permutations and any
+ // rank, but cannot express out-of-bounds accesses.
+ // TODO: add support for OutOfBound access.
+ if (isOutOfBounds)
+ return failure();
+ return lowerToScatteredLoadOp(readOp, rewriter);
}
};
@@ -727,47 +753,51 @@ struct TransferWriteLowering
return success();
}
- // TODO:This check needs to be replaced with proper uArch capability check
+ // TODO: This check needs to be replaced with proper uArch capability check.
auto chip = xegpu::getChipStr(writeOp);
- // Lower to scattered store Op if the target HW doesn't have 2d block
- // store support and the memref is not SLM.
- if ((chip != "pvc" && chip != "bmg" && chip != "cri") ||
- writeOp.getVectorType().getRank() > 2) {
-
- // TODO: add support for OutOfBound access
- if (writeOp.hasOutOfBoundsDim())
- return failure();
- return lowerToScatteredStoreOp(writeOp, rewriter);
- }
-
- if (failed(storeLoadPreconditions(rewriter, writeOp, vecTy, writeMemTy)))
- return failure();
+ bool hasBlockStoreSupport =
+ (chip == "pvc" || chip == "bmg" || chip == "cri");
+ // Prefer an nd block store. It requires HW block-store support, a >1D
+ // vector of a scalar element type backed by a scalar-element memref, and a
+ // minor-identity map (block stores have no transpose support). Out-of-bounds
+ // writes are handled by the descriptor's boundary check.
AffineMap map = writeOp.getPermutationMap();
- if (!map.isMinorIdentity())
- return rewriter.notifyMatchFailure(writeOp, "Expects identity map");
-
- auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
- rewriter, loc, writeOp.getBase(),
- getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
-
- auto descType = xegpu::TensorDescType::get(
- vecTy.getShape(), vecTy.getElementType(),
- /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
- xegpu::MemorySpace::Global);
- // By default, no specific caching policy is assigned.
- xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
- rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
-
- auto storeOp = xegpu::StoreNdOp::create(rewriter, loc, writeOp.getVector(),
- ndDesc, indices,
- /*l1_hint=*/hint,
- /*l2_hint=*/hint, /*l3_hint=*/hint,
- /*layout=*/nullptr);
- rewriter.replaceOp(writeOp, storeOp);
+ bool canLowerToStoreNd =
+ hasBlockStoreSupport && vecTy.getRank() > 1 && map.isMinorIdentity() &&
+ vecTy.getElementType().isIntOrFloat() &&
+ writeMemTy.getElementType().isIntOrFloat();
+
+ if (canLowerToStoreNd) {
+ auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
+ rewriter, loc, writeOp.getBase(),
+ getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
+
+ auto descType = xegpu::TensorDescType::get(
+ vecTy.getShape(), vecTy.getElementType(),
+ /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
+ xegpu::MemorySpace::Global);
+ // By default, no specific caching policy is assigned.
+ xegpu::CachePolicyAttr hint = nullptr;
+ xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
+
+ auto storeOp =
+ xegpu::StoreNdOp::create(rewriter, loc, writeOp.getVector(), ndDesc,
+ indices,
+ /*l1_hint=*/hint,
+ /*l2_hint=*/hint, /*l3_hint=*/hint,
+ /*layout=*/nullptr);
+ rewriter.replaceOp(writeOp, storeOp);
+ return success();
+ }
- return success();
+ // Fall back to a scattered store. It supports arbitrary permutations and any
+ // rank, but cannot express out-of-bounds accesses.
+ // TODO: add support for OutOfBound access.
+ if (writeOp.hasOutOfBoundsDim())
+ return failure();
+ return lowerToScatteredStoreOp(writeOp, rewriter);
}
};
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 3354868cd6cb9..ed495729f9add 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -182,6 +182,33 @@ gpu.func @load_transpose_3d_memref(%source: memref<32x64x128xf32>,
}
+// -----
+// A high-dim load whose innermost two dims are transposed lowers to an
+// (untransposed) nd block load followed by a vector.transpose of the last two
+// dims.
+gpu.module @xevm_module {
+gpu.func @load_high_dim_transposed(%source: memref<2x2x64x128xf16>,
+ %offset: index) -> vector<1x1x64x128xf16> {
+ %c0 = arith.constant 0.0 : f16
+ %0 = vector.transfer_read %source[%offset, %offset, %offset, %offset], %c0
+ {permutation_map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3, d2)>,
+ in_bounds = [true, true, true, true]}
+ : memref<2x2x64x128xf16>, vector<1x1x64x128xf16>
+ gpu.return %0 : vector<1x1x64x128xf16>
+}
+
+// LOAD-ND-LABEL: @load_high_dim_transposed(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<2x2x64x128xf16>,
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<2x2x64x128xf16>
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x1x128x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
+// LOAD-ND-SAME: -> vector<1x1x128x64xf16>
+// LOAD-ND: vector.transpose %[[VEC]], [0, 1, 3, 2] : vector<1x1x128x64xf16> to vector<1x1x64x128xf16>
+
+// LOAD-GATHER-LABEL: @load_high_dim_transposed(
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} : i64, vector<1x1x64x128xindex>, vector<1x1x64x128xi1> -> vector<1x1x64x128xf16>
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
@@ -275,21 +302,31 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
gpu.return %0 : vector<2x4x8x16xf32>
}
-// CHECK-LABEL: @load_dynamic_source3(
-// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<2x4x8x16xi1>
-// CHECK: memref.extract_strided_metadata %[[SRC]] : memref<?x?x?x?x?xf32> -> memref<f32>, index, index, index, index, index, index, index, index, index, index, index
-// CHECK-COUNT4: vector.step
-// CHECK-COUNT3: vector.broadcast
-// CHECK-COUNT4: vector.shape_cast
-// CHECK-COUNT4: vector.broadcast {{.*}} : vector<2x4x8x16xindex>
-// CHECK-COUNT3: arith.addi {{.*}} : vector<2x4x8x16xindex>
-// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x4x8x16xindex>
-// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x4x8x16xindex>
-// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?x?x?xf32> -> index
-// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<2x4x8x16xindex>, vector<2x4x8x16xi1> -> vector<2x4x8x16xf32>
-// CHECK: return %[[VEC]]
+// LOAD-ND-LABEL: @load_dynamic_source3(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
+// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
+// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:4, %[[STRIDES:.*]]:4 = memref.extract_strided_metadata %[[SUBVIEW]]
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<2x4x8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
+// LOAD-ND-SAME: -> vector<2x4x8x16xf32>
+// LOAD-ND: return %[[VEC]]
+
+// LOAD-GATHER-LABEL: @load_dynamic_source3(
+// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
+// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<2x4x8x16xi1>
+// LOAD-GATHER: memref.extract_strided_metadata %[[SRC]] : memref<?x?x?x?x?xf32> -> memref<f32>, index, index, index, index, index, index, index, index, index, index, index
+// LOAD-GATHER-COUNT4: vector.step
+// LOAD-GATHER-COUNT3: vector.broadcast
+// LOAD-GATHER-COUNT4: vector.shape_cast
+// LOAD-GATHER-COUNT4: vector.broadcast {{.*}} : vector<2x4x8x16xindex>
+// LOAD-GATHER-COUNT3: arith.addi {{.*}} : vector<2x4x8x16xindex>
+// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<2x4x8x16xindex>
+// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], {{.*}} : vector<2x4x8x16xindex>
+// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?x?x?xf32> -> index
+// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<2x4x8x16xindex>, vector<2x4x8x16xi1> -> vector<2x4x8x16xf32>
+// LOAD-GATHER: return %[[VEC]]
}
// -----
@@ -302,21 +339,27 @@ gpu.func @load_high_dim_vector(%source: memref<16x32x64xf32>,
gpu.return %0 : vector<8x16x32xf32>
}
-// CHECK-LABEL: @load_high_dim_vector(
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8x16x32xi1>
-// CHECK: %[[CST_0:.+]] = arith.constant dense<64> : vector<16xindex>
-// CHECK: %[[CST_1:.+]] = arith.constant dense<2048> : vector<8xindex>
-// CHECK: %[[C2048:.+]] = arith.constant 2048 : index
-// CHECK: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-COUNT3: vector.step
-// CHECK-COUNT3: vector.shape_cast
-// CHECK-COUNT3: vector.broadcast {{.*}} : vector<8x16x32xindex>
-// CHECK-COUNT2: arith.addi {{.*}} : vector<8x16x32xindex>
-// CHECK: %[[BCASTOFF:.+]] = vector.broadcast {{.*}} : index to vector<8x16x32xindex>
-// CHECK: %[[IDX:.+]] = arith.addi %[[BCASTOFF]], {{.*}} : vector<8x16x32xindex>
-// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %arg0 : memref<16x32x64xf32> -> index
-// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16x32xindex>, vector<8x16x32xi1> -> vector<8x16x32xf32>
+// LOAD-ND-LABEL: @load_high_dim_vector(
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %arg0 : memref<16x32x64xf32>
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<8x16x32xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
+// LOAD-ND-SAME: -> vector<8x16x32xf32>
+
+// LOAD-GATHER-LABEL: @load_high_dim_vector(
+// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16x32xi1>
+// LOAD-GATHER: %[[CST_0:.+]] = arith.constant dense<64> : vector<16xindex>
+// LOAD-GATHER: %[[CST_1:.+]] = arith.constant dense<2048> : vector<8xindex>
+// LOAD-GATHER: %[[C2048:.+]] = arith.constant 2048 : index
+// LOAD-GATHER: %[[C64:.+]] = arith.constant 64 : index
+// LOAD-GATHER-COUNT3: vector.step
+// LOAD-GATHER-COUNT3: vector.shape_cast
+// LOAD-GATHER-COUNT3: vector.broadcast {{.*}} : vector<8x16x32xindex>
+// LOAD-GATHER-COUNT2: arith.addi {{.*}} : vector<8x16x32xindex>
+// LOAD-GATHER: %[[BCASTOFF:.+]] = vector.broadcast {{.*}} : index to vector<8x16x32xindex>
+// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[BCASTOFF]], {{.*}} : vector<8x16x32xindex>
+// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %arg0 : memref<16x32x64xf32> -> index
+// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16x32xindex>, vector<8x16x32xi1> -> vector<8x16x32xf32>
}
@@ -330,18 +373,25 @@ gpu.func @load_8D_vector(%source: memref<2x2x2x2x2x2x2x2xf32>,
gpu.return %0 : vector<2x2x2x2x2x2x2x2xf32>
}
-// CHECK-LABEL: @load_8D_vector(
-// CHECK-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>,
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<2x2x2x2x2x2x2x2xi1>
-// CHECK-COUNT8: vector.step
-// CHECK-COUNT7: vector.shape_cast
-// CHECK-COUNT8: vector.broadcast {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// CHECK-COUNT7: 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>
+// 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>
}
@@ -728,23 +778,38 @@ gpu.func @transpose_1x1024x24x64(
gpu.return
}
-// 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
+// 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
// 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).
-// CHECK: arith.muli %{{.+}}, %[[C1536]] : index
-// CHECK: arith.muli %block_id_x, %[[C64]] : index
-// CHECK: xegpu.load {{.*}} -> vector<1x1x16x8xf16>
+// LOAD-GATHER: arith.muli %{{.+}}, %[[C1536]] : index
+// LOAD-GATHER: arith.muli %block_id_x, %[[C64]] : index
+// LOAD-GATHER: 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).
-// CHECK: arith.muli %block_id_x, %[[C65536]] : index
-// CHECK: arith.muli %{{.+}}, %[[C64]] : index
-// CHECK: xegpu.store {{.*}}
+// LOAD-GATHER: arith.muli %block_id_x, %[[C65536]] : index
+// LOAD-GATHER: arith.muli %{{.+}}, %[[C64]] : index
+// LOAD-GATHER: 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 45bda46f738b2..ad7332097cc8f 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -11,36 +11,20 @@ gpu.func @store_1D_vector(%vec: vector<8xf32>,
gpu.return
}
-// STORE-ND-LABEL: @store_1D_vector(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]
-// STORE-ND: %[[BASE_BUFFER:.+]], %[[OFFSET1:.+]], %[[SIZES:.+]], %[[STRIDES:.+]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// STORE-ND-SAME: : memref<32xf32, strided<[1], offset: ?>> -> memref<f32>, index, index, index
-// STORE-ND: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// STORE-ND-SAME: : memref<f32> -> index
-// STORE-ND: %[[MUL:.+]] = arith.muli %[[OFFSET1]], %[[ELEM_BYTES]] : index
-// STORE-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// STORE-ND: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [32],
-// STORE-ND-SAME: strides : [1] : i64 -> !xegpu.tensor_desc<8xf32,
-// STORE-ND-SAME: boundary_check = false
-// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]]] : vector<8xf32>
-
-// STORE-SCATTER-LABEL: @store_1D_vector(
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf32>,
-// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// STORE-SCATTER-DAG: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// STORE-SCATTER-DAG: %[[STEP:.+]] = vector.step
-// STORE-SCATTER-COUNT2: arith.muli {{.*}} : index
-// STORE-SCATTER-COUNT2: arith.addi {{.*}} : index
-// STORE-SCATTER-DAG: %[[BCAST:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
-// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST]], %{{.*}} : vector<8xindex>
-// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
-// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>
+// A 1D store cannot use an nd block store, so it lowers through the scattered
+// path identically on both the target and non-target runs.
+// CHECK-LABEL: @store_1D_vector(
+// CHECK-SAME: %[[VEC:.+]]: vector<8xf32>,
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-DAG: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// CHECK-DAG: %[[STEP:.+]] = vector.step
+// CHECK-COUNT2: arith.muli {{.*}} : index
+// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK-DAG: %[[BCAST:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
+// CHECK-DAG: %[[IDX:.+]] = arith.addi %[[BCAST]], %{{.*}} : vector<8xindex>
+// CHECK-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
+// CHECK-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>
}
// -----
@@ -150,7 +134,7 @@ gpu.func @store_out_of_bounds(%vec: vector<8x16xf32>,
// -----
gpu.module @xevm_module {
-gpu.func @no_store_transposed(%vec: vector<8x16xf32>,
+gpu.func @store_transposed(%vec: vector<8x16xf32>,
%source: memref<32x64xf32>, %offset: index) {
vector.transfer_write %vec, %source[%offset, %offset]
{permutation_map = affine_map<(d0, d1) -> (d1, d0)>,
@@ -159,22 +143,38 @@ gpu.func @no_store_transposed(%vec: vector<8x16xf32>,
gpu.return
}
-// STORE-ND-LABEL: @no_store_transposed(
-// STORE-ND: vector.transfer_write
+// An nd block store cannot transpose, so a transposed write lowers through the
+// scattered path identically on both the target and non-target runs.
+// CHECK-LABEL: @store_transposed(
+// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>,
+// CHECK-SAME: %[[SRC:.+]]: memref<32x64xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
+// CHECK-COUNT2: %[[STEP:.+]] = vector.step
+// CHECK-COUNT2: vector.shape_cast {{.*}}
+// CHECK-COUNT2: vector.broadcast {{.*}} : vector<8x16xindex>
+// CHECK-DAG: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>
+// CHECK-DAG: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}} : vector<8x16xindex>
+// CHECK-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32x64xf32> -> index
+// CHECK-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>
+}
-// STORE-SCATTER-LABEL: @no_store_transposed(
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16xf32>,
-// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<32x64xf32>,
-// STORE-SCATTER-SAME: %[[OFFSET:.+]]: index
-// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
-// STORE-SCATTER-COUNT2: %[[STEP:.+]] = vector.step
-// STORE-SCATTER-COUNT2: vector.shape_cast {{.*}}
-// STORE-SCATTER-COUNT2: vector.broadcast {{.*}} : vector<8x16xindex>
-// STORE-SCATTER-DAG: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>
-// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}} : vector<8x16xindex>
-// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32x64xf32> -> index
-// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>
+// -----
+// A high-dim transposed write cannot use an nd block store (no transpose
+// support), so it lowers through the scattered path on both runs.
+gpu.module @xevm_module {
+gpu.func @store_high_dim_transposed(%vec: vector<8x16x4xf32>,
+ %source: memref<16x32x64xf32>, %offset: index) {
+ vector.transfer_write %vec, %source[%offset, %offset, %offset]
+ {permutation_map = affine_map<(d0, d1, d2) -> (d0, d2, d1)>,
+ in_bounds = [true, true, true]}
+ : vector<8x16x4xf32>, memref<16x32x64xf32>
+ gpu.return
+}
+
+// CHECK-LABEL: @store_high_dim_transposed(
+// CHECK: xegpu.store {{.*}} : vector<8x16x4xf32>, i64, vector<8x16x4xindex>, vector<8x16x4xi1>
}
// -----
@@ -187,23 +187,31 @@ gpu.func @store_high_dim_vector(%vec: vector<8x16x32xf32>,
gpu.return
}
-// CHECK-LABEL: @store_high_dim_vector(
-// CHECK-SAME: %[[VEC:.+]]: vector<8x16x32xf32>,
-// CHECK-SAME: %[[SRC:.+]]: memref<16x32x64xf32>
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8x16x32xi1>
-// CHECK: %[[CST_0:.+]] = arith.constant dense<64> : vector<16xindex>
-// CHECK: %[[CST_1:.+]] = arith.constant dense<2048> : vector<8xindex>
-// CHECK: %[[C2048:.+]] = arith.constant 2048 : index
-// CHECK: %[[C64:.+]] = arith.constant 64 : index
-// CHECK-COUNT3: vector.step
-// CHECK-COUNT3: vector.shape_cast
-// CHECK-COUNT3: vector.broadcast {{.*}} : vector<8x16x32xindex>
-// CHECK-COUNT2: arith.addi {{.*}} : vector<8x16x32xindex>
-// CHECK: %[[BCASTOFF:.+]] = vector.broadcast {{.*}} : index to vector<8x16x32xindex>
-// CHECK: %[[IDX:.+]] = arith.addi %[[BCASTOFF]], {{.*}} : vector<8x16x32xindex>
-// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<16x32x64xf32> -> index
-// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : vector<8x16x32xf32>, i64, vector<8x16x32xindex>, vector<8x16x32xi1>
+// STORE-ND-LABEL: @store_high_dim_vector(
+// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16x32xf32>,
+// STORE-ND-SAME: %[[SRC:.+]]: memref<16x32x64xf32>
+// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<16x32x64xf32>
+// STORE-ND-SAME: -> !xegpu.tensor_desc<8x16x32xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]]
+// STORE-ND-SAME: : vector<8x16x32xf32>, !xegpu.tensor_desc<8x16x32xf32
+
+// STORE-SCATTER-LABEL: @store_high_dim_vector(
+// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16x32xf32>,
+// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<16x32x64xf32>
+// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16x32xi1>
+// STORE-SCATTER: %[[CST_0:.+]] = arith.constant dense<64> : vector<16xindex>
+// STORE-SCATTER: %[[CST_1:.+]] = arith.constant dense<2048> : vector<8xindex>
+// STORE-SCATTER: %[[C2048:.+]] = arith.constant 2048 : index
+// STORE-SCATTER: %[[C64:.+]] = arith.constant 64 : index
+// STORE-SCATTER-COUNT3: vector.step
+// STORE-SCATTER-COUNT3: vector.shape_cast
+// STORE-SCATTER-COUNT3: vector.broadcast {{.*}} : vector<8x16x32xindex>
+// STORE-SCATTER-COUNT2: arith.addi {{.*}} : vector<8x16x32xindex>
+// STORE-SCATTER: %[[BCASTOFF:.+]] = vector.broadcast {{.*}} : index to vector<8x16x32xindex>
+// STORE-SCATTER: %[[IDX:.+]] = arith.addi %[[BCASTOFF]], {{.*}} : vector<8x16x32xindex>
+// STORE-SCATTER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<16x32x64xf32> -> index
+// STORE-SCATTER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : vector<8x16x32xf32>, i64, vector<8x16x32xindex>, vector<8x16x32xi1>
}
// -----
@@ -216,19 +224,27 @@ gpu.func @store_8D_vector(%vec: vector<2x2x2x2x2x2x2x2xf32>,
gpu.return
}
-// CHECK-LABEL: @store_8D_vector(
-// CHECK-SAME: %[[VEC:.+]]: vector<2x2x2x2x2x2x2x2xf32>,
-// CHECK-SAME: %[[SRC:.+]]: memref<2x2x2x2x2x2x2x2xf32>
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<2x2x2x2x2x2x2x2xi1>
-// CHECK-COUNT8: vector.step
-// CHECK-COUNT7: vector.shape_cast
-// CHECK-COUNT8: vector.broadcast {{.*}} : vector<2x2x2x2x2x2x2x2xindex>
-// CHECK-COUNT7: 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>
+// 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>
}
// -----
@@ -330,41 +346,28 @@ gpu.func @store_to_subview(%vec: vector<8xf16>,
: vector<8xf16>, memref<256x256xf16, strided<[4096, 1], offset: ?>>
gpu.return
}
-// STORE-ND-LABEL: @store_to_subview(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf16>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
-// STORE-ND-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// STORE-ND: %[[ELEM_BYTES:.+]] = arith.constant 2 : index
-// STORE-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SUBVIEW]][%[[OFF2]], 0]
-// STORE-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.*]], %[[STRIDES:.*]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// STORE-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// STORE-ND: %[[MUL:.+]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// STORE-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// STORE-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// STORE-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [256], strides : [1] : i64 ->
-// STORE-ND-SAME: !xegpu.tensor_desc<8xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
-// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF2]]] : vector<8xf16>
-
-// STORE-SCATTER-LABEL: @store_to_subview(
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf16>,
-// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
-// STORE-SCATTER-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// STORE-SCATTER: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]
-// STORE-SCATTER-SAME: : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// STORE-SCATTER: %[[BB:.+]], %[[OFFSET:.+]], {{.*}}, {{.*}} = memref.extract_strided_metadata %[[SUBVIEW]]
-// STORE-SCATTER-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
-// STORE-SCATTER: %[[STEP:.+]] = vector.step : vector<8xindex>
-// STORE-SCATTER: arith.muli {{.*}} : index
-// STORE-SCATTER: arith.addi %[[OFFSET]]{{.*}} : index
-// STORE-SCATTER: arith.addi {{.*}} : index
-// STORE-SCATTER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
-// STORE-SCATTER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
-// STORE-SCATTER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]]
-// STORE-SCATTER-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
-// STORE-SCATTER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>
+// A 1D store cannot use an nd block store, so it lowers through the scattered
+// path identically on both runs, even when the destination is a strided
+// subview.
+// CHECK-LABEL: @store_to_subview(
+// CHECK-SAME: %[[VEC:.+]]: vector<8xf16>,
+// CHECK-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
+// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]
+// CHECK-SAME: : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
+// CHECK: %[[BB:.+]], %[[OFFSET:.+]], {{.*}}, {{.*}} = memref.extract_strided_metadata %[[SUBVIEW]]
+// CHECK-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
+// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi %[[OFFSET]]{{.*}} : index
+// CHECK: arith.addi {{.*}} : index
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]]
+// CHECK-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>
}
// -----
>From 0c33ce78ecc38f5ab40ebc922a0ddc5306369df5 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 18 Jul 2026 21:54:01 +0000
Subject: [PATCH 2/2] [mlir][xegpu] Refactor transfer lowering; inline
storeLoadPreconditions
Restructure TransferReadLowering/TransferWriteLowering to try the nd block
op first via a single canLowerTo{Load,Store}Nd predicate, then fall back to
the scattered path (guarding out-of-bounds, which scatter cannot express).
Remove storeLoadPreconditions and inline its remaining checks into the plain
vector.load/store patterns. Drop the redundant vector element-type check since
it always matches the memref element type.
No functional change.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 86 +++++++------------
1 file changed, 30 insertions(+), 56 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 306ce6f1e8008..80bd88beece6d 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -77,42 +77,11 @@ static bool isInnermostTwoDimsTransposed(AffineMap map) {
if (map.getResult(i) != getAffineDimExpr(numInputs - numResults + i, ctx))
return false;
// The innermost two results must be the last two input dims, swapped.
- return map.getResult(numResults - 2) == getAffineDimExpr(numInputs - 1, ctx) &&
+ return map.getResult(numResults - 2) ==
+ getAffineDimExpr(numInputs - 1, ctx) &&
map.getResult(numResults - 1) == getAffineDimExpr(numInputs - 2, ctx);
}
-static LogicalResult storeLoadPreconditions(PatternRewriter &rewriter,
- Operation *op, VectorType vecTy,
- MemRefType memTy,
- bool allowHighDimVector = false) {
- // Validate only vector as the basic vector store and load ops guarantee
- // XeGPU-compatible memref source.
- unsigned vecRank = vecTy.getRank();
- // nd block loads/stores support N-D block transfers, so callers that lower to
- // them (transfer_read/transfer_write) allow any non-zero rank. The plain
- // vector.load/store lowering keeps the 1D/2D restriction.
- if (allowHighDimVector) {
- if (vecRank == 0)
- return rewriter.notifyMatchFailure(op, "Expects non-0D vector");
- } else if (!(vecRank == 1 || vecRank == 2)) {
- return rewriter.notifyMatchFailure(op, "Expects 1D or 2D vector");
- }
-
- if (!vecTy.getElementType().isIntOrFloat())
- return rewriter.notifyMatchFailure(
- op, "Expected scalar type with known bitwidth");
-
- // XeGPU requires the memref to have a scalar integer or float element type.
- // Memrefs with vector element types (e.g. memref<?xvector<4xf32>>) are not
- // supported because createNdDescriptor computes byte offsets using
- // getElementTypeBitWidth(), which asserts on non-integer/float types.
- if (!memTy.getElementType().isIntOrFloat())
- return rewriter.notifyMatchFailure(
- op, "Unsupported memref element type: expected integer or float");
-
- return success();
-}
-
static LogicalResult transferPreconditions(PatternRewriter &rewriter,
VectorTransferOpInterface xferOp) {
if (xferOp.getMask())
@@ -647,13 +616,12 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
// Prefer an nd block load. It requires HW block-load support, a >1D vector
- // of a scalar element type backed by a scalar-element memref, and a map the
- // block load can realize. Out-of-bounds reads are allowed as long as the
- // padding matches load_nd's implicit zero padding.
+ // backed by a scalar-element memref, and a map the block load can realize.
+ // Out-of-bounds reads are allowed as long as the padding matches load_nd's
+ // implicit zero padding.
bool canLowerToLoadNd =
hasBlockLoadSupport && loadedVecTy.getRank() > 1 &&
(readMap.isMinorIdentity() || isTransposeLoad) &&
- loadedVecTy.getElementType().isIntOrFloat() &&
readMemTy.getElementType().isIntOrFloat() &&
(!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
@@ -759,14 +727,13 @@ struct TransferWriteLowering
(chip == "pvc" || chip == "bmg" || chip == "cri");
// Prefer an nd block store. It requires HW block-store support, a >1D
- // vector of a scalar element type backed by a scalar-element memref, and a
- // minor-identity map (block stores have no transpose support). Out-of-bounds
- // writes are handled by the descriptor's boundary check.
+ // vector backed by a scalar-element memref, and a minor-identity map (block
+ // stores have no transpose support). Out-of-bounds writes are handled by
+ // the descriptor's boundary check.
AffineMap map = writeOp.getPermutationMap();
- bool canLowerToStoreNd =
- hasBlockStoreSupport && vecTy.getRank() > 1 && map.isMinorIdentity() &&
- vecTy.getElementType().isIntOrFloat() &&
- writeMemTy.getElementType().isIntOrFloat();
+ bool canLowerToStoreNd = hasBlockStoreSupport && vecTy.getRank() > 1 &&
+ map.isMinorIdentity() &&
+ writeMemTy.getElementType().isIntOrFloat();
if (canLowerToStoreNd) {
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
@@ -782,18 +749,17 @@ struct TransferWriteLowering
xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
- auto storeOp =
- xegpu::StoreNdOp::create(rewriter, loc, writeOp.getVector(), ndDesc,
- indices,
- /*l1_hint=*/hint,
- /*l2_hint=*/hint, /*l3_hint=*/hint,
- /*layout=*/nullptr);
+ auto storeOp = xegpu::StoreNdOp::create(
+ rewriter, loc, writeOp.getVector(), ndDesc, indices,
+ /*l1_hint=*/hint,
+ /*l2_hint=*/hint, /*l3_hint=*/hint,
+ /*layout=*/nullptr);
rewriter.replaceOp(writeOp, storeOp);
return success();
}
- // Fall back to a scattered store. It supports arbitrary permutations and any
- // rank, but cannot express out-of-bounds accesses.
+ // Fall back to a scattered store. It supports arbitrary permutations and
+ // any rank, but cannot express out-of-bounds accesses.
// TODO: add support for OutOfBound access.
if (writeOp.hasOutOfBoundsDim())
return failure();
@@ -878,8 +844,12 @@ struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
VectorType vecTy = loadOp.getResult().getType();
MemRefType memTy = loadOp.getBase().getType();
- if (failed(storeLoadPreconditions(rewriter, loadOp, vecTy, memTy)))
- return failure();
+ // The plain vector.load lowering only supports 1D/2D block loads.
+ if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
+ return rewriter.notifyMatchFailure(loadOp, "Expects 1D or 2D vector");
+ if (!memTy.getElementType().isIntOrFloat())
+ return rewriter.notifyMatchFailure(
+ loadOp, "Unsupported memref element type: expected integer or float");
// Boundary check is available only for block instructions.
bool boundaryCheck = vecTy.getRank() > 1;
@@ -918,8 +888,12 @@ struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
TypedValue<VectorType> vector = storeOp.getValueToStore();
VectorType vecTy = vector.getType();
MemRefType memTy = storeOp.getBase().getType();
- if (failed(storeLoadPreconditions(rewriter, storeOp, vecTy, memTy)))
- return failure();
+ // The plain vector.store lowering only supports 1D/2D block stores.
+ if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
+ return rewriter.notifyMatchFailure(storeOp, "Expects 1D or 2D vector");
+ if (!memTy.getElementType().isIntOrFloat())
+ return rewriter.notifyMatchFailure(
+ storeOp, "Unsupported memref element type: expected integer or float");
// Boundary check is available only for block instructions.
bool boundaryCheck = vecTy.getRank() > 1;
More information about the Mlir-commits
mailing list