[Mlir-commits] [mlir] 400cdf6 - [mlir][VectorToXeGPU] Lower out-of-bounds 1-D transfers (#220490)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 18 15:46:20 PDT 2026
Author: Jianhui Li
Date: 2026-09-18T15:46:15-07:00
New Revision: 400cdf6f379c0177e6259eeaf65ea8efa93326b3
URL: https://github.com/llvm/llvm-project/commit/400cdf6f379c0177e6259eeaf65ea8efa93326b3
DIFF: https://github.com/llvm/llvm-project/commit/400cdf6f379c0177e6259eeaf65ea8efa93326b3.diff
LOG: [mlir][VectorToXeGPU] Lower out-of-bounds 1-D transfers (#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.
It also lowers a transfer_read of a single-element vector (any rank) to
a scalar xegpu.load, when that vector is only ever extracted to a
scalar. The load takes a scalar offset and a scalar mask, and the mask
is what keeps an out-of-bounds transfer from touching memory; because a
masked-off load is unspecified, the padding is applied with an
arith.select, as the scattered path already does. This avoids
special-casing single-element vectors in layout propagation: a scalar
consumer requests no layout, so the one-lane xegpu.load would otherwise
be left without one.
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 500e5d055e4ad..5547eaa96f794 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -109,9 +109,6 @@ 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.");
AffineMap map = xferOp.getPermutationMap();
if (!map.isProjectedPermutation(/*allowZeroInResults=*/false))
@@ -221,6 +218,19 @@ computeMemrefMeta(OpType xferOp, PatternRewriter &rewriter) {
return {strides, offsetVal};
}
+// Adds the transfer's indices to `baseOffset`, the memref's own element offset:
+// the element at `indices` sits at `baseOffset + sum(indices[d] * strides[d])`.
+static Value computeBaseOffset(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter,
+ ArrayRef<Value> strides, Value baseOffset) {
+ Location loc = xferOp.getLoc();
+ for (auto [index, stride] : llvm::zip_equal(xferOp.getIndices(), strides)) {
+ Value contrib = arith::MulIOp::create(rewriter, loc, index, stride);
+ baseOffset = arith::AddIOp::create(rewriter, loc, baseOffset, contrib);
+ }
+ return baseOffset;
+}
+
// This function compute the vectors of localOffsets for scattered load/stores.
// It is used in the lowering of vector.transfer_read/write to
// load_gather/store_scatter Example:
@@ -254,8 +264,6 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
Value baseOffset) {
Location loc = xferOp.getLoc();
VectorType vectorType = xferOp.getVectorType();
- SmallVector<Value> indices(xferOp.getIndices().begin(),
- xferOp.getIndices().end());
ArrayRef<int64_t> vectorShape = vectorType.getShape();
// Create vector.step operations for each dimension
@@ -313,21 +321,112 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
localOffsets =
arith::AddIOp::create(rewriter, loc, localOffsets, broadcasted[i]);
- // Compute base offset from transfer read indices
- for (size_t i = 0; i < indices.size(); ++i) {
- Value strideVal = strides[i];
- Value offsetContrib =
- arith::MulIOp::create(rewriter, loc, indices[i], strideVal);
- baseOffset =
- arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
- }
// Broadcast base offset to match vector shape
+ baseOffset = computeBaseOffset(xferOp, rewriter, strides, baseOffset);
Value bcastBase = vector::BroadcastOp::create(
rewriter, loc, fullIndexVectorType, baseOffset);
localOffsets = arith::AddIOp::create(rewriter, loc, bcastBase, localOffsets);
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
+// 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());
+ 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 = 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]);
+ 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);
+}
+
+// Builds that mask as a scalar `i1`, for a transfer of a single element. The
+// element sits at `indices`, so `i < dim(d) - indices[d]` at `i == 0` reduces
+// to `indices[d] < dim(d)`, which needs neither the limit subtraction nor the
+// step vector. The dimensions checked are the same.
+//
+// Example, for a `vector<1xf32>` read of a `memref<?xf32>` at `%off`:
+// %dim = memref.dim %src, %c0
+// %mask = arith.cmpi slt, %off, %dim : index
+static Value computeUnitInBoundsMask(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter) {
+ Location loc = xferOp.getLoc();
+ AffineMap map = xferOp.getPermutationMap();
+ OperandRange indices = xferOp.getIndices();
+
+ Value mask;
+ 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 dimMask = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::slt, indices[d], bound);
+ mask = mask
+ ? arith::AndIOp::create(rewriter, loc, mask, dimMask).getResult()
+ : dimMask;
+ }
+ if (mask)
+ return mask;
+ return arith::ConstantOp::create(rewriter, loc, rewriter.getBoolAttr(true));
+}
+
// Compute the element-wise offsets for vector.gather or vector.scatter ops.
//
// This function linearizes the base offsets of the gather/scatter operation
@@ -438,12 +537,66 @@ static Value memrefToIndexPtr(OpType xferOp, PatternRewriter &rewriter) {
.getResult();
}
+// Returns true if every use of `vec` extracts a scalar element from it. An
+// unused value does not qualify.
+static bool isUsedAsScalar(Value vec) {
+ if (vec.use_empty())
+ return false;
+ return llvm::all_of(vec.getUsers(), [](Operation *user) {
+ auto extractOp = dyn_cast<vector::ExtractOp>(user);
+ return extractOp && !isa<VectorType>(extractOp.getResult().getType());
+ });
+}
+
+// Lowers a transfer of a single element to a scalar `xegpu.load`.
+//
+// The transfer touches one location, which its own indices name, so the load
+// takes a scalar offset and a scalar mask - no descriptor, no lane vectors -
+// and the result is broadcast back to the transfer's unit-size vector type.
+//
+// %off = arith.addi %base, %contrib : index
+// %inb = arith.cmpi slt, %idx, %dim : index
+// %val = xegpu.load %src[%off], %inb : i64, index, i1 -> f32
+// %vec = vector.broadcast %val : f32 to vector<1xf32>
+static LogicalResult lowerToScalarLoadOp(vector::TransferReadOp readOp,
+ PatternRewriter &rewriter) {
+ Location loc = readOp.getLoc();
+ VectorType vectorType = readOp.getVectorType();
+ if (!isa<MemRefType>(readOp.getShapedType()))
+ return rewriter.notifyMatchFailure(readOp, "Expected memref source");
+
+ auto meta = computeMemrefMeta(readOp, rewriter);
+ if (meta.first.empty())
+ return rewriter.notifyMatchFailure(readOp, "Failed to compute strides");
+
+ Value offset = computeBaseOffset(readOp, rewriter, meta.first, meta.second);
+ Value flatMemref = memrefToIndexPtr(readOp, rewriter);
+
+ Value mask = computeUnitInBoundsMask(readOp, rewriter);
+ auto loadOp = xegpu::LoadGatherOp::create(
+ rewriter, loc, vectorType.getElementType(), flatMemref, offset, mask,
+ /*l1_hint=*/xegpu::CachePolicyAttr{},
+ /*l2_hint=*/xegpu::CachePolicyAttr{},
+ /*l3_hint=*/xegpu::CachePolicyAttr{},
+ /*layout=*/nullptr, /*contiguity=*/nullptr);
+
+ // A masked-off xegpu.load is unspecified, so the padding has to be applied
+ // explicitly, as in lowerToScatteredLoadOp. A poison padding is "don't care".
+ Value scalar = loadOp.getResult();
+ if (readOp.hasOutOfBoundsDim() &&
+ !readOp.getPadding().getDefiningOp<ub::PoisonOp>())
+ scalar = arith::SelectOp::create(rewriter, loc, mask, scalar,
+ readOp.getPadding());
+
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vectorType, scalar);
+ return success();
+}
+
static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
PatternRewriter &rewriter) {
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");
@@ -457,9 +610,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,
/*l1_hint=*/xegpu::CachePolicyAttr{},
@@ -467,7 +618,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();
}
@@ -475,9 +637,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");
@@ -491,9 +650,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,
/*l1_hint=*/xegpu::CachePolicyAttr{},
@@ -552,6 +711,13 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
return success();
}
+ // A transfer of a single element that is only ever extracted to a scalar
+ // becomes a scalar load, at any rank: the vector is a wrapper its consumers
+ // undo, so neither an nd descriptor nor a lane-vector gather buys anything.
+ // A unit-size vector genuinely used as a vector keeps the paths below.
+ if (loadedVecTy.getNumElements() == 1 && isUsedAsScalar(readOp.getResult()))
+ return lowerToScalarLoadOp(readOp, rewriter);
+
// TODO: This check needs to be replaced with proper uArch capability check.
auto chip = xegpu::getChipStr(readOp);
bool hasBlockLoadSupport =
@@ -619,10 +785,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.
- // TODO: add support for OutOfBound access.
- if (isOutOfBounds)
- return failure();
+ // rank, and masks off the out-of-bounds elements.
return lowerToScatteredLoadOp(readOp, rewriter);
}
};
@@ -650,6 +813,10 @@ struct TransferWriteLowering
if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
return rewriter.notifyMatchFailure(
writeOp, "Only 1D and 2D vector stores are supported for SLM");
+ // Out of bounds case is not supported for SLM stores.
+ if (writeOp.hasOutOfBoundsDim())
+ return rewriter.notifyMatchFailure(
+ writeOp, "Out-of-bounds access is not supported for SLM stores");
// Create mem_desc for SLM
auto memDescType =
xegpu::MemDescType::get(rewriter.getContext(), writeMemTy.getShape(),
@@ -710,10 +877,7 @@ struct TransferWriteLowering
}
// 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();
+ // 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 1464aa1ace477..9eeb0c05de48f 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,13 +417,16 @@ 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
}
// -----
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 +434,142 @@ 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 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]]
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_unit_1D_vector(%source: memref<8x16x32xf32>,
+ %offset: index) -> f32 {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset, %offset, %offset], %c0
+ {in_bounds = [true]} : memref<8x16x32xf32>, vector<1xf32>
+ %1 = vector.extract %0[0] : f32 from vector<1xf32>
+ gpu.return %1 : f32
+}
+
+// A single-element transfer whose result is only extracted to a scalar touches
+// one location, which the transfer's own indices name, so the load takes a
+// scalar offset and mask - no descriptor, no lane vectors.
+// CHECK-LABEL: @load_unit_1D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK-NOT: vector.step
+// CHECK-DAG: %[[MASK:.+]] = arith.constant true
+// CHECK: %[[PTR:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
+// CHECK: %[[BASE:.+]] = arith.index_cast %[[PTR]] : index to i64
+// CHECK: %[[RES:.+]] = xegpu.load %[[BASE]][%{{.+}}], %[[MASK]] : i64, index, i1 -> f32
+// CHECK: return %[[RES]]
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_out_of_bounds_unit_1D_vector(%source: memref<8x16x32xf32>,
+ %offset: index) -> f32 {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset, %offset, %offset], %c0
+ {in_bounds = [false]} : memref<8x16x32xf32>, vector<1xf32>
+ %1 = vector.extract %0[0] : f32 from vector<1xf32>
+ gpu.return %1 : f32
+}
+
+// The mask keeps an out-of-bounds transfer from touching memory. A masked-off
+// load is unspecified, so the padding is applied with a select.
+// 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: %[[MASK:.+]] = arith.cmpi slt, %[[OFFSET]], %[[C32]] : index
+// CHECK: %[[VAL:.+]] = xegpu.load %{{.+}}[%{{.+}}], %[[MASK]] : i64, index, i1 -> f32
+// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VAL]], %[[PAD]] : f32
+// CHECK: return %[[RES]]
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @no_scalar_load_unit_1D_vector_used_as_vector(
+ %source: memref<8x16x32xf32>, %offset: index) -> vector<1xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset, %offset, %offset], %c0
+ : memref<8x16x32xf32>, vector<1xf32>
+ gpu.return %0 : vector<1xf32>
+}
+
+// A unit-size vector genuinely used as a vector keeps the lane-vector paths -
+// only a read its consumers unwrap to a scalar becomes a scalar load.
+// CHECK-LABEL: @no_scalar_load_unit_1D_vector_used_as_vector(
+// CHECK: vector.step : vector<1xindex>
+// CHECK: xegpu.load {{.*}} : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
+}
+
+// -----
+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
+}
+
+// A poison padding leaves the masked-off value "don't care", so no select is
+// needed. The extract of the broadcast folds away, leaving just the 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: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[DIM:.+]] = memref.dim %[[SRC]], %[[C0]] : memref<?xi32, strided<[1], offset: ?>>
+// CHECK: %[[MASK:.+]] = arith.cmpi slt, %[[OFFSET]], %[[DIM]] : index
+// CHECK: %[[RES:.+]] = xegpu.load %{{.+}}[%{{.+}}], %[[MASK]] : i64, index, i1 -> i32
+// CHECK-NOT: arith.select
+// CHECK-NOT: vector.broadcast
+// CHECK: return %[[RES]]
+}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_out_of_bounds_unit_2D_vector(%source: memref<8x16xf32>,
+ %offset: index) -> f32 {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%offset, %offset], %c0
+ : memref<8x16xf32>, vector<1x1xf32>
+ %1 = vector.extract %0[0, 0] : f32 from vector<1x1xf32>
+ gpu.return %1 : f32
+}
+
+// A single element is a scalar load at any rank - a 1x1 block descriptor buys
+// nothing. Each not-in-bounds dimension contributes a term to the mask.
+// CHECK-LABEL: @load_out_of_bounds_unit_2D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16xf32>,
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK-NOT: xegpu.load_nd
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[M0:.+]] = arith.cmpi slt, %[[OFFSET]], %[[C8]] : index
+// CHECK: %[[M1:.+]] = arith.cmpi slt, %[[OFFSET]], %[[C16]] : index
+// CHECK: %[[MASK:.+]] = arith.andi %[[M0]], %[[M1]] : i1
+// CHECK: %[[VAL:.+]] = xegpu.load %{{.+}}[%{{.+}}], %[[MASK]] : i64, index, i1 -> f32
+// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VAL]], %[[PAD]] : f32
+// 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 c6432378b496a..9842d8efcbe38 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>
}
// -----
@@ -316,7 +321,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 +329,21 @@ 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 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
}
// -----
@@ -400,6 +418,21 @@ gpu.func @store_1D_vector_addrspace3(%vec: vector<8xf32>,
// CHECK: gpu.return
}
+// -----
+gpu.module @xevm_module {
+gpu.func @no_store_addrspace3_out_of_bounds(%vec: vector<8xf32>,
+ %source: memref<32xf32, 3>, %offset: index) {
+ vector.transfer_write %vec, %source[%offset] {in_bounds = [false]}
+ : vector<8xf32>, memref<32xf32, 3>
+ gpu.return
+}
+
+// CHECK-LABEL: @no_store_addrspace3_out_of_bounds
+// CHECK-NOT: xegpu.create_mem_desc
+// CHECK-NOT: xegpu.store_matrix
+// CHECK: vector.transfer_write
+}
+
// -----
gpu.module @xevm_module {
gpu.func @store_0D_vector_unsupported(%vec: vector<f32>,
More information about the Mlir-commits
mailing list