[Mlir-commits] [mlir] [mlir][VectorToXeGPU] Lower out-of-bounds 1-D transfers (PR #220490)
Jianhui Li
llvmlistbot at llvm.org
Tue Sep 1 22:52:07 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/220490
A vector.transfer_read / vector.transfer_write whose vector is 1-D and whose access is out of bounds was not lowered.
This PR now builds a real in-bounds mask instead of handing xegpu.load/xegpu.store an all-ones one. An element is in bounds when its index is still inside the memref dimension it reads, i.e. i < dim(d) - indices[d], and dimensions the transfer already declares in-bounds are left alone. Since masked-off lanes of an xegpu.load hold unspecified values rather than zeros, a read adds an arith.select to fill them with the transfer's padding.
A 1-D transfer of a single element skips that path and becomes a scalar memref.load, wrapped in an scf.if that yields the padding when the index is out of bounds. This avoid special handling the layout propagation for single element vector.
>From ade4e239e2621dd30ce0bd01ce329def970c43aa Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 23 Aug 2026 23:31:29 +0000
Subject: [PATCH 1/3] [mlir][VectorToXeGPU] Lower a 1-D out-of-bounds transfer
via load_nd/store_nd
A `vector.transfer_read` / `transfer_write` whose vector is 1-D and whose
access is out of bounds was not lowered at all, even on a target with 2-D
block load/store support: `transferPreconditions` rejected it outright with
"Boundary check is available only for block instructions." The op then
survived to `xegpu-wg-to-sg-distribute`, which has no transfer pattern and
marks unknown ops legal, so it stayed at the workgroup shape while its
consumers were distributed and the module failed to verify.
The premise of the bail is right - only a 2-D block instruction has a
boundary check. The scattered `xegpu.load`/`store` path has no notion of
bounds, and a rank-1 `xegpu.load_nd` takes the 1-D branch of
`LoadStorePrefetchNdToXeVMPattern`, which emits `xevm.blockload` with no
boundary check. But rejecting is unnecessary: a 1-D transfer can be given a
boundary check by viewing it as a single row of a 2-D surface.
Promote such a transfer to a single-row (`1xN`) block access. The descriptor
gains a leading unit dimension and `boundary_check = true`; the loaded value
is `vector.shape_cast` down to the original 1-D type, and the stored value up
to `1xN`. This matches transfer semantics exactly - an element is out of
bounds when its index is past the memref dim size, which is what the 2-D
block boundary check clamps on.
For a memref of rank > 1 the existing `convertMemrefAndOffsetsToTargetRank`
already yields the rank-2 view. A rank-1 memref cannot back a rank-2
descriptor (`create_nd_tdesc` rejects a TensorDesc rank above the source
rank, and forbids explicit shape/strides on a memref source), so it is
expanded with `memref.expand_shape` into a `1xD` surface - a contiguous
D-element buffer is one row of D elements, so its pitch is D as well. The
pattern bails out if that type cannot be expressed.
%dim = memref.dim %src, %c0 : memref<?xi32, strided<[1], offset: ?>>
%row = memref.expand_shape %src [[0, 1]] output_shape [1, %dim]
: memref<?xi32, strided<[1], offset: ?>> into memref<1x?xi32, strided<[?, 1], offset: ?>>
%desc = xegpu.create_nd_tdesc %row : memref<1x?xi32, ...> -> !xegpu.tensor_desc<1x128xi32>
%vec = xegpu.load_nd %desc[0, %off] : !xegpu.tensor_desc<1x128xi32> -> vector<1x128xi32>
The two `TODO: add support for OutOfBound access` bail-outs in the scattered
fallbacks become `notifyMatchFailure`, since that path still cannot express
out-of-bounds; only the nd path can.
The two `@no_*_out_of_bounds_1D_vector` negative tests become positive
LOAD-ND / STORE-ND tests - they still stay as transfers under the no-chip
runs, where block load/store is unavailable - plus new cases covering the
rank-1 memref row view.
Fixes #1899.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 125 ++++++++++++++----
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 47 ++++++-
.../transfer-write-to-xegpu.mlir | 46 ++++++-
3 files changed, 185 insertions(+), 33 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9863206f14fe1..779d415f59549 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -108,9 +108,9 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
unsigned vecRank = vecTy.getRank();
if (vecRank == 0)
return rewriter.notifyMatchFailure(xferOp, "0D vectors are not supported");
- if (xferOp.hasOutOfBoundsDim() && vecRank < 2)
- return rewriter.notifyMatchFailure(
- xferOp, "Boundary check is available only for block instructions.");
+ // A 1-D out-of-bounds transfer is not rejected here: it is promoted to a
+ // single-row 2-D block access, the only one with a boundary check. The
+ // individual patterns bail out if they cannot take that nd path.
AffineMap map = xferOp.getPermutationMap();
if (!map.isProjectedPermutation(/*allowZeroInResults=*/false))
@@ -421,6 +421,19 @@ convertMemrefAndOffsetsToTargetRank(PatternRewriter &rewriter, Location loc,
return {subviewOp.getResult(), newOffsets};
}
+// Reassociation viewing a 1-D memref as a `1xD` one, i.e. as a surface of a
+// single row holding all of its elements.
+static SmallVector<ReassociationIndices> getRowViewReassociation() {
+ return {ReassociationIndices{0, 1}};
+}
+
+// The `1xD` type of that view, or failure if it cannot be expressed.
+static FailureOr<MemRefType> getRowViewType(MemRefType memrefTy) {
+ assert(memrefTy.getRank() == 1 && "Expected a 1-D memref");
+ return memref::ExpandShapeOp::computeExpandedType(
+ memrefTy, {1, memrefTy.getDimSize(0)}, getRowViewReassociation());
+}
+
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
@@ -565,13 +578,29 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
AffineMap readMap = readOp.getPermutationMap();
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
+ // A 1-D out-of-bounds read has to take the nd path as well: the scattered
+ // load has no notion of bounds, and a 1-D nd load lowers to a 1-D block
+ // load, which has no boundary check either. Reading the vector as a single
+ // row of a 2-D surface puts it on the 2-D block load, whose boundary check
+ // is exactly the transfer's out-of-bounds semantics. A 1-D memref needs a
+ // `1xD` view to describe that surface; a higher-rank one already provides
+ // it below.
+ bool promoteToRow = loadedVecTy.getRank() == 1 && isOutOfBounds &&
+ readMap.isMinorIdentity();
+ FailureOr<MemRefType> rowViewTy = failure();
+ if (promoteToRow && readMemTy.getRank() == 1) {
+ rowViewTy = getRowViewType(readMemTy);
+ promoteToRow = succeeded(rowViewTy);
+ }
+
// 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.
+ // realize. In-bounds 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 &&
+ hasBlockLoadSupport && (loadedVecTy.getRank() > 1 || promoteToRow) &&
(readMap.isMinorIdentity() || isTransposeLoad) &&
readMemTy.getElementType().isIntOrFloat() &&
(!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
@@ -580,6 +609,10 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
auto elementType = loadedVecTy.getElementType();
SmallVector<int64_t> descShape(loadedVecTy.getShape());
+ if (promoteToRow) {
+ descShape.insert(descShape.begin(), 1);
+ loadedVecTy = VectorType::get(descShape, elementType);
+ }
if (isTransposeLoad) {
// If load is transposed, simply swap the last two dimensions of the
// loaded vector type to get the descriptor shape.
@@ -594,6 +627,12 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, readOp.getBase(),
getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
+ if (succeeded(rowViewTy)) {
+ src = memref::ExpandShapeOp::create(rewriter, loc, *rowViewTy, src,
+ getRowViewReassociation());
+ // The promoted row is the only row of the surface.
+ indices.insert(indices.begin(), rewriter.getIndexAttr(0));
+ }
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
@@ -605,6 +644,11 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
/*l1_hint=*/hint,
/*l2_hint=*/hint, /*l3_hint=*/hint,
/*layout=*/nullptr);
+ if (promoteToRow) {
+ // Drop the unit row dimension the promotion introduced.
+ loadedOp = vector::ShapeCastOp::create(
+ rewriter, loc, readOp.getVectorType(), loadedOp->getResult(0));
+ }
if (isTransposeLoad) {
// Undo the innermost-two-dims swap with a trailing vector.transpose:
// keep the leading dimensions in place and interchange only the last
@@ -621,9 +665,9 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
// 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 rewriter.notifyMatchFailure(
+ readOp, "Out-of-bounds reads require an nd block load");
return lowerToScatteredLoadOp(readOp, rewriter);
}
};
@@ -677,44 +721,71 @@ struct TransferWriteLowering
bool hasBlockStoreSupport =
(chip == "pvc" || chip == "bmg" || chip == "cri");
+ AffineMap map = writeOp.getPermutationMap();
+ bool isOutOfBounds = writeOp.hasOutOfBoundsDim();
+
+ // Mirrors the read side: a 1-D out-of-bounds write is stored as a single
+ // row of a 2-D surface, the only access with a boundary check.
+ bool promoteToRow =
+ vecTy.getRank() == 1 && isOutOfBounds && map.isMinorIdentity();
+ FailureOr<MemRefType> rowViewTy = failure();
+ if (promoteToRow && writeMemTy.getRank() == 1) {
+ rowViewTy = getRowViewType(writeMemTy);
+ promoteToRow = succeeded(rowViewTy);
+ }
+
// 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.
- AffineMap map = writeOp.getPermutationMap();
- bool canLowerToStoreNd = hasBlockStoreSupport && vecTy.getRank() > 1 &&
- map.isMinorIdentity() &&
- writeMemTy.getElementType().isIntOrFloat();
+ // (block stores have no transpose support). In-bounds 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.
+ bool canLowerToStoreNd =
+ hasBlockStoreSupport && (vecTy.getRank() > 1 || promoteToRow) &&
+ map.isMinorIdentity() && writeMemTy.getElementType().isIntOrFloat();
if (canLowerToStoreNd) {
+ Value storedVec = writeOp.getVector();
+ SmallVector<int64_t> descShape(vecTy.getShape());
+ if (promoteToRow) {
+ descShape.insert(descShape.begin(), 1);
+ storedVec = vector::ShapeCastOp::create(
+ rewriter, loc, VectorType::get(descShape, vecTy.getElementType()),
+ storedVec);
+ }
+
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, writeOp.getBase(),
- getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
+ getAsOpFoldResult(writeOp.getIndices()), descShape.size());
+ if (succeeded(rowViewTy)) {
+ src = memref::ExpandShapeOp::create(rewriter, loc, *rowViewTy, src,
+ getRowViewReassociation());
+ // The promoted row is the only row of the surface.
+ indices.insert(indices.begin(), rewriter.getIndexAttr(0));
+ }
auto descType = xegpu::TensorDescType::get(
- vecTy.getShape(), vecTy.getElementType(),
- /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
+ descShape, vecTy.getElementType(),
+ /*array_length=*/1, /*boundary_check=*/isOutOfBounds,
xegpu::MemorySpace::Global);
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
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, storedVec, 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.
- // TODO: add support for OutOfBound access.
- if (writeOp.hasOutOfBoundsDim())
- return failure();
+ if (isOutOfBounds)
+ return rewriter.notifyMatchFailure(
+ writeOp, "Out-of-bounds writes require an nd block store");
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 1464aa1ace477..08c75790ecf6e 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -417,7 +417,7 @@ gpu.func @no_load_out_of_bounds_non_zero_pad(%source: memref<32x64xf32>,
// -----
gpu.module @xevm_module {
-gpu.func @no_load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
+gpu.func @load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
%offset: index) -> vector<8xf32> {
%c0 = arith.constant 0.0 : f32
%0 = vector.transfer_read %source[%offset, %offset, %offset], %c0
@@ -425,8 +425,49 @@ gpu.func @no_load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
gpu.return %0 : vector<8xf32>
}
-// CHECK-LABEL: @no_load_out_of_bounds_1D_vector(
-// CHECK: vector.transfer_read
+// A 1D out-of-bounds read is promoted to a single-row block load, the only
+// access with a boundary check. The rank-3 source already collapses to a 2D
+// view, so no expand_shape is needed.
+// LOAD-ND-LABEL: @load_out_of_bounds_1D_vector(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// LOAD-ND-SAME: %[[OFFSET:.+]]: index
+// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SUBVIEW]]
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<1x8xf32>
+// LOAD-ND: %[[RES:.+]] = vector.shape_cast %[[VEC]] : vector<1x8xf32> to vector<8xf32>
+// LOAD-ND: return %[[RES]]
+
+// LOAD-GATHER-LABEL: @load_out_of_bounds_1D_vector(
+// LOAD-GATHER: vector.transfer_read
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_out_of_bounds_1D_memref(%source: memref<?xf32, strided<[1], offset: ?>>,
+ %offset: index) -> vector<8xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset], %c0
+ : memref<?xf32, strided<[1], offset: ?>>, vector<8xf32>
+ gpu.return %0 : vector<8xf32>
+}
+
+// A 1D memref is a surface of a single row holding all of its elements, so it
+// is expanded to a 1xD view to describe it.
+// LOAD-ND-LABEL: @load_out_of_bounds_1D_memref(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<?xf32, strided<[1], offset: ?>>,
+// LOAD-ND-SAME: %[[OFFSET:.+]]: index
+// LOAD-ND: %[[DIM:.+]] = memref.dim %[[SRC]]
+// LOAD-ND: %[[ROW:.+]] = memref.expand_shape %[[SRC]] {{\[\[}}0, 1{{\]\]}} output_shape [1, %[[DIM]]]
+// LOAD-ND-SAME: into memref<1x?xf32, strided<[?, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[ROW]]
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][0, %[[OFFSET]]]{{.*}}-> vector<1x8xf32>
+// LOAD-ND: %[[RES:.+]] = vector.shape_cast %[[VEC]] : vector<1x8xf32> to vector<8xf32>
+// LOAD-ND: return %[[RES]]
+
+// LOAD-GATHER-LABEL: @load_out_of_bounds_1D_memref(
+// LOAD-GATHER: vector.transfer_read
}
// -----
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index c6432378b496a..8b243b5738f83 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -316,7 +316,7 @@ gpu.func @no_store_unsupported_map(%vec: vector<8x16xf32>,
// -----
gpu.module @xevm_module {
-gpu.func @no_store_out_of_bounds_1D_vector(%vec: vector<8xf32>,
+gpu.func @store_out_of_bounds_1D_vector(%vec: vector<8xf32>,
%source: memref<8x16x32xf32>, %offset: index) {
vector.transfer_write %vec, %source[%offset, %offset, %offset]
{in_bounds = [false]}
@@ -324,8 +324,48 @@ gpu.func @no_store_out_of_bounds_1D_vector(%vec: vector<8xf32>,
gpu.return
}
-// CHECK-LABEL: @no_store_out_of_bounds_1D_vector(
-// CHECK: vector.transfer_write
+// A 1D out-of-bounds write is promoted to a single-row block store, the only
+// access with a boundary check. The rank-3 source already collapses to a 2D
+// view, so no expand_shape is needed.
+// STORE-ND-LABEL: @store_out_of_bounds_1D_vector(
+// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,
+// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// STORE-ND-SAME: %[[OFFSET:.+]]: index
+// STORE-ND: %[[ROW:.+]] = vector.shape_cast %[[VEC]] : vector<8xf32> to vector<1x8xf32>
+// STORE-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
+// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SUBVIEW]]
+// STORE-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
+// STORE-ND: xegpu.store_nd %[[ROW]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<1x8xf32>
+
+// STORE-SCATTER-LABEL: @store_out_of_bounds_1D_vector(
+// STORE-SCATTER: vector.transfer_write
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @store_out_of_bounds_1D_memref(%vec: vector<8xf32>,
+ %source: memref<?xf32, strided<[1], offset: ?>>, %offset: index) {
+ vector.transfer_write %vec, %source[%offset]
+ : vector<8xf32>, memref<?xf32, strided<[1], offset: ?>>
+ gpu.return
+}
+
+// A 1D memref is a surface of a single row holding all of its elements, so it
+// is expanded to a 1xD view to describe it.
+// STORE-ND-LABEL: @store_out_of_bounds_1D_memref(
+// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,
+// STORE-ND-SAME: %[[SRC:.+]]: memref<?xf32, strided<[1], offset: ?>>,
+// STORE-ND-SAME: %[[OFFSET:.+]]: index
+// STORE-ND: %[[ROW:.+]] = vector.shape_cast %[[VEC]] : vector<8xf32> to vector<1x8xf32>
+// STORE-ND: %[[DIM:.+]] = memref.dim %[[SRC]]
+// STORE-ND: %[[VIEW:.+]] = memref.expand_shape %[[SRC]] {{\[\[}}0, 1{{\]\]}} output_shape [1, %[[DIM]]]
+// STORE-ND-SAME: into memref<1x?xf32, strided<[?, 1], offset: ?>>
+// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[VIEW]]
+// STORE-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
+// STORE-ND: xegpu.store_nd %[[ROW]], %[[DESC]][0, %[[OFFSET]]] : vector<1x8xf32>
+
+// STORE-SCATTER-LABEL: @store_out_of_bounds_1D_memref(
+// STORE-SCATTER: vector.transfer_write
}
// -----
>From 032c71b5f7ebbaac13447ea41c1cb91b2bdd6c17 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 31 Aug 2026 21:50:49 +0000
Subject: [PATCH 2/3] [mlir][VectorToXeGPU] Mask off out-of-bounds elements in
the scattered path
A `vector.transfer_read` / `transfer_write` whose vector is 1-D and whose
access is out of bounds was not lowered at all: `transferPreconditions`
rejected it with "Boundary check is available only for block instructions."
The op then survived to `xegpu-wg-to-sg-distribute`, which has no transfer
pattern and marks unknown ops legal, so it stayed at the workgroup shape while
its consumers were distributed and the module failed to verify with
error: 'arith.divsi' op requires the same type for all operands and results
note: ... : (vector<128xi32>, vector<32xi32>) -> vector<32xi32>
Since `in_bounds` defaults to "out-of-bounds" when absent, this was the
default path for 1-D transfers, not a narrow corner.
Only a 2-D block instruction has a hardware boundary check, so the nd path
could not take a 1-D transfer. But the scattered path does not need one: its
mask already decides, per lane, whether memory is touched. Compute an
in-bounds mask instead of the `vector.constant_mask` all-ones. The element at
position `i` along vector dimension `v` accesses memref dimension `d` (the one
`v` maps to) at `indices[d] + i`, so it is in bounds iff
`i < dim(d) - indices[d]`:
%dim = memref.dim %src, %c0
%limit = arith.subi %dim, %off
%step = vector.step : vector<8xindex>
%bcast = vector.broadcast %limit : index to vector<8xindex>
%mask = arith.cmpi slt, %step, %bcast : vector<8xindex>
Dimensions the transfer declares in-bounds are skipped, so an all-in-bounds
transfer keeps the constant mask it had. Non-vector dimensions are always in
bounds by the op's contract. Per-dimension masks are spread over the full
vector shape and combined, mirroring how `computeOffsets` spreads offsets. A
negative limit masks the dimension off entirely, which is the right answer for
a starting index past the end.
The masked-off lanes of an `xegpu.load` are unspecified, so a read applies the
transfer's padding explicitly with `arith.select`, matching what
`GatherLowering` already does for `vector.gather`'s pass-thru. A poison padding
leaves those lanes "don't care", so the select is skipped. A write needs no
counterpart - masked lanes simply are not stored.
This also fixes a rank >= 2 out-of-bounds transfer with non-zero padding,
which the nd path declines (`load_nd` pads with zero) and the scattered path
previously could not take either.
Both `TODO: add support for OutOfBound access` bail-outs are removed.
Fixes #1899.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 214 +++++++++---------
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 76 +++----
.../transfer-write-to-xegpu.mlir | 64 ++----
3 files changed, 152 insertions(+), 202 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 779d415f59549..528be9137ac6c 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -108,9 +108,9 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
unsigned vecRank = vecTy.getRank();
if (vecRank == 0)
return rewriter.notifyMatchFailure(xferOp, "0D vectors are not supported");
- // A 1-D out-of-bounds transfer is not rejected here: it is promoted to a
- // single-row 2-D block access, the only one with a boundary check. The
- // individual patterns bail out if they cannot take that nd path.
+ // An out-of-bounds transfer is not rejected here. A block instruction gets a
+ // boundary check from its descriptor, and the scattered path masks off the
+ // out-of-bounds elements instead.
AffineMap map = xferOp.getPermutationMap();
if (!map.isProjectedPermutation(/*allowZeroInResults=*/false))
@@ -327,6 +327,68 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
return localOffsets;
}
+// Builds the mask a scattered access needs to stay within the source bounds.
+//
+// The element at vector position `i` along vector dimension `v` accesses memref
+// dimension `d` (the one `v` maps to) at `indices[d] + i`, so it is in bounds
+// iff `i < dim(d) - indices[d]`. Dimensions the transfer declares in-bounds are
+// skipped; non-vector dimensions are always in bounds by the op's contract. The
+// per-dimension masks are spread over the full vector shape and combined, the
+// same way `computeOffsets` spreads the per-dimension offsets.
+//
+// Example, for a `vector<8xf32>` read of a `memref<?xf32>` at `%off`:
+// %dim = memref.dim %src, %c0
+// %limit = arith.subi %dim, %off
+// %step = vector.step : vector<8xindex>
+// %bcast = vector.broadcast %limit : index to vector<8xindex>
+// %mask = arith.cmpi slt, %step, %bcast : vector<8xindex>
+static Value computeInBoundsMask(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter) {
+ Location loc = xferOp.getLoc();
+ ArrayRef<int64_t> vectorShape = xferOp.getVectorType().getShape();
+ auto maskType = VectorType::get(vectorShape, rewriter.getI1Type());
+ auto memrefTy = cast<MemRefType>(xferOp.getShapedType());
+ AffineMap map = xferOp.getPermutationMap();
+ OperandRange indices = xferOp.getIndices();
+
+ Value mask;
+ for (unsigned v = 0, e = vectorShape.size(); v < e; ++v) {
+ if (xferOp.isDimInBounds(v))
+ continue;
+ unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
+ Value bound =
+ memrefTy.isDynamicDim(d)
+ ? memref::DimOp::create(rewriter, loc, xferOp.getBase(), d)
+ .getResult()
+ : arith::ConstantIndexOp::create(rewriter, loc,
+ memrefTy.getDimSize(d))
+ .getResult();
+ // A negative limit masks the dimension off entirely, which is what a
+ // starting index past the end of the dimension should do.
+ Value limit = arith::SubIOp::create(rewriter, loc, bound, indices[d]);
+ auto stepType = VectorType::get({vectorShape[v]}, rewriter.getIndexType());
+ Value step = vector::StepOp::create(rewriter, loc, stepType);
+ Value limitVec =
+ vector::BroadcastOp::create(rewriter, loc, stepType, limit);
+ Value dimMask = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::slt, step, limitVec);
+ if (vectorShape.size() > 1) {
+ SmallVector<int64_t> expandedShape(vectorShape.size(), 1);
+ expandedShape[v] = vectorShape[v];
+ dimMask = vector::ShapeCastOp::create(
+ rewriter, loc, VectorType::get(expandedShape, rewriter.getI1Type()),
+ dimMask);
+ dimMask = vector::BroadcastOp::create(rewriter, loc, maskType, dimMask);
+ }
+ mask = mask
+ ? arith::AndIOp::create(rewriter, loc, mask, dimMask).getResult()
+ : dimMask;
+ }
+ if (mask)
+ return mask;
+ return vector::ConstantMaskOp::create(rewriter, loc, maskType, vectorShape);
+}
+
// Compute the element-wise offsets for vector.gather or vector.scatter ops.
//
// This function linearizes the base offsets of the gather/scatter operation
@@ -421,19 +483,6 @@ convertMemrefAndOffsetsToTargetRank(PatternRewriter &rewriter, Location loc,
return {subviewOp.getResult(), newOffsets};
}
-// Reassociation viewing a 1-D memref as a `1xD` one, i.e. as a surface of a
-// single row holding all of its elements.
-static SmallVector<ReassociationIndices> getRowViewReassociation() {
- return {ReassociationIndices{0, 1}};
-}
-
-// The `1xD` type of that view, or failure if it cannot be expressed.
-static FailureOr<MemRefType> getRowViewType(MemRefType memrefTy) {
- assert(memrefTy.getRank() == 1 && "Expected a 1-D memref");
- return memref::ExpandShapeOp::computeExpandedType(
- memrefTy, {1, memrefTy.getDimSize(0)}, getRowViewReassociation());
-}
-
template <
typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
@@ -455,7 +504,6 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
Location loc = readOp.getLoc();
VectorType vectorType = readOp.getVectorType();
- ArrayRef<int64_t> vectorShape = vectorType.getShape();
auto memrefType = dyn_cast<MemRefType>(readOp.getShapedType());
if (!memrefType)
return rewriter.notifyMatchFailure(readOp, "Expected memref source");
@@ -469,9 +517,7 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
Value flatMemref = memrefToIndexPtr(readOp, rewriter);
- Value mask = vector::ConstantMaskOp::create(
- rewriter, loc, VectorType::get(vectorShape, rewriter.getI1Type()),
- vectorShape);
+ Value mask = computeInBoundsMask(readOp, rewriter);
auto gatherOp = xegpu::LoadGatherOp::create(
rewriter, loc, vectorType, flatMemref, localOffsets, mask,
/*chunk_size=*/IntegerAttr{},
@@ -480,7 +526,18 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
/*l3_hint=*/xegpu::CachePolicyAttr{},
/*layout=*/nullptr, /*contiguity=*/nullptr);
- rewriter.replaceOp(readOp, gatherOp.getResult());
+ // The masked-off lanes of an xegpu.load are unspecified, so the transfer's
+ // padding has to be applied explicitly. A poison padding leaves them "don't
+ // care", so the select can be skipped.
+ Value result = gatherOp.getResult();
+ if (readOp.hasOutOfBoundsDim() &&
+ !readOp.getPadding().getDefiningOp<ub::PoisonOp>()) {
+ Value padding = vector::BroadcastOp::create(rewriter, loc, vectorType,
+ readOp.getPadding());
+ result = arith::SelectOp::create(rewriter, loc, mask, result, padding);
+ }
+
+ rewriter.replaceOp(readOp, result);
return success();
}
@@ -488,9 +545,6 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
PatternRewriter &rewriter) {
Location loc = writeOp.getLoc();
- VectorType vectorType = writeOp.getVectorType();
- ArrayRef<int64_t> vectorShape = vectorType.getShape();
-
auto memrefType = dyn_cast<MemRefType>(writeOp.getShapedType());
if (!memrefType)
return rewriter.notifyMatchFailure(writeOp, "Expected memref source");
@@ -504,9 +558,9 @@ static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
Value flatMemref = memrefToIndexPtr(writeOp, rewriter);
- Value mask = vector::ConstantMaskOp::create(
- rewriter, loc, VectorType::get(vectorShape, rewriter.getI1Type()),
- vectorShape);
+ // Out-of-bounds elements are simply not stored, so the mask is all this
+ // needs - no counterpart to the read's padding.
+ Value mask = computeInBoundsMask(writeOp, rewriter);
xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
localOffsets, mask,
/*chunk_size=*/IntegerAttr{},
@@ -578,29 +632,13 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
AffineMap readMap = readOp.getPermutationMap();
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
- // A 1-D out-of-bounds read has to take the nd path as well: the scattered
- // load has no notion of bounds, and a 1-D nd load lowers to a 1-D block
- // load, which has no boundary check either. Reading the vector as a single
- // row of a 2-D surface puts it on the 2-D block load, whose boundary check
- // is exactly the transfer's out-of-bounds semantics. A 1-D memref needs a
- // `1xD` view to describe that surface; a higher-rank one already provides
- // it below.
- bool promoteToRow = loadedVecTy.getRank() == 1 && isOutOfBounds &&
- readMap.isMinorIdentity();
- FailureOr<MemRefType> rowViewTy = failure();
- if (promoteToRow && readMemTy.getRank() == 1) {
- rowViewTy = getRowViewType(readMemTy);
- promoteToRow = succeeded(rowViewTy);
- }
-
// 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. In-bounds 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.
+ // 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.
bool canLowerToLoadNd =
- hasBlockLoadSupport && (loadedVecTy.getRank() > 1 || promoteToRow) &&
+ hasBlockLoadSupport && loadedVecTy.getRank() > 1 &&
(readMap.isMinorIdentity() || isTransposeLoad) &&
readMemTy.getElementType().isIntOrFloat() &&
(!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
@@ -609,10 +647,6 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
auto elementType = loadedVecTy.getElementType();
SmallVector<int64_t> descShape(loadedVecTy.getShape());
- if (promoteToRow) {
- descShape.insert(descShape.begin(), 1);
- loadedVecTy = VectorType::get(descShape, elementType);
- }
if (isTransposeLoad) {
// If load is transposed, simply swap the last two dimensions of the
// loaded vector type to get the descriptor shape.
@@ -627,12 +661,6 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, readOp.getBase(),
getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
- if (succeeded(rowViewTy)) {
- src = memref::ExpandShapeOp::create(rewriter, loc, *rowViewTy, src,
- getRowViewReassociation());
- // The promoted row is the only row of the surface.
- indices.insert(indices.begin(), rewriter.getIndexAttr(0));
- }
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
@@ -644,11 +672,6 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
/*l1_hint=*/hint,
/*l2_hint=*/hint, /*l3_hint=*/hint,
/*layout=*/nullptr);
- if (promoteToRow) {
- // Drop the unit row dimension the promotion introduced.
- loadedOp = vector::ShapeCastOp::create(
- rewriter, loc, readOp.getVectorType(), loadedOp->getResult(0));
- }
if (isTransposeLoad) {
// Undo the innermost-two-dims swap with a trailing vector.transpose:
// keep the leading dimensions in place and interchange only the last
@@ -664,10 +687,7 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
}
// Fall back to a scattered load. It supports arbitrary permutations and any
- // rank, but cannot express out-of-bounds accesses.
- if (isOutOfBounds)
- return rewriter.notifyMatchFailure(
- readOp, "Out-of-bounds reads require an nd block load");
+ // rank, and masks off the out-of-bounds elements.
return lowerToScatteredLoadOp(readOp, rewriter);
}
};
@@ -721,71 +741,41 @@ struct TransferWriteLowering
bool hasBlockStoreSupport =
(chip == "pvc" || chip == "bmg" || chip == "cri");
- AffineMap map = writeOp.getPermutationMap();
- bool isOutOfBounds = writeOp.hasOutOfBoundsDim();
-
- // Mirrors the read side: a 1-D out-of-bounds write is stored as a single
- // row of a 2-D surface, the only access with a boundary check.
- bool promoteToRow =
- vecTy.getRank() == 1 && isOutOfBounds && map.isMinorIdentity();
- FailureOr<MemRefType> rowViewTy = failure();
- if (promoteToRow && writeMemTy.getRank() == 1) {
- rowViewTy = getRowViewType(writeMemTy);
- promoteToRow = succeeded(rowViewTy);
- }
-
// 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). In-bounds 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.
- bool canLowerToStoreNd =
- hasBlockStoreSupport && (vecTy.getRank() > 1 || promoteToRow) &&
- map.isMinorIdentity() && writeMemTy.getElementType().isIntOrFloat();
+ // (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.
+ AffineMap map = writeOp.getPermutationMap();
+ bool canLowerToStoreNd = hasBlockStoreSupport && vecTy.getRank() > 1 &&
+ map.isMinorIdentity() &&
+ writeMemTy.getElementType().isIntOrFloat();
if (canLowerToStoreNd) {
- Value storedVec = writeOp.getVector();
- SmallVector<int64_t> descShape(vecTy.getShape());
- if (promoteToRow) {
- descShape.insert(descShape.begin(), 1);
- storedVec = vector::ShapeCastOp::create(
- rewriter, loc, VectorType::get(descShape, vecTy.getElementType()),
- storedVec);
- }
-
auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
rewriter, loc, writeOp.getBase(),
- getAsOpFoldResult(writeOp.getIndices()), descShape.size());
- if (succeeded(rowViewTy)) {
- src = memref::ExpandShapeOp::create(rewriter, loc, *rowViewTy, src,
- getRowViewReassociation());
- // The promoted row is the only row of the surface.
- indices.insert(indices.begin(), rewriter.getIndexAttr(0));
- }
+ getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
auto descType = xegpu::TensorDescType::get(
- descShape, vecTy.getElementType(),
- /*array_length=*/1, /*boundary_check=*/isOutOfBounds,
+ 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 = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
- auto storeOp =
- xegpu::StoreNdOp::create(rewriter, loc, storedVec, 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.
- if (isOutOfBounds)
- return rewriter.notifyMatchFailure(
- writeOp, "Out-of-bounds writes require an nd block store");
+ // any rank, and masks off the out-of-bounds elements.
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 08c75790ecf6e..3a9f86bee78bf 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -73,8 +73,14 @@ gpu.func @load_zero_pad_out_of_bounds(%source: memref<32x64xf32>,
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<8x16xf32>
// LOAD-ND: return %[[VEC]]
+// Only the outer dimension is out of bounds, so its 1D mask is spread over the
+// full vector shape.
// LOAD-GATHER-LABEL: @load_zero_pad_out_of_bounds(
-// LOAD-GATHER: vector.transfer_read
+// LOAD-GATHER: %[[CMP:.+]] = arith.cmpi slt, {{.*}} : vector<8xindex>
+// LOAD-GATHER: %[[CAST:.+]] = vector.shape_cast %[[CMP]] : vector<8xi1> to vector<8x1xi1>
+// LOAD-GATHER: %[[MASK:.+]] = vector.broadcast %[[CAST]] : vector<8x1xi1> to vector<8x16xi1>
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
+// LOAD-GATHER: arith.select %[[MASK]], %[[VEC]], %{{.+}} : vector<8x16xi1>, vector<8x16xf32>
}
// -----
@@ -401,7 +407,7 @@ gpu.func @load_transpose_f16(%source: memref<32x64xf16>,
// -----
gpu.module @xevm_module {
-gpu.func @no_load_out_of_bounds_non_zero_pad(%source: memref<32x64xf32>,
+gpu.func @load_out_of_bounds_non_zero_pad(%source: memref<32x64xf32>,
%offset: index, %arg2: index, %pad: f32) -> (vector<8x16xf32>, vector<8x16xf32>) {
%c1 = arith.constant 1.0 : f32
%0 = vector.transfer_read %source[%offset, %arg2], %c1
@@ -411,8 +417,11 @@ gpu.func @no_load_out_of_bounds_non_zero_pad(%source: memref<32x64xf32>,
gpu.return %0, %1 : vector<8x16xf32>, vector<8x16xf32>
}
-// CHECK-LABEL: @no_load_out_of_bounds_non_zero_pad(
-// CHECK-COUNT-2: vector.transfer_read
+// A non-zero padding does not match load_nd's implicit zero padding, so this
+// takes the scattered path, where the padding is applied explicitly.
+// CHECK-LABEL: @load_out_of_bounds_non_zero_pad(
+// CHECK-COUNT-2: arith.select
+// CHECK-NOT: vector.transfer_read
}
// -----
@@ -425,49 +434,22 @@ gpu.func @load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
gpu.return %0 : vector<8xf32>
}
-// A 1D out-of-bounds read is promoted to a single-row block load, the only
-// access with a boundary check. The rank-3 source already collapses to a 2D
-// view, so no expand_shape is needed.
-// LOAD-ND-LABEL: @load_out_of_bounds_1D_vector(
-// LOAD-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// LOAD-ND-SAME: %[[OFFSET:.+]]: index
-// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SUBVIEW]]
-// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<1x8xf32>
-// LOAD-ND: %[[RES:.+]] = vector.shape_cast %[[VEC]] : vector<1x8xf32> to vector<8xf32>
-// LOAD-ND: return %[[RES]]
-
-// LOAD-GATHER-LABEL: @load_out_of_bounds_1D_vector(
-// LOAD-GATHER: vector.transfer_read
-}
-
-// -----
-gpu.module @xevm_module {
-gpu.func @load_out_of_bounds_1D_memref(%source: memref<?xf32, strided<[1], offset: ?>>,
- %offset: index) -> vector<8xf32> {
- %c0 = arith.constant 0.0 : f32
- %0 = vector.transfer_read %source[%offset], %c0
- : memref<?xf32, strided<[1], offset: ?>>, vector<8xf32>
- gpu.return %0 : vector<8xf32>
-}
-
-// A 1D memref is a surface of a single row holding all of its elements, so it
-// is expanded to a 1xD view to describe it.
-// LOAD-ND-LABEL: @load_out_of_bounds_1D_memref(
-// LOAD-ND-SAME: %[[SRC:.+]]: memref<?xf32, strided<[1], offset: ?>>,
-// LOAD-ND-SAME: %[[OFFSET:.+]]: index
-// LOAD-ND: %[[DIM:.+]] = memref.dim %[[SRC]]
-// LOAD-ND: %[[ROW:.+]] = memref.expand_shape %[[SRC]] {{\[\[}}0, 1{{\]\]}} output_shape [1, %[[DIM]]]
-// LOAD-ND-SAME: into memref<1x?xf32, strided<[?, 1], offset: ?>>
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[ROW]]
-// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][0, %[[OFFSET]]]{{.*}}-> vector<1x8xf32>
-// LOAD-ND: %[[RES:.+]] = vector.shape_cast %[[VEC]] : vector<1x8xf32> to vector<8xf32>
-// LOAD-ND: return %[[RES]]
-
-// LOAD-GATHER-LABEL: @load_out_of_bounds_1D_memref(
-// LOAD-GATHER: vector.transfer_read
+// A 1D vector has no nd block instruction to get a boundary check from, so the
+// out-of-bounds elements are masked off in the scattered path instead. The
+// masked-off lanes of an xegpu.load are unspecified, hence the select applying
+// the transfer's padding.
+// CHECK-LABEL: @load_out_of_bounds_1D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK-DAG: %[[PAD:.+]] = arith.constant dense<0.000000e+00> : vector<8xf32>
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK: %[[LIMIT:.+]] = arith.subi %[[C32]], %[[OFFSET]] : index
+// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
+// CHECK: %[[LIMIT_V:.+]] = vector.broadcast %[[LIMIT]] : index to vector<8xindex>
+// CHECK: %[[MASK:.+]] = arith.cmpi slt, %[[STEP]], %[[LIMIT_V]] : vector<8xindex>
+// CHECK: %[[VEC:.+]] = xegpu.load {{.*}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
+// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PAD]] : vector<8xi1>, vector<8xf32>
+// CHECK: return %[[RES]]
}
// -----
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index 8b243b5738f83..48de892974bc2 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -137,8 +137,13 @@ gpu.func @store_out_of_bounds(%vec: vector<8x16xf32>,
// STORE-ND-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32>
// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>
+// Only the outer dimension is out of bounds, so its 1D mask is spread over the
+// full vector shape.
// STORE-SCATTER-LABEL: @store_out_of_bounds(
-// STORE-SCATTER: vector.transfer_write
+// STORE-SCATTER: %[[CMP:.+]] = arith.cmpi slt, {{.*}} : vector<8xindex>
+// STORE-SCATTER: %[[CAST:.+]] = vector.shape_cast %[[CMP]] : vector<8xi1> to vector<8x1xi1>
+// STORE-SCATTER: %[[MASK:.+]] = vector.broadcast %[[CAST]] : vector<8x1xi1> to vector<8x16xi1>
+// STORE-SCATTER: xegpu.store {{.*}}, %[[MASK]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>
}
// -----
@@ -324,48 +329,21 @@ gpu.func @store_out_of_bounds_1D_vector(%vec: vector<8xf32>,
gpu.return
}
-// A 1D out-of-bounds write is promoted to a single-row block store, the only
-// access with a boundary check. The rank-3 source already collapses to a 2D
-// view, so no expand_shape is needed.
-// STORE-ND-LABEL: @store_out_of_bounds_1D_vector(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[ROW:.+]] = vector.shape_cast %[[VEC]] : vector<8xf32> to vector<1x8xf32>
-// STORE-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SUBVIEW]]
-// STORE-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
-// STORE-ND: xegpu.store_nd %[[ROW]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<1x8xf32>
-
-// STORE-SCATTER-LABEL: @store_out_of_bounds_1D_vector(
-// STORE-SCATTER: vector.transfer_write
-}
-
-// -----
-gpu.module @xevm_module {
-gpu.func @store_out_of_bounds_1D_memref(%vec: vector<8xf32>,
- %source: memref<?xf32, strided<[1], offset: ?>>, %offset: index) {
- vector.transfer_write %vec, %source[%offset]
- : vector<8xf32>, memref<?xf32, strided<[1], offset: ?>>
- gpu.return
-}
-
-// A 1D memref is a surface of a single row holding all of its elements, so it
-// is expanded to a 1xD view to describe it.
-// STORE-ND-LABEL: @store_out_of_bounds_1D_memref(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<?xf32, strided<[1], offset: ?>>,
-// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[ROW:.+]] = vector.shape_cast %[[VEC]] : vector<8xf32> to vector<1x8xf32>
-// STORE-ND: %[[DIM:.+]] = memref.dim %[[SRC]]
-// STORE-ND: %[[VIEW:.+]] = memref.expand_shape %[[SRC]] {{\[\[}}0, 1{{\]\]}} output_shape [1, %[[DIM]]]
-// STORE-ND-SAME: into memref<1x?xf32, strided<[?, 1], offset: ?>>
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[VIEW]]
-// STORE-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
-// STORE-ND: xegpu.store_nd %[[ROW]], %[[DESC]][0, %[[OFFSET]]] : vector<1x8xf32>
-
-// STORE-SCATTER-LABEL: @store_out_of_bounds_1D_memref(
-// STORE-SCATTER: vector.transfer_write
+// A 1D vector has no nd block instruction to get a boundary check from, so the
+// out-of-bounds elements are masked off in the scattered path instead. Unlike a
+// read, a store needs no padding counterpart - the masked lanes just are not
+// written.
+// CHECK-LABEL: @store_out_of_bounds_1D_vector(
+// CHECK-SAME: %[[VEC:.+]]: vector<8xf32>,
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK: %[[LIMIT:.+]] = arith.subi %[[C32]], %[[OFFSET]] : index
+// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
+// CHECK: %[[LIMIT_V:.+]] = vector.broadcast %[[LIMIT]] : index to vector<8xindex>
+// CHECK: %[[MASK:.+]] = arith.cmpi slt, %[[STEP]], %[[LIMIT_V]] : vector<8xindex>
+// CHECK: xegpu.store %[[VEC]], {{.*}}, %[[MASK]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>
+// CHECK-NOT: arith.select
}
// -----
>From c781a727a97a860320b5aa8b516c1f4555447e8f Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 2 Sep 2026 05:38:10 +0000
Subject: [PATCH 3/3] [mlir][VectorToXeGPU] Lower a single-element transfer to
a scalar load
A `vector.transfer_read` of a 1-D single-element vector addresses exactly
one location, and the transfer's own indices already name it. It needs
neither a block descriptor nor a gather, but the scattered path was the
only one that would take it (`load_nd` requires rank >= 2), so it became a
one-lane `xegpu.load` with a one-lane mask and offset vector.
That is not just verbose, it does not survive the pipeline. Such a read
typically feeds a `vector.extract` to a scalar - the shape a frontend
leaves behind when it vectorizes a lookup like `expert_ids[pid_m]` - and a
scalar consumer requests no layout. `xegpu-propagate-layout` then leaves
the load, an `AnchorLayoutInterface` op, without a layout, while its
offset and mask operands do pick one up by forward fill (a
`vector.step : vector<1xindex>` CSE-shared with blocked loads that do have
consumers). `ResolveLayoutConflicts` trusts an anchor to declare a layout
per operand and fails the pass on the null:
error: No consumer layout found for vector operand.
note: see current operation: %289 = xegpu.load %285[%284], %288
: i64, vector<1xindex>, vector<1xi1> -> vector<1xi32>
Take the single-element case before the block/gather decision and emit a
scalar `memref.load` plus a `vector.broadcast`. The broadcast folds into
the extract, so nothing vector-shaped is left for layout propagation to
reason about.
An out-of-bounds transfer must still not touch memory - `in_bounds`
defaults to out-of-bounds, so this is the common case, not a corner - so
the load is guarded by an `scf.if` yielding the transfer's padding on the
out-of-bounds side:
%dim = memref.dim %src, %c0
%inb = arith.cmpi slt, %off, %dim : index
%val = scf.if %inb -> (i32) {
%0 = memref.load %src[%off] : memref<?xi32, strided<[1], offset: ?>>
scf.yield %0 : i32
} else {
scf.yield %pad : i32
}
%vec = vector.broadcast %val : i32 to vector<1xi32>
The guard is what keeps this from being a plain `memref.load`: clamping the
index instead would speculate a load the transfer never performs, which
faults on a zero-sized dimension. This is also why the upstream
`RewriteScalarExtractOfTransferRead` pattern declines an out-of-bounds
transfer outright. A fully in-bounds transfer gets no predicate and lowers
to a bare load and broadcast.
Only rank 1 is taken here; a higher-rank unit-size vector can still be a
block access. The dim-size computation moves out of `computeInBoundsMask`
into `getMemrefDimSize`, shared with the new predicate.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
mlir/include/mlir/Conversion/Passes.td | 2 +-
.../Conversion/VectorToXeGPU/CMakeLists.txt | 1 +
.../VectorToXeGPU/VectorToXeGPU.cpp | 112 ++++++++++++++++--
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 80 +++++++++++++
4 files changed, 186 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index f13cb9a801139..7e2225165f27c 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1758,7 +1758,7 @@ def ConvertVectorToXeGPU : Pass<"convert-vector-to-xegpu"> {
let summary = "Lower the operations from the vector dialect into the XeGPU "
"dialect";
let dependentDialects = [
- "memref::MemRefDialect", "arith::ArithDialect",
+ "memref::MemRefDialect", "arith::ArithDialect", "scf::SCFDialect",
"vector::VectorDialect", "xegpu::XeGPUDialect"
];
}
diff --git a/mlir/lib/Conversion/VectorToXeGPU/CMakeLists.txt b/mlir/lib/Conversion/VectorToXeGPU/CMakeLists.txt
index dd699f5b1880e..f526427b43bb5 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/CMakeLists.txt
+++ b/mlir/lib/Conversion/VectorToXeGPU/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_conversion_library(MLIRVectorToXeGPU
LINK_LIBS PUBLIC
MLIRArithDialect
MLIRMemRefDialect
+ MLIRSCFDialect
MLIRTransforms
MLIRVectorDialect
MLIRVectorToGPU
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 72ff735878b5b..3a66fb92c8a13 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/SCF/IR/SCF.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
@@ -328,6 +329,18 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
return localOffsets;
}
+// Returns the size of memref dimension `dim` as a value.
+static Value getMemrefDimSize(VectorTransferOpInterface xferOp, unsigned dim,
+ PatternRewriter &rewriter) {
+ Location loc = xferOp.getLoc();
+ auto memrefTy = cast<MemRefType>(xferOp.getShapedType());
+ if (memrefTy.isDynamicDim(dim))
+ return memref::DimOp::create(rewriter, loc, xferOp.getBase(), dim)
+ .getResult();
+ return arith::ConstantIndexOp::create(rewriter, loc, memrefTy.getDimSize(dim))
+ .getResult();
+}
+
// Builds the mask a scattered access needs to stay within the source bounds.
//
// The element at vector position `i` along vector dimension `v` accesses memref
@@ -348,7 +361,6 @@ static Value computeInBoundsMask(VectorTransferOpInterface xferOp,
Location loc = xferOp.getLoc();
ArrayRef<int64_t> vectorShape = xferOp.getVectorType().getShape();
auto maskType = VectorType::get(vectorShape, rewriter.getI1Type());
- auto memrefTy = cast<MemRefType>(xferOp.getShapedType());
AffineMap map = xferOp.getPermutationMap();
OperandRange indices = xferOp.getIndices();
@@ -357,13 +369,7 @@ static Value computeInBoundsMask(VectorTransferOpInterface xferOp,
if (xferOp.isDimInBounds(v))
continue;
unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
- Value bound =
- memrefTy.isDynamicDim(d)
- ? memref::DimOp::create(rewriter, loc, xferOp.getBase(), d)
- .getResult()
- : arith::ConstantIndexOp::create(rewriter, loc,
- memrefTy.getDimSize(d))
- .getResult();
+ Value bound = getMemrefDimSize(xferOp, d, rewriter);
// A negative limit masks the dimension off entirely, which is what a
// starting index past the end of the dimension should do.
Value limit = arith::SubIOp::create(rewriter, loc, bound, indices[d]);
@@ -500,6 +506,88 @@ static Value memrefToIndexPtr(OpType xferOp, PatternRewriter &rewriter) {
.getResult();
}
+// Builds the predicate that a single-element transfer is in bounds.
+//
+// The one element it touches sits at `indices`, so it is in bounds iff
+// `indices[d] < dim(d)` for every memref dimension `d` that a vector dimension
+// not declared in-bounds maps to. Dimensions the transfer declares in-bounds
+// are skipped; non-vector dimensions are always in bounds by the op's contract.
+// Returns a null value when every dimension is in bounds.
+static Value computeUnitInBoundsPredicate(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter) {
+ Location loc = xferOp.getLoc();
+ AffineMap map = xferOp.getPermutationMap();
+ OperandRange indices = xferOp.getIndices();
+
+ Value pred;
+ for (unsigned v = 0, e = xferOp.getVectorType().getRank(); v < e; ++v) {
+ if (xferOp.isDimInBounds(v))
+ continue;
+ unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
+ Value bound = getMemrefDimSize(xferOp, d, rewriter);
+ Value dimPred = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::slt, indices[d], bound);
+ pred = pred
+ ? arith::AndIOp::create(rewriter, loc, pred, dimPred).getResult()
+ : dimPred;
+ }
+ return pred;
+}
+
+// Lowers a transfer of a single element to a scalar `memref.load`.
+//
+// A `vector<1xT>` transfer addresses exactly one location, which the transfer's
+// own indices already name, so it needs neither a block descriptor nor a
+// gather - a scalar load is both shorter and the form the consumers of such a
+// read (typically a `vector.extract` feeding scalar index math) want anyway.
+//
+// An out-of-bounds transfer must not touch memory, so the load is guarded by an
+// `scf.if` that yields the transfer's padding on the out-of-bounds side. This
+// is what keeps the lowering from being a plain `memref.load`: clamping the
+// index instead would speculate a load that the transfer never performs, which
+// faults on a zero-sized dimension.
+//
+// %dim = memref.dim %src, %c0
+// %inb = arith.cmpi slt, %off, %dim : index
+// %val = scf.if %inb -> (f32) {
+// %0 = memref.load %src[%off] : memref<?xf32>
+// scf.yield %0 : f32
+// } else {
+// scf.yield %pad : f32
+// }
+// %vec = vector.broadcast %val : f32 to vector<1xf32>
+static LogicalResult lowerToScalarLoadOp(vector::TransferReadOp readOp,
+ PatternRewriter &rewriter) {
+ Location loc = readOp.getLoc();
+ if (!isa<MemRefType>(readOp.getShapedType()))
+ return rewriter.notifyMatchFailure(readOp, "Expected memref source");
+
+ Value inBounds = computeUnitInBoundsPredicate(readOp, rewriter);
+ Value scalar;
+ if (!inBounds) {
+ scalar = memref::LoadOp::create(rewriter, loc, readOp.getBase(),
+ readOp.getIndices())
+ .getResult();
+ } else {
+ auto ifOp = scf::IfOp::create(
+ rewriter, loc, inBounds,
+ [&](OpBuilder &builder, Location loc) {
+ Value loaded = memref::LoadOp::create(builder, loc, readOp.getBase(),
+ readOp.getIndices())
+ .getResult();
+ scf::YieldOp::create(builder, loc, loaded);
+ },
+ [&](OpBuilder &builder, Location loc) {
+ scf::YieldOp::create(builder, loc, readOp.getPadding());
+ });
+ scalar = ifOp.getResult(0);
+ }
+
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
+ readOp, readOp.getVectorType(), scalar);
+ return success();
+}
+
static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
PatternRewriter &rewriter) {
@@ -621,6 +709,14 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
return success();
}
+ // A 1-D transfer of a single element is a scalar load, not a one-lane
+ // gather. Taking it here keeps such a read - an out-of-bounds one in
+ // particular, which the scattered path would give a one-lane mask - out of
+ // the vector paths entirely. Only rank 1 is taken: a higher-rank unit-size
+ // vector can still be a block access.
+ if (loadedVecTy.getRank() == 1 && loadedVecTy.getNumElements() == 1)
+ return lowerToScalarLoadOp(readOp, rewriter);
+
// TODO: This check needs to be replaced with proper uArch capability check.
auto chip = xegpu::getChipStr(readOp);
bool hasBlockLoadSupport =
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 3a9f86bee78bf..4605c5c2f3344 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -452,6 +452,86 @@ gpu.func @load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
// CHECK: return %[[RES]]
}
+// -----
+gpu.module @xevm_module {
+gpu.func @load_unit_1D_vector(%source: memref<8x16x32xf32>,
+ %offset: index) -> vector<1xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset, %offset, %offset], %c0
+ {in_bounds = [true]} : memref<8x16x32xf32>, vector<1xf32>
+ gpu.return %0 : vector<1xf32>
+}
+
+// A single-element transfer addresses one location, which the transfer's own
+// indices already name, so it is a scalar load - no descriptor, no gather.
+// CHECK-LABEL: @load_unit_1D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK: %[[VAL:.+]] = memref.load %[[SRC]][%[[OFFSET]], %[[OFFSET]], %[[OFFSET]]] : memref<8x16x32xf32>
+// CHECK: %[[RES:.+]] = vector.broadcast %[[VAL]] : f32 to vector<1xf32>
+// CHECK-NOT: xegpu.load
+// CHECK: return %[[RES]]
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_out_of_bounds_unit_1D_vector(%source: memref<8x16x32xf32>,
+ %offset: index) -> vector<1xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset, %offset, %offset], %c0
+ {in_bounds = [false]} : memref<8x16x32xf32>, vector<1xf32>
+ gpu.return %0 : vector<1xf32>
+}
+
+// An out-of-bounds transfer must not touch memory, so the scalar load is
+// guarded and the padding is yielded on the out-of-bounds side.
+// CHECK-LABEL: @load_out_of_bounds_unit_1D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK: %[[INB:.+]] = arith.cmpi slt, %[[OFFSET]], %[[C32]] : index
+// CHECK: %[[VAL:.+]] = scf.if %[[INB]] -> (f32) {
+// CHECK: %[[LOAD:.+]] = memref.load %[[SRC]][%[[OFFSET]], %[[OFFSET]], %[[OFFSET]]] : memref<8x16x32xf32>
+// CHECK: scf.yield %[[LOAD]] : f32
+// CHECK: } else {
+// CHECK: scf.yield %[[PAD]] : f32
+// CHECK: }
+// CHECK: %[[RES:.+]] = vector.broadcast %[[VAL]] : f32 to vector<1xf32>
+// CHECK-NOT: xegpu.load
+// CHECK: return %[[RES]]
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_out_of_bounds_unit_1D_vector_dynamic(
+ %source: memref<?xi32, strided<[1], offset: ?>>, %offset: index) -> i32 {
+ %pad = ub.poison : i32
+ %0 = vector.transfer_read %source[%offset], %pad
+ : memref<?xi32, strided<[1], offset: ?>>, vector<1xi32>
+ %1 = vector.extract %0[0] : i32 from vector<1xi32>
+ gpu.return %1 : i32
+}
+
+// The extract of the broadcast folds away, leaving just the guarded scalar
+// load - the form a lookup like this one had before it was vectorized.
+// CHECK-LABEL: @load_out_of_bounds_unit_1D_vector_dynamic(
+// CHECK-SAME: %[[SRC:.+]]: memref<?xi32, strided<[1], offset: ?>>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK-DAG: %[[PAD:.+]] = ub.poison : i32
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[DIM:.+]] = memref.dim %[[SRC]], %[[C0]] : memref<?xi32, strided<[1], offset: ?>>
+// CHECK: %[[INB:.+]] = arith.cmpi slt, %[[OFFSET]], %[[DIM]] : index
+// CHECK: %[[RES:.+]] = scf.if %[[INB]] -> (i32) {
+// CHECK: %[[LOAD:.+]] = memref.load %[[SRC]][%[[OFFSET]]] : memref<?xi32, strided<[1], offset: ?>>
+// CHECK: scf.yield %[[LOAD]] : i32
+// CHECK: } else {
+// CHECK: scf.yield %[[PAD]] : i32
+// CHECK: }
+// CHECK-NOT: vector.broadcast
+// CHECK: return %[[RES]]
+}
+
// -----
gpu.module @xevm_module {
gpu.func @no_load_masked(%source : memref<4xf32>,
More information about the Mlir-commits
mailing list