[Mlir-commits] [mlir] [MLIR][XeGPU] Lower transfer_read with a broadcast permutation map (PR #225957)
Jianhui Li
llvmlistbot at llvm.org
Thu Sep 24 22:49:00 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/225957
>From 13ba1e2871da61fccaa384f7654b56cf54130589 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 23 Sep 2026 22:12:40 +0000
Subject: [PATCH 1/3] [MLIR][XeGPU] Fix stride order for non-involutive
transfer permutations
The scattered load/store path walks each vector dimension at the stride of
the memref dimension the permutation map has it read.
adjustStridesForPermutation derived that order via
inverseAndBroadcastProjectedPermutation and
isPermutationOfMinorIdentityWithBroadcasting, which only agrees with the map
for involutive permutations, i.e. a swap of two dimensions. For a rotation
such as (d0, d1, d2) -> (d2, d0, d1) it paired the vector dimensions with the
wrong strides, so the lowering accessed the wrong elements.
Replace it with getStridesInVectorOrder, which takes the stride straight off
each map result. Every pre-existing permutation test swaps two dimensions,
where the old code happened to agree, so add @load_rotated_dims to cover a
rotation.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 43 ++++++-------------
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 26 +++++++++++
2 files changed, 39 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 5547eaa96f7947..a827216cc207d9 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -16,7 +16,6 @@
#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"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
@@ -124,31 +123,19 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
return success();
}
-// Adjusts the strides of a memref according to a given permutation map for
-// vector operations.
-//
-// This function updates the innermost strides in the `strides` array to
-// reflect the permutation specified by `permMap`. The permutation is computed
-// using the inverse and broadcasting-aware version of the permutation map,
-// and is applied to the relevant strides. This ensures that memory accesses
-// are consistent with the logical permutation of vector elements.
+// Returns the memref strides in vector-dimension order: the stride at vector
+// dimension `v` is the stride of the memref dimension the permutation map has
+// `v` read, so that stepping one element along `v` steps that far in memory.
//
// Example:
-// Suppose we have a memref of rank 4 with strides `[s0, s1, s2, s3]`.
-// If the permutation map swaps the last two dimensions (e.g., [0, 1] -> [1,
-// 0]), then after calling this function, the last two strides will be
-// swapped:
-// Original strides: [s0, s1, s2, s3]
-// After permutation: [s0, s1, s3, s2]
+// For a memref of rank 3 with strides `[s0, s1, s2]` read into a rank-2
+// vector through `(d0, d1, d2) -> (d2, d1)`, the result is `[s2, s1]`.
//
-static void adjustStridesForPermutation(AffineMap permMap,
- SmallVectorImpl<Value> &strides) {
-
- AffineMap invMap = inverseAndBroadcastProjectedPermutation(permMap);
- SmallVector<unsigned> perms;
- invMap.isPermutationOfMinorIdentityWithBroadcasting(perms);
- SmallVector<int64_t> perms64(perms.begin(), perms.end());
- strides = applyPermutation(strides, perms64);
+static SmallVector<Value> getStridesInVectorOrder(AffineMap permMap,
+ ArrayRef<Value> strides) {
+ return llvm::map_to_vector(permMap.getResults(), [&](AffineExpr expr) {
+ return strides[cast<AffineDimExpr>(expr).getPosition()];
+ });
}
// Computes memory strides and a memref offset for vector transfer operations,
@@ -275,18 +262,14 @@ static Value computeOffsets(VectorTransferOpInterface xferOp,
return stepOp;
});
- // Local offsets are indexed in vector order, so permute strides; the base
- // offset below uses the original memref-order strides.
- SmallVector<Value> permutedStrides(strides.begin(), strides.end());
- adjustStridesForPermutation(xferOp.getPermutationMap(), permutedStrides);
+ SmallVector<Value> vectorStrides =
+ getStridesInVectorOrder(xferOp.getPermutationMap(), strides);
// Multiply step vectors by corresponding strides
- size_t memrefRank = permutedStrides.size();
size_t vectorRank = vectorShape.size();
SmallVector<Value> strideMultiplied;
for (size_t i = 0; i < vectorRank; ++i) {
- size_t memrefDim = memrefRank - vectorRank + i;
- Value strideValue = permutedStrides[memrefDim];
+ Value strideValue = vectorStrides[i];
auto mulType = dyn_cast<VectorType>(stepVectors[i].getType());
auto bcastOp =
vector::BroadcastOp::create(rewriter, loc, mulType, strideValue);
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 9eeb0c05de48f6..f449f0bb59a6bd 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -177,6 +177,32 @@ gpu.func @load_high_dim_transposed(%source: memref<2x2x64x128xf16>,
// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} : i64, vector<1x1x64x128xindex>, vector<1x1x64x128xi1> -> vector<1x1x64x128xf16>
}
+// -----
+// A permutation the nd block load cannot realize falls back to the scattered
+// load, which walks each vector dimension at the stride of the memref dimension
+// the map has it read: 1 for d2, 8192 for d0 and 128 for d1 here.
+gpu.module @xevm_module {
+gpu.func @load_rotated_dims(%source: memref<32x64x128xf32>,
+ %i: index, %j: index, %k: index) -> vector<2x4x8xf32> {
+ %c0 = arith.constant 0.0 : f32
+ %0 = vector.transfer_read %source[%i, %j, %k], %c0
+ {permutation_map = affine_map<(d0, d1, d2) -> (d2, d0, d1)>,
+ in_bounds = [true, true, true]} : memref<32x64x128xf32>, vector<2x4x8xf32>
+ gpu.return %0 : vector<2x4x8xf32>
+}
+
+// CHECK-LABEL: @load_rotated_dims(
+// CHECK: %[[STRIDE_D1:.+]] = arith.constant dense<128> : vector<8xindex>
+// CHECK: %[[STRIDE_D0:.+]] = arith.constant dense<8192> : vector<4xindex>
+// CHECK: %[[STEP_D2:.+]] = vector.step : vector<2xindex>
+// CHECK: %[[STEP_D0:.+]] = vector.step : vector<4xindex>
+// CHECK: %[[STEP_D1:.+]] = vector.step : vector<8xindex>
+// CHECK: arith.muli %[[STEP_D0]], %[[STRIDE_D0]] : vector<4xindex>
+// CHECK: arith.muli %[[STEP_D1]], %[[STRIDE_D1]] : vector<8xindex>
+// CHECK: %[[VEC:.+]] = xegpu.load {{.*}} -> vector<2x4x8xf32>
+// CHECK: return %[[VEC]]
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
>From d5f9ab5715783c2e2460e33c5019ae8a0d7731dd Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 23 Sep 2026 22:13:38 +0000
Subject: [PATCH 2/3] [MLIR][XeGPU] Lower transfer_read with a broadcast
permutation map
Grouped dequantization produces reads whose permutation map broadcasts some
vector dimensions, e.g. a (d0, d1, d2) -> (d2, 0) read of vector<64x64xf16>,
which loads 64 contiguous elements and repeats them across the second
dimension. vector-to-xegpu rejected those maps outright.
Lower them to a transfer of just the data in memory plus a vector.broadcast. A
broadcast result names no source dimension, so reading one element along it
reads all of the data, and it can be given any source dimension the map leaves
unnamed - the dimension is only read at its own index anyway. That turns the
map into a plain permutation the existing paths handle, so the example becomes
an nd block load of <1x64>, a transpose and a broadcast:
%narrow = xegpu.load_nd %desc[%g, %n] : ... -> vector<1x64xf16>
%t = vector.transpose %narrow, [1, 0] : vector<1x64xf16> to vector<64x1xf16>
%r = vector.broadcast %t : vector<64x1xf16> to vector<64x64xf16>
The peeling runs before every other path in TransferReadLowering, so the
scattered, SLM and scalar lowerings keep seeing a map that names a source
dimension per vector dimension. Transfer writes need no counterpart, as the
vector dialect forbids broadcast dimensions there.
Also reject memrefs of vector elements, whose maps have fewer results than the
vector's rank - previously only rejected by accident, through an unsigned
underflow in the innermost-dimension check.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 85 ++++++++++++++++++-
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 61 +++++++++++++
2 files changed, 145 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index a827216cc207d9..583fe6d3bf7b08 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -91,6 +91,9 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
auto srcTy = dyn_cast<MemRefType>(xferOp.getShapedType());
if (!srcTy)
return rewriter.notifyMatchFailure(xferOp, "Expects memref source");
+ if (isa<VectorType>(srcTy.getElementType()))
+ return rewriter.notifyMatchFailure(xferOp,
+ "Expects a memref of scalar elements");
// Validate further transfer op semantics.
SmallVector<int64_t> strides;
@@ -110,11 +113,14 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
return rewriter.notifyMatchFailure(xferOp, "0D vectors are not supported");
AffineMap map = xferOp.getPermutationMap();
- if (!map.isProjectedPermutation(/*allowZeroInResults=*/false))
+ if (!map.isProjectedPermutation(/*allowZeroInResults=*/true))
return rewriter.notifyMatchFailure(xferOp, "Unsupported permutation map");
unsigned numInputDims = map.getNumInputs();
for (AffineExpr expr : map.getResults().take_back(vecRank)) {
+ // A broadcast result names no source dimension, so it constrains nothing.
auto dim = dyn_cast<AffineDimExpr>(expr);
+ if (!dim)
+ continue;
if (dim.getPosition() < (numInputDims - vecRank))
return rewriter.notifyMatchFailure(
xferOp, "Only the innermost dimensions can be accessed");
@@ -616,6 +622,79 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
return success();
}
+// Lowers a transfer whose permutation map broadcasts some of the vector
+// dimensions into a narrower transfer plus a vector.broadcast. For example
+//
+// %r = vector.transfer_read %src[%i0, %i1, %i2], %pad
+// {permutation_map = affine_map<(d0, d1, d2) -> (d2, 0)>,
+// in_bounds = [false, false]} : memref<256x32x512xf16>, vector<64x64xf16>
+//
+// becomes
+//
+// %v = vector.transfer_read %src[%i0, %i1, %i2], %pad
+// {permutation_map = affine_map<(d0, d1, d2) -> (d2, d1)>,
+// in_bounds = [false, true]} : memref<256x32x512xf16>, vector<64x1xf16>
+// %r = vector.broadcast %v : vector<64x1xf16> to vector<64x64xf16>
+//
+// Each map result is either a source dimension d<k>, which reads real data, or
+// the constant 0 - a broadcast, along which that data only repeats. Only the
+// innermost `vecRank` source dimensions can be named, `dimOffset` being the
+// first of them and `isNamed` tracking which are taken. Result i of the
+// narrower transfer is built as:
+//
+// d<k>: readExprs[i] keeps d<k>, readShape[i] keeps the vector's extent, and
+// readInBounds[i] keeps the transfer's own bounds check.
+// 0: readExprs[i] takes the lowest source dimension `isNamed` still has
+// free, so that the map stays a permutation; readShape[i] becomes 1,
+// and readInBounds[i] becomes true, as a single element at that
+// dimension's own index needs no bounds check.
+//
+// Any mix works: the map names one source dimension per non-broadcast result,
+// so a spare one is always left for each broadcast result. What remains is a
+// plain permutation, which the other paths lower as usual.
+static LogicalResult lowerBroadcastRead(vector::TransferReadOp readOp,
+ PatternRewriter &rewriter) {
+ VectorType vecTy = readOp.getVectorType();
+ AffineMap map = readOp.getPermutationMap();
+ int64_t vecRank = vecTy.getRank();
+ int64_t dimOffset = map.getNumInputs() - vecRank;
+
+ SmallVector<int64_t> readShape(vecTy.getShape());
+ SmallVector<bool> readInBounds(vecRank, true);
+ SmallVector<AffineExpr> readExprs(map.getResults());
+ SmallVector<bool> isNamed(vecRank, false);
+ for (int64_t i = 0; i < vecRank; ++i) {
+ if (auto dim = dyn_cast<AffineDimExpr>(readExprs[i])) {
+ readInBounds[i] = readOp.isDimInBounds(i);
+ isNamed[dim.getPosition() - dimOffset] = true;
+ }
+ }
+
+ for (int64_t i = 0, leftoverDim = 0; i < vecRank; ++i) {
+ if (isa<AffineDimExpr>(readExprs[i]))
+ continue;
+ while (isNamed[leftoverDim])
+ ++leftoverDim;
+ isNamed[leftoverDim] = true;
+ readShape[i] = 1;
+ readExprs[i] = rewriter.getAffineDimExpr(dimOffset + leftoverDim);
+ }
+
+ auto readTy = VectorType::get(readShape, vecTy.getElementType());
+ Value read = vector::TransferReadOp::create(
+ rewriter, readOp.getLoc(), readTy, readOp.getBase(), readOp.getIndices(),
+ AffineMapAttr::get(
+ AffineMap::get(map.getNumInputs(), 0, readExprs, map.getContext())),
+ readOp.getPadding(), /*mask=*/Value(),
+ rewriter.getBoolArrayAttr(readInBounds));
+
+ if (readTy == vecTy)
+ rewriter.replaceOp(readOp, read);
+ else
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vecTy, read);
+ return success();
+}
+
static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
PatternRewriter &rewriter) {
@@ -655,6 +734,10 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
if (failed(transferPreconditions(rewriter, readOp)))
return failure();
+
+ if (!readOp.getPermutationMap().getBroadcastDims().empty())
+ return lowerBroadcastRead(readOp, rewriter);
+
auto readMemTy = cast<MemRefType>(readOp.getShapedType());
VectorType loadedVecTy = readOp.getVectorType();
bool isOutOfBounds = readOp.hasOutOfBoundsDim();
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index f449f0bb59a6bd..dc577093ef1179 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -203,6 +203,67 @@ gpu.func @load_rotated_dims(%source: memref<32x64x128xf32>,
// CHECK: return %[[VEC]]
}
+// -----
+// A broadcast dimension holds no data of its own, so the load shrinks to a
+// single element along it - here a 1x64 block over (d1, d2), transposed into
+// the vector order the map asks for - and a vector.broadcast stretches it back.
+gpu.module @xevm_module {
+gpu.func @load_broadcast_innermost_dim(%source: memref<256x32x512xf16>,
+ %e: index, %g: index, %n: index) -> vector<64x64xf16> {
+ %c0 = arith.constant 0.0 : f16
+ %0 = vector.transfer_read %source[%e, %g, %n], %c0
+ {permutation_map = affine_map<(d0, d1, d2) -> (d2, 0)>,
+ in_bounds = [true, true]} : memref<256x32x512xf16>, vector<64x64xf16>
+ gpu.return %0 : vector<64x64xf16>
+}
+
+// LOAD-ND-LABEL: @load_broadcast_innermost_dim(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<256x32x512xf16>,
+// LOAD-ND-SAME: %[[E:.+]]: index, %[[G:.+]]: index, %[[N:.+]]: index
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[E]], 0, 0] [1, 32, 512] [1, 1, 1] : memref<256x32x512xf16> to memref<32x512xf16, strided<[512, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[G]], %[[N]]]
+// LOAD-ND-SAME: -> vector<1x64xf16>
+// LOAD-ND: %[[TRANS:.+]] = vector.transpose %[[VEC]], [1, 0] : vector<1x64xf16> to vector<64x1xf16>
+// LOAD-ND: %[[BCAST:.+]] = vector.broadcast %[[TRANS]] : vector<64x1xf16> to vector<64x64xf16>
+// LOAD-ND: return %[[BCAST]]
+
+// LOAD-GATHER-LABEL: @load_broadcast_innermost_dim(
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} : i64, vector<64x1xindex>, vector<64x1xi1> -> vector<64x1xf16>
+// LOAD-GATHER: %[[BCAST:.+]] = vector.broadcast %[[VEC]] : vector<64x1xf16> to vector<64x64xf16>
+// LOAD-GATHER: return %[[BCAST]]
+}
+
+// -----
+// The dimension that does hold data need not be the innermost one; the load is
+// then a column, which needs no transpose.
+gpu.module @xevm_module {
+gpu.func @load_broadcast_outer_dim(%source: memref<256x512x32xf16>,
+ %e: index, %n: index, %g: index) -> vector<64x16xf16> {
+ %c0 = arith.constant 0.0 : f16
+ %0 = vector.transfer_read %source[%e, %n, %g], %c0
+ {permutation_map = affine_map<(d0, d1, d2) -> (d1, 0)>,
+ in_bounds = [true, true]} : memref<256x512x32xf16>, vector<64x16xf16>
+ gpu.return %0 : vector<64x16xf16>
+}
+
+// LOAD-ND-LABEL: @load_broadcast_outer_dim(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<256x512x32xf16>,
+// LOAD-ND-SAME: %[[E:.+]]: index, %[[N:.+]]: index, %[[G:.+]]: index
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[E]], 0, 0] [1, 512, 32] [1, 1, 1]
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<64x1xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[N]], %[[G]]]
+// LOAD-ND-SAME: -> vector<64x1xf16>
+// LOAD-ND-NOT: vector.transpose
+// LOAD-ND: vector.broadcast %[[VEC]] : vector<64x1xf16> to vector<64x16xf16>
+
+// LOAD-GATHER-LABEL: @load_broadcast_outer_dim(
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} -> vector<64x1xf16>
+// LOAD-GATHER: vector.broadcast %[[VEC]] : vector<64x1xf16> to vector<64x16xf16>
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
>From 30ce1cbd8c336e1cac926d995d1fb06be5f06aad Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 25 Sep 2026 05:47:49 +0000
Subject: [PATCH 3/3] [MLIR][XeGPU] Reduce broadcast reads to 1D transfers and
block-load those
Rework the broadcast lowering to drop the broadcast results from the map
instead of handing each one a source dimension the map left unnamed. What is
left is the transfer that reads the data - usually 1D - and the dropped
results come back as unit dimensions afterwards:
%v = vector.transfer_read %src[%i0, %i1, %i2], %pad
{permutation_map = affine_map<(d0, d1, d2) -> (d2)>} : vector<64xf16>
%c = vector.shape_cast %v : vector<64xf16> to vector<1x64xf16>
%t = vector.transpose %c, [1, 0] : vector<1x64xf16> to vector<64x1xf16>
%r = vector.broadcast %t : vector<64x1xf16> to vector<64x64xf16>
The shape_cast only appends unit dimensions, which leaves the dimension the
data came in on innermost, and the transpose moves them to the positions
their results held. Only unit dimensions move, so it permutes dimensions
without shuffling elements between them. Broadcasting first and transposing
after would instead leave a full-size transpose that cannot fold, as
FoldTransposeBroadcast needs the permutation to be local to unit-dim groups.
For that to still reach a block load, teach the nd path to take a 1D
transfer by giving it a leading unit dimension and shape_casting the result
back. Only a transfer along the innermost source dimension qualifies: that
dimension is contiguous, which transferPreconditions already checks, so the
<1xN> block is one the hardware can fetch. A transfer along any other
dimension would need <Nx1>, which the layouts downstream cannot distribute -
xegpu-propagate-layout rejects such a descriptor even with nothing else in
the function - so those keep the scattered path.
The precondition's innermost-dimension window gets one dimension of slack at
rank 1, as dropping the broadcast results of a map like (d1, 0) leaves a
rank-1 transfer that still reads d1, and the scattered path reads along any
dimension anyway.
1D reads reach load_nd now, so the 1D tests split per configuration.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 159 ++++++++++++-----
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 168 +++++++++++-------
2 files changed, 218 insertions(+), 109 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 583fe6d3bf7b08..5b9dbb96b18899 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -115,13 +115,20 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
AffineMap map = xferOp.getPermutationMap();
if (!map.isProjectedPermutation(/*allowZeroInResults=*/true))
return rewriter.notifyMatchFailure(xferOp, "Unsupported permutation map");
+ // The transfer is built from a run of innermost source dimensions, as wide as
+ // the vector, with one dimension of slack for a 1D transfer: dropping the
+ // broadcast results of a map like (d1, 0) leaves a rank-1 transfer that still
+ // reads d1, and the scattered path reads along any dimension anyway.
unsigned numInputDims = map.getNumInputs();
+ unsigned numAccessibleDims = std::max(vecRank, 2u);
+ unsigned firstAccessibleDim =
+ numInputDims > numAccessibleDims ? numInputDims - numAccessibleDims : 0;
for (AffineExpr expr : map.getResults().take_back(vecRank)) {
// A broadcast result names no source dimension, so it constrains nothing.
auto dim = dyn_cast<AffineDimExpr>(expr);
if (!dim)
continue;
- if (dim.getPosition() < (numInputDims - vecRank))
+ if (dim.getPosition() < firstAccessibleDim)
return rewriter.notifyMatchFailure(
xferOp, "Only the innermost dimensions can be accessed");
}
@@ -622,76 +629,131 @@ static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
return success();
}
+// An nd block load fetches a 2D block, which a 1D transfer does not ask for -
+// but a leading unit dimension gives it one:
+//
+// %v = vector.transfer_read %src[%i0, %i1] : memref<32x64xf32>, vector<8xf32>
+// becomes
+// %b = vector.transfer_read %src[%i0, %i1]
+// : memref<32x64xf32>, vector<1x8xf32>
+// %v = vector.shape_cast %b : vector<1x8xf32> to vector<8xf32>
+//
+// Only a transfer running along the innermost source dimension qualifies. That
+// dimension is contiguous, which transferPreconditions has already checked, so
+// the <1xN> block it asks for is one the hardware can fetch. A transfer along
+// any other dimension would need the unit dimension on the other side, and a
+// <Nx1> block is not a shape the layouts downstream can distribute - the
+// scattered path is the one that fits it.
+static LogicalResult lowerToUnitDimRead(vector::TransferReadOp readOp,
+ PatternRewriter &rewriter) {
+ VectorType vecTy = readOp.getVectorType();
+ AffineMap map = readOp.getPermutationMap();
+ int64_t numInputs = map.getNumInputs();
+ int64_t pos = cast<AffineDimExpr>(map.getResult(0)).getPosition();
+ if (numInputs < 2 || pos != numInputs - 1)
+ return rewriter.notifyMatchFailure(
+ readOp, "1D transfer does not run along the innermost dimension");
+
+ Value block = vector::TransferReadOp::create(
+ rewriter, readOp.getLoc(),
+ VectorType::get({1, vecTy.getDimSize(0)}, vecTy.getElementType()),
+ readOp.getBase(), readOp.getIndices(),
+ AffineMapAttr::get(
+ AffineMap::getMinorIdentityMap(numInputs, 2, map.getContext())),
+ readOp.getPadding(), /*mask=*/Value(),
+ rewriter.getBoolArrayAttr({true, readOp.isDimInBounds(0)}));
+
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, vecTy, block);
+ return success();
+}
+
// Lowers a transfer whose permutation map broadcasts some of the vector
-// dimensions into a narrower transfer plus a vector.broadcast. For example
+// dimensions into a transfer of just the data in memory - the map with its
+// broadcast results dropped - plus the vector ops that expand it back:
//
// %r = vector.transfer_read %src[%i0, %i1, %i2], %pad
// {permutation_map = affine_map<(d0, d1, d2) -> (d2, 0)>,
// in_bounds = [false, false]} : memref<256x32x512xf16>, vector<64x64xf16>
-//
// becomes
-//
// %v = vector.transfer_read %src[%i0, %i1, %i2], %pad
-// {permutation_map = affine_map<(d0, d1, d2) -> (d2, d1)>,
-// in_bounds = [false, true]} : memref<256x32x512xf16>, vector<64x1xf16>
-// %r = vector.broadcast %v : vector<64x1xf16> to vector<64x64xf16>
-//
-// Each map result is either a source dimension d<k>, which reads real data, or
-// the constant 0 - a broadcast, along which that data only repeats. Only the
-// innermost `vecRank` source dimensions can be named, `dimOffset` being the
-// first of them and `isNamed` tracking which are taken. Result i of the
-// narrower transfer is built as:
-//
-// d<k>: readExprs[i] keeps d<k>, readShape[i] keeps the vector's extent, and
-// readInBounds[i] keeps the transfer's own bounds check.
-// 0: readExprs[i] takes the lowest source dimension `isNamed` still has
-// free, so that the map stays a permutation; readShape[i] becomes 1,
-// and readInBounds[i] becomes true, as a single element at that
-// dimension's own index needs no bounds check.
+// {permutation_map = affine_map<(d0, d1, d2) -> (d2)>,
+// in_bounds = [false]} : memref<256x32x512xf16>, vector<64xf16>
+// %c = vector.shape_cast %v : vector<64xf16> to vector<1x64xf16>
+// %t = vector.transpose %c, [1, 0] : vector<1x64xf16> to vector<64x1xf16>
+// %r = vector.broadcast %t : vector<64x1xf16> to vector<64x64xf16>
//
-// Any mix works: the map names one source dimension per non-broadcast result,
-// so a spare one is always left for each broadcast result. What remains is a
-// plain permutation, which the other paths lower as usual.
+// A broadcast result names no source dimension, so only the dimensions the map
+// does name hold data; dropping the rest leaves the transfer that reads it, at
+// the extents and bounds checks those dimensions had. The dropped results come
+// back as unit dimensions, which the transpose moves to the positions those
+// results held, and vector.broadcast stretches to their final extents.
static LogicalResult lowerBroadcastRead(vector::TransferReadOp readOp,
PatternRewriter &rewriter) {
+ Location loc = readOp.getLoc();
VectorType vecTy = readOp.getVectorType();
+ ArrayRef<int64_t> vecShape = vecTy.getShape();
AffineMap map = readOp.getPermutationMap();
int64_t vecRank = vecTy.getRank();
- int64_t dimOffset = map.getNumInputs() - vecRank;
- SmallVector<int64_t> readShape(vecTy.getShape());
- SmallVector<bool> readInBounds(vecRank, true);
- SmallVector<AffineExpr> readExprs(map.getResults());
- SmallVector<bool> isNamed(vecRank, false);
+ SmallVector<int64_t> dataPos, bcastPos;
for (int64_t i = 0; i < vecRank; ++i) {
- if (auto dim = dyn_cast<AffineDimExpr>(readExprs[i])) {
- readInBounds[i] = readOp.isDimInBounds(i);
- isNamed[dim.getPosition() - dimOffset] = true;
- }
+ if (isa<AffineDimExpr>(map.getResult(i)))
+ dataPos.push_back(i);
+ else
+ bcastPos.push_back(i);
}
- for (int64_t i = 0, leftoverDim = 0; i < vecRank; ++i) {
- if (isa<AffineDimExpr>(readExprs[i]))
- continue;
- while (isNamed[leftoverDim])
- ++leftoverDim;
- isNamed[leftoverDim] = true;
- readShape[i] = 1;
- readExprs[i] = rewriter.getAffineDimExpr(dimOffset + leftoverDim);
+ SmallVector<AffineExpr> readExprs;
+ SmallVector<int64_t> readShape;
+ SmallVector<bool> readInBounds;
+ for (int64_t i : dataPos) {
+ readExprs.push_back(map.getResult(i));
+ readShape.push_back(vecShape[i]);
+ readInBounds.push_back(readOp.isDimInBounds(i));
+ }
+ // An all-broadcast map leaves no data dimension to read. Read one element of
+ // the innermost source dimension instead - a 0D transfer is not supported,
+ // and one element is all the data such a transfer holds.
+ if (readExprs.empty()) {
+ readExprs.push_back(rewriter.getAffineDimExpr(map.getNumInputs() - 1));
+ readShape.push_back(1);
+ readInBounds.push_back(true);
}
auto readTy = VectorType::get(readShape, vecTy.getElementType());
Value read = vector::TransferReadOp::create(
- rewriter, readOp.getLoc(), readTy, readOp.getBase(), readOp.getIndices(),
+ rewriter, loc, readTy, readOp.getBase(), readOp.getIndices(),
AffineMapAttr::get(
AffineMap::get(map.getNumInputs(), 0, readExprs, map.getContext())),
readOp.getPadding(), /*mask=*/Value(),
rewriter.getBoolArrayAttr(readInBounds));
- if (readTy == vecTy)
- rewriter.replaceOp(readOp, read);
+ // Grow the transfer back to the requested rank, one unit dimension per
+ // dropped result. The shape_cast adds them innermost, where they leave the
+ // dimension the data came in on innermost too, and the transpose moves them
+ // to the positions their results held. Only unit dimensions move, so the
+ // transpose permutes dimensions without shuffling elements between them.
+ SmallVector<int64_t> order(bcastPos);
+ llvm::append_range(order, dataPos);
+ SmallVector<int64_t> castShape(bcastPos.size(), 1);
+ for (int64_t i : dataPos)
+ castShape.push_back(vecShape[i]);
+ auto castTy = VectorType::get(castShape, vecTy.getElementType());
+ Value cast = castTy == readTy
+ ? read
+ : vector::ShapeCastOp::create(rewriter, loc, castTy, read);
+
+ if (!llvm::is_sorted(order)) {
+ SmallVector<int64_t> perm(vecRank);
+ for (int64_t i = 0; i < vecRank; ++i)
+ perm[order[i]] = i;
+ cast = vector::TransposeOp::create(rewriter, loc, cast, perm);
+ }
+
+ if (cast.getType() == vecTy)
+ rewriter.replaceOp(readOp, cast);
else
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vecTy, read);
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vecTy, cast);
return success();
}
@@ -796,6 +858,15 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
AffineMap readMap = readOp.getPermutationMap();
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
+ // A 1D transfer has no 2D block for an nd load to fetch, but it can borrow
+ // a unit dimension to get one. Only worth it where a block load is what we
+ // are heading for; otherwise the scattered path takes the 1D vector as is.
+ if (hasBlockLoadSupport && loadedVecTy.getRank() == 1 &&
+ readMemTy.getElementType().isIntOrFloat() &&
+ (!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding())) &&
+ succeeded(lowerToUnitDimRead(readOp, rewriter)))
+ return success();
+
// 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
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index dc577093ef1179..8bf0dc2cd6841f 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -9,17 +9,28 @@ gpu.func @load_1D_vector(%source: memref<8x16x32xf32>, %offset: index) -> vector
gpu.return %0 : vector<8xf32>
}
-// CHECK-LABEL: @load_1D_vector(
-// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
-// CHECK-COUNT-2: arith.muli {{.*}} : index
-// CHECK-COUNT-2: 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 %[[SRC]] : memref<8x16x32xf32> -> index
-// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
+// LOAD-ND-LABEL: @load_1D_vector(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// LOAD-ND-SAME: %[[OFFSET:.+]]: index
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0] [1, 16, 32] [1, 1, 1]
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]
+// LOAD-ND-SAME: -> vector<1x8xf32>
+// LOAD-ND: %[[RES:.+]] = vector.shape_cast %[[VEC]] : vector<1x8xf32> to vector<8xf32>
+// LOAD-ND: return %[[RES]]
+
+// LOAD-GATHER-LABEL: @load_1D_vector(
+// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// LOAD-GATHER: %[[STEP:.+]] = vector.step : vector<8xindex>
+// LOAD-GATHER-COUNT-2: arith.muli {{.*}} : index
+// LOAD-GATHER-COUNT-2: arith.addi {{.*}} : index
+// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
+// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
+// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
+// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
}
// -----
@@ -230,14 +241,18 @@ gpu.func @load_broadcast_innermost_dim(%source: memref<256x32x512xf16>,
// LOAD-ND: return %[[BCAST]]
// LOAD-GATHER-LABEL: @load_broadcast_innermost_dim(
-// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} : i64, vector<64x1xindex>, vector<64x1xi1> -> vector<64x1xf16>
-// LOAD-GATHER: %[[BCAST:.+]] = vector.broadcast %[[VEC]] : vector<64x1xf16> to vector<64x64xf16>
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} : i64, vector<64xindex>, vector<64xi1> -> vector<64xf16>
+// LOAD-GATHER: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<64xf16> to vector<1x64xf16>
+// LOAD-GATHER: %[[TRANS:.+]] = vector.transpose %[[CAST]], [1, 0] : vector<1x64xf16> to vector<64x1xf16>
+// LOAD-GATHER: %[[BCAST:.+]] = vector.broadcast %[[TRANS]] : vector<64x1xf16> to vector<64x64xf16>
// LOAD-GATHER: return %[[BCAST]]
}
// -----
-// The dimension that does hold data need not be the innermost one; the load is
-// then a column, which needs no transpose.
+// The dimension that holds the data need not be the innermost one. It is then
+// strided, so the read falls to the scattered path: a <Nx1> block is the only
+// one an nd load could fetch, and that is not a shape the layouts downstream can
+// distribute.
gpu.module @xevm_module {
gpu.func @load_broadcast_outer_dim(%source: memref<256x512x32xf16>,
%e: index, %n: index, %g: index) -> vector<64x16xf16> {
@@ -248,20 +263,16 @@ gpu.func @load_broadcast_outer_dim(%source: memref<256x512x32xf16>,
gpu.return %0 : vector<64x16xf16>
}
-// LOAD-ND-LABEL: @load_broadcast_outer_dim(
-// LOAD-ND-SAME: %[[SRC:.+]]: memref<256x512x32xf16>,
-// LOAD-ND-SAME: %[[E:.+]]: index, %[[N:.+]]: index, %[[G:.+]]: index
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[E]], 0, 0] [1, 512, 32] [1, 1, 1]
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]
-// LOAD-ND-SAME: -> !xegpu.tensor_desc<64x1xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[N]], %[[G]]]
-// LOAD-ND-SAME: -> vector<64x1xf16>
-// LOAD-ND-NOT: vector.transpose
-// LOAD-ND: vector.broadcast %[[VEC]] : vector<64x1xf16> to vector<64x16xf16>
-
-// LOAD-GATHER-LABEL: @load_broadcast_outer_dim(
-// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}} -> vector<64x1xf16>
-// LOAD-GATHER: vector.broadcast %[[VEC]] : vector<64x1xf16> to vector<64x16xf16>
+// CHECK-LABEL: @load_broadcast_outer_dim(
+// CHECK-SAME: %[[SRC:.+]]: memref<256x512x32xf16>,
+// CHECK-SAME: %[[E:.+]]: index, %[[N:.+]]: index, %[[G:.+]]: index
+// CHECK: %[[STRIDE:.+]] = arith.constant dense<32> : vector<64xindex>
+// CHECK: %[[STEP:.+]] = vector.step : vector<64xindex>
+// CHECK: arith.muli %[[STEP]], %[[STRIDE]] : vector<64xindex>
+// CHECK: %[[VEC:.+]] = xegpu.load {{.*}} -> vector<64xf16>
+// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<64xf16> to vector<1x64xf16>
+// CHECK: %[[TRANS:.+]] = vector.transpose %[[CAST]], [1, 0] : vector<1x64xf16> to vector<64x1xf16>
+// CHECK: vector.broadcast %[[TRANS]] : vector<64x1xf16> to vector<64x16xf16>
}
// -----
@@ -521,22 +532,34 @@ gpu.func @load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
gpu.return %0 : vector<8xf32>
}
-// 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]]
+// A 1D transfer borrows a unit dimension to reach an nd block load, whose
+// descriptor carries the boundary check; the zero padding the transfer asks for
+// is what load_nd fills out-of-bounds elements with anyway.
+// LOAD-ND-LABEL: @load_out_of_bounds_1D_vector(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// LOAD-ND-SAME: %[[OFFSET:.+]]: index
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf32>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]
+// LOAD-ND-SAME: -> vector<1x8xf32>
+// LOAD-ND: %[[RES:.+]] = vector.shape_cast %[[VEC]] : vector<1x8xf32> to vector<8xf32>
+// LOAD-ND: return %[[RES]]
+
+// Without block-load support 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.
+// LOAD-GATHER-LABEL: @load_out_of_bounds_1D_vector(
+// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// LOAD-GATHER-SAME: %[[OFFSET:.+]]: index
+// LOAD-GATHER-DAG: %[[PAD:.+]] = arith.constant dense<0.000000e+00> : vector<8xf32>
+// LOAD-GATHER-DAG: %[[C32:.+]] = arith.constant 32 : index
+// LOAD-GATHER: %[[LIMIT:.+]] = arith.subi %[[C32]], %[[OFFSET]] : index
+// LOAD-GATHER: %[[STEP:.+]] = vector.step : vector<8xindex>
+// LOAD-GATHER: %[[LIMIT_V:.+]] = vector.broadcast %[[LIMIT]] : index to vector<8xindex>
+// LOAD-GATHER: %[[MASK:.+]] = arith.cmpi slt, %[[STEP]], %[[LIMIT_V]] : vector<8xindex>
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load {{.*}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
+// LOAD-GATHER: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PAD]] : vector<8xi1>, vector<8xf32>
+// LOAD-GATHER: return %[[RES]]
}
// -----
@@ -598,11 +621,16 @@ gpu.func @no_scalar_load_unit_1D_vector_used_as_vector(
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>
+// A unit-size vector genuinely used as a vector keeps the vector paths - only a
+// read its consumers unwrap to a scalar becomes a scalar load.
+// LOAD-ND-LABEL: @no_scalar_load_unit_1D_vector_used_as_vector(
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd
+// LOAD-ND-SAME: -> vector<1x1xf32>
+// LOAD-ND: vector.shape_cast %[[VEC]] : vector<1x1xf32> to vector<1xf32>
+
+// LOAD-GATHER-LABEL: @no_scalar_load_unit_1D_vector_used_as_vector(
+// LOAD-GATHER: vector.step : vector<1xindex>
+// LOAD-GATHER: xegpu.load {{.*}} : i64, vector<1xindex>, vector<1xi1> -> vector<1xf32>
}
// -----
@@ -728,21 +756,31 @@ gpu.func @load_from_subview_1D(%source: memref<4096x4096xf16>, %off1: index, %of
gpu.return %0 : vector<8xf16>
}
-// CHECK-LABEL: @load_from_subview_1D(
-// 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] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// CHECK: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : 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]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
-// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf16>
+// LOAD-ND-LABEL: @load_from_subview_1D(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
+// LOAD-ND-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
+// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SUBVIEW]]
+// LOAD-ND-SAME: -> !xegpu.tensor_desc<1x8xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF2]], %[[OFF2]]]
+// LOAD-ND-SAME: -> vector<1x8xf16>
+// LOAD-ND: vector.shape_cast %[[VEC]] : vector<1x8xf16> to vector<8xf16>
+
+// LOAD-GATHER-LABEL: @load_from_subview_1D(
+// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
+// LOAD-GATHER-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
+// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// LOAD-GATHER: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
+// LOAD-GATHER: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
+// LOAD-GATHER: %[[STEP:.+]] = vector.step : vector<8xindex>
+// LOAD-GATHER: arith.muli {{.*}} : index
+// LOAD-GATHER: arith.addi %[[OFFSET]]{{.*}} : index
+// LOAD-GATHER: arith.addi {{.*}} : index
+// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
+// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
+// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
+// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf16>
}
// -----
More information about the Mlir-commits
mailing list