[Mlir-commits] [mlir] 8dba613 - [mlir][xegpu] Lower dynamic high-D nd load/store via base-pointer fold (#215711)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 22 14:35:40 PDT 2026
Author: Jianhui Li
Date: 2026-08-22T14:35:34-07:00
New Revision: 8dba613e379f9f5b4a58676ace073ba5927015aa
URL: https://github.com/llvm/llvm-project/commit/8dba613e379f9f5b4a58676ace073ba5927015aa
DIFF: https://github.com/llvm/llvm-project/commit/8dba613e379f9f5b4a58676ace073ba5927015aa.diff
LOG: [mlir][xegpu] Lower dynamic high-D nd load/store via base-pointer fold (#215711)
Reworks lowering of >2D (batched)
`xegpu.create_nd_tdesc`/`load_nd`/`store_nd`/ `prefetch_nd` (batched
GEMM, rank-4 flash-attention) to keep the full high-D memref as the
descriptor source and carry the leading (batch) offsets as a row offset
into a flattened 2D plane, instead of slicing a per-batch
`memref.subview` during blocking.
For example, the wg-level IR read a 4d vector out of a dynamic shaped
memref.
```mlir
%0 = vector.transfer_read %source[%i, %j, %k, %l], %c0
{in_bounds = [true, true, true, true]} : memref<?x?x8x16xf32>,
vector<2x4x8x16xf32>
```
This lowers with no subview — the full memref is the descriptor source
and all four offsets stay on the load:
```mlir
%desc = xegpu.create_nd_tdesc %source
: memref<?x?x8x16xf32> -> !xegpu.tensor_desc<2x4x8x16xf32, #...<boundary_check = false>>
%vec = xegpu.load_nd %desc[%i, %j, %k, %l] -> vector<2x4x8x16xf32>
```
After blocking/unrolling, the 2x4 batch dims become unit tiles and their
offsets ride on each load. At XeVM lowering the source is viewed as a
single flattened 2D plane: base_height is the product of all dims but
the innermost, so the 2D-block surface spans every batch plane at once,
and the batch offsets become a row offset into it:
offset_h = offset[R-2] + Σ_{d<R-2} offset[d] * (stride[d] / stride[R-2])
```mlir
%desc = xegpu.create_nd_tdesc %source : memref<?x?x8x16xf32> ->
!xegpu.tensor_desc<1x1x8x16xf32>
%t0 = xegpu.load_nd %desc[%i, %j, %k, %l] -> vector<1x1x8x16xf32>
%t1 = xegpu.load_nd %desc[%i, %j + 1, %k, %l] -> vector<1x1x8x16xf32>
%t2 = xegpu.load_nd %desc[%i, %j + 2, %k, %l] -> vector<1x1x8x16xf32>
%t3 = xegpu.load_nd %desc[%i, %j + 3, %k, %l] -> vector<1x1x8x16xf32>
%t4 = xegpu.load_nd %desc[%i + 1, %j, %k, %l] -> vector<1x1x8x16xf32>
%t5 = xegpu.load_nd %desc[%i + 1, %j + 1, %k, %l] -> vector<1x1x8x16xf32>
%t6 = xegpu.load_nd %desc[%i + 1, %j + 2, %k, %l] -> vector<1x1x8x16xf32>
%t7 = xegpu.load_nd %desc[%i + 1, %j + 3, %k, %l] -> vector<1x1x8x16xf32>
// %t0..%t7 reassembled into vector<2x4x8x16xf32> via vector.insert_strided_slice
```
This supersedes the subview approach (#201725), which could not produce a valid base for a dynamic-shape memref.
Collapsing the leading dims into one tall surface does mean base_height grows to the product of the leading dims, and can exceed the HW 2D-block surface height for large attention. That is the accepted cost of keeping base_ptr valid so the boundary check stays meaningful. Two further limitations are documented in the pass: each leading stride must be a whole number of rows, so a source with gaps between planes (stride[d] % stride[R-2] != 0) is not lowered; and plane boundaries are invisible to the boundary check, so a tile whose rows run past size[R-2] reads the next plane's rows instead of zero padding.
assited-by-claude
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
mlir/test/Dialect/XeGPU/invalid.mlir
mlir/test/Dialect/XeGPU/ops.mlir
mlir/test/Dialect/XeGPU/peephole-optimize.mlir
mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 17be6d1539414..d1d337699671b 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -74,20 +74,20 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
Arguments:
- `source`: an object representing (starting address/pointer of) a memory region.
It can be either a memref object, or simply a pointer represented by uint64_t type.
- For the case of dynamic memrefs or pointer, the shape and layout information of the
- memory region should be explicitly passed via `shape` and `strides` parameters.
+ For a memref source (static or dynamic shape) the shape and strides are taken from
+ the memref itself and must not be passed explicitly; only a pointer source requires
+ the `shape` and `strides` parameters.
- - `shape`: the shape information of the memory region pointed by the "source". It is
- typically encoded via the MemRefType of the source, e.g., memref<4096x4096xf16>.
- But if "source" is simply a pointer represented as uint64_t type, or a memref
- type without shape information e.g., memref<?x?xf16>, the shape information has
- to be explicitly passed via the "shape" and "const_shape" arguments.
+ - `shape`: the shape information of the memory region pointed by the "source". For a
+ memref source it is encoded via the MemRefType, e.g., memref<4096x4096xf16> or a
+ dynamic memref<?x?xf16> (whose dynamic dims are recovered at lowering time). Only
+ when "source" is a pointer represented as uint64_t type must the shape be passed
+ explicitly via the "shape" and "const_shape" arguments.
- `strides`: the strides of the memory region pointed by the "source". Similar to shape,
- it is typically encoded via the MemRefType of the source too. But if "source" is
- simply a pointer represented as uint64_t type, or a memref type without shape
- information e.g., memref<?x?xf16>, the strides information has to be explicitly
- passed via the "strides" and "const_strides" argument.
+ for a memref source it is encoded via the MemRefType, including a dynamic
+ memref<?x?xf16>. Only when "source" is a pointer represented as uint64_t type must
+ the strides be passed explicitly via the "strides" and "const_strides" arguments.
Results:
- `res`: nd tensor descriptor
@@ -98,14 +98,14 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
%1 = xegpu.create_nd_tdesc %0 : memref<1024x1024xf32> -> TensorDesc<8x16xf32>
```
- Example 2 (suppose the tensor shape inferred by the compiler is 8x16):
+ Example 2 (dynamic memref; shape/strides are inferred from the memref, not
+ passed explicitly):
```mlir
%0 = memref.alloc(%h, %w) : memref<?x?xf32>
- %c1 = arith.constant 1 : index
- %1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: memref<?x?xf32> -> TensorDesc<8x16xf32>
+ %1 = xegpu.create_nd_tdesc %0 : memref<?x?xf32> -> TensorDesc<8x16xf32>
```
- Example 3 (suppose the tensor shape inferred by the compiler is 8x16):
+ Example 3 (a pointer source must supply shape/strides explicitly):
```mlir
%0 = ... : ui64
%c1 = arith.constant 1 : index
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 0125dfc44196b..07a5cd59a39ba 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -14,6 +14,7 @@
#include "mlir/IR/OpDefinition.h"
#include "llvm/ADT/SetVector.h"
#include <functional>
+#include <optional>
namespace mlir {
@@ -209,8 +210,13 @@ template <typename T,
void setTemporaryLayout(const T &operandOrResult,
const DistributeLayoutAttr layout);
-/// Helper function to check if the layout is packed. Layout is packed if it is
-/// 2D and lane_data[0] != 1 (data packed from col dimension).
+/// Returns the innermost 2 entries of `vals` if it is at least 2D and all of
+/// its leading entries are unit; std::nullopt otherwise.
+std::optional<SmallVector<int64_t>>
+getInner2DIfUnitLeadingDims(ArrayRef<int64_t> vals);
+
+/// Helper function to check if the layout is packed. Layout is packed if
+/// lane_data[rank-2] != 1 (data packed from col dimension).
/// TODO: Move to target info.
bool requirePacked(const DistributeLayoutAttr layout);
@@ -218,6 +224,9 @@ bool requirePacked(const DistributeLayoutAttr layout);
bool requireTranspose(const DistributeLayoutAttr layout,
const uArch::uArch *uArch);
+/// Returns true if `type` has a static shape and static strides.
+bool hasStaticShapeAndStrides(MemRefType type);
+
// Check if dst shape is an expansion of src shape by inserting unit dimensions.
bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
SmallVector<int64_t> &expandedUnitDims);
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 8a45836426931..9863206f14fe1 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -126,56 +126,6 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
return success();
}
-static xegpu::CreateNdDescOp createNdDescriptor(PatternRewriter &rewriter,
- Location loc,
- xegpu::TensorDescType descType,
- TypedValue<MemRefType> src) {
- MemRefType srcTy = src.getType();
- assert(srcTy.isStrided() && "Expected strided memref type");
- auto [strides, offset] = srcTy.getStridesAndOffset();
- bool isStatic = true;
-
- // Memref is dynamic if any of its shape, offset or strides is dynamic.
- if (!srcTy.hasStaticShape())
- isStatic = false;
-
- if (!ShapedType::isStatic(offset))
- isStatic = false;
-
- for (auto stride : strides) {
- if (!ShapedType::isStatic(stride)) {
- isStatic = false;
- break;
- }
- }
-
- xegpu::CreateNdDescOp ndDesc;
- if (isStatic) {
- ndDesc = xegpu::CreateNdDescOp::create(rewriter, loc, descType, src);
- } else {
- // In case of ranked dynamic memref, instead of passing on the memref,
- // i64 base address, source's offset, shape and strides have to be
- // explicitly provided.
- auto meta = memref::ExtractStridedMetadataOp::create(rewriter, loc, src);
- auto baseAddrIndex = memref::ExtractAlignedPointerAsIndexOp::create(
- rewriter, loc, meta.getBaseBuffer());
- auto offset = meta.getOffset();
- auto elemByteSize = srcTy.getElementTypeBitWidth() / 8;
- auto offsetInBytes = arith::MulIOp::create(
- rewriter, loc, offset,
- arith::ConstantIndexOp::create(rewriter, loc, elemByteSize));
- auto adjustedBaseAddr = arith::AddIOp::create(
- rewriter, loc, baseAddrIndex.getResult(), offsetInBytes);
- auto adjustedAddrI64 = arith::IndexCastOp::create(
- rewriter, loc, rewriter.getI64Type(), adjustedBaseAddr);
- ndDesc = xegpu::CreateNdDescOp::create(
- rewriter, loc, descType, adjustedAddrI64,
- meta.getConstifiedMixedSizes(), meta.getConstifiedMixedStrides());
- }
-
- return ndDesc;
-}
-
// Adjusts the strides of a memref according to a given permutation map for
// vector operations.
//
@@ -646,7 +596,7 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
Operation *loadedOp =
@@ -748,7 +698,7 @@ struct TransferWriteLowering
xegpu::MemorySpace::Global);
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
auto storeOp = xegpu::StoreNdOp::create(
@@ -866,7 +816,7 @@ struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
vecTy.getShape(), vecTy.getElementType(), /*array_length=*/1,
boundaryCheck, xegpu::MemorySpace::Global);
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
auto loadNdOp =
xegpu::LoadNdOp::create(rewriter, loc, vecTy, ndDesc, indices,
@@ -911,7 +861,7 @@ struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
// By default, no specific caching policy is assigned.
xegpu::CachePolicyAttr hint = nullptr;
- xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
+ xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
auto storeNdOp =
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 76d6a73387aa1..f782be5502a8a 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -51,12 +51,19 @@ static constexpr int32_t executionSize{16};
// Offsets to individual fields of the 8xi32 layout nd tensor descriptor.
enum class NdTdescOffset : uint32_t {
- BasePtr = 0, // Base pointer (i64)
- BaseShapeW = 2, // Base shape width (i32)
- BaseShapeH = 3, // Base shape height (i32)
- BasePitch = 4, // Base pitch/stride of dim rank-2 (i32)
+ BasePtr = 0, // Base pointer (i64)
+ BaseShapeW = 2, // Base shape width (i32)
+ BaseShapeH = 3, // Base shape height (i32)
+ BasePitch = 4, // Base pitch/stride of dim rank-2 (i32)
+ LeadingStride0 = 5, // Row strides of the leading (batch) dims of a >2D
+ LeadingStride1 = 6, // descriptor (i32); added into offset_h by the load/store
+ LeadingStride2 = 7, // lowering. Left at 0 for 2D descriptors.
};
+// Spare payload slots above, and the resulting max lowerable descriptor rank.
+static constexpr int64_t maxNdTdescLeadingDims{3};
+static constexpr int64_t maxNdTdescRank{2 + maxNdTdescLeadingDims};
+
static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
switch (xeGpuMemspace) {
case xegpu::MemorySpace::Global:
@@ -205,6 +212,31 @@ translateStoreXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
// emulated
//
+//
+// High-D (>2D) nd descriptors are lowered by viewing the source as a single
+// flattened 2D plane: `base_height` is the product of all dims but the
+// innermost, so the 2D-block surface covers every leading (batch) plane at
+// once, and a batch position becomes a row offset into it. Leaving `base_ptr`
+// at the true base means an out-of-range batch index lands past `base_height`,
+// where the HW boundary check handles it, instead of aiming the surface at
+// unmapped memory. Encoding the leading strides as row counts (`stride[d] /
+// stride[R-2]`) also makes them dimensionless, so they survive element-type
+// repacking (e.g. the f16 -> i32 transpose repack) without a unit conversion.
+//
+// Limitations of the flattened-plane view:
+// 1. Each leading stride must be a whole number of rows. A source with gaps
+// between planes (`stride[d] % stride[R-2] != 0`) is not lowered. This can
+// only be checked when the strides are static; for dynamic strides the
+// divisibility is assumed.
+// 2. `base_height` grows to the product of the leading dims, so a source with
+// a large batch x head x sequence extent can exceed the HW 2D-block
+// surface height.
+// 3. Plane boundaries are invisible to the boundary check: a tile whose rows
+// run past `size[R-2]` reads the next plane's rows instead of the zero
+// padding a per-plane surface would return. This only matters when
+// `size[R-2]` is not a multiple of the tile height.
+//
+
class CreateNdDescToXeVMPattern
: public OpConversionPattern<xegpu::CreateNdDescOp> {
using OpConversionPattern::OpConversionPattern;
@@ -214,12 +246,91 @@ class CreateNdDescToXeVMPattern
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto source = op.getSource();
+
+ // Check all failure conditions before generating any IR, so nothing has to
+ // be rolled back.
+ int64_t rank = op.getType().getRank();
+ int64_t sourceRank;
+ auto memrefTy = dyn_cast<MemRefType>(source.getType());
+ if (memrefTy) {
+ if (!memrefTy.isStrided())
+ return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
+ sourceRank = memrefTy.getRank();
+ } else if (isa<IntegerType>(source.getType())) {
+ sourceRank = op.getMixedSizes().size();
+ } else {
+ return rewriter.notifyMatchFailure(op,
+ "Expected ranked Memref or integer.");
+ }
+ if (sourceRank != rank)
+ return rewriter.notifyMatchFailure(
+ op, "Expected descriptor rank to match source rank; subview the "
+ "source down to the descriptor rank.");
+ if (rank > maxNdTdescRank)
+ return rewriter.notifyMatchFailure(
+ op, "Batched nd descriptor supports at most " +
+ std::to_string(maxNdTdescLeadingDims) +
+ " leading dims (rank <= " + std::to_string(maxNdTdescRank) +
+ ").");
+ // Limitation 1 above; dynamic strides are assumed to divide evenly.
+ if (rank > 2) {
+ SmallVector<std::optional<int64_t>> constStrides(rank, std::nullopt);
+ if (memrefTy) {
+ SmallVector<int64_t> staticStrides;
+ int64_t staticOffset;
+ if (succeeded(
+ memrefTy.getStridesAndOffset(staticStrides, staticOffset)))
+ for (int64_t d = 0; d < rank; ++d)
+ if (!ShapedType::isDynamic(staticStrides[d]))
+ constStrides[d] = staticStrides[d];
+ } else {
+ SmallVector<OpFoldResult> mixed = op.getMixedStrides();
+ for (int64_t d = 0; d < rank; ++d)
+ constStrides[d] = getConstantIntValue(mixed[d]);
+ }
+ if (std::optional<int64_t> pitch = constStrides[rank - 2]) {
+ for (int64_t d = 0; d < rank - 2; ++d) {
+ std::optional<int64_t> leading = constStrides[d];
+ if (leading && (*pitch == 0 || *leading % *pitch != 0))
+ return rewriter.notifyMatchFailure(
+ op, "Expected each leading (batch) stride to be a multiple of "
+ "the row stride; the source has gaps between planes.");
+ }
+ }
+ }
+
+ Type payloadElemTy = rewriter.getI32Type();
+ Type i64Ty = rewriter.getI64Type();
+
+ // Access the adaptor only after the failure checks, so a bail-out leaves no
+ // materialization cast behind.
+ Value baseAddr = adaptor.getSource();
+ if (isa<IntegerType>(source.getType()) && baseAddr.getType() != i64Ty) {
+ // Pointer type may be i32. Cast to i64 if needed.
+ baseAddr = arith::ExtUIOp::create(rewriter, loc, i64Ty, baseAddr);
+ }
+ // 1D tensor descriptor is just the base address.
+ if (rank == 1) {
+ rewriter.replaceOp(op, baseAddr);
+ return success();
+ }
+
+ SmallVector<OpFoldResult> mixedSizes;
+ SmallVector<OpFoldResult> mixedStrides;
+ if (memrefTy && !xegpu::hasStaticShapeAndStrides(memrefTy)) {
+ auto meta =
+ memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
+ mixedSizes = meta.getConstifiedMixedSizes();
+ mixedStrides = meta.getConstifiedMixedStrides();
+ } else {
+ mixedSizes = op.getMixedSizes();
+ mixedStrides = op.getMixedStrides();
+ }
+
// Op is lowered to a code sequence that populates payload.
// Payload is a 8xi32 vector. Offset to individual fields are defined in
// NdTdescOffset enum.
- Type payloadElemTy = rewriter.getI32Type();
VectorType payloadTy = VectorType::get(8, payloadElemTy);
- Type i64Ty = rewriter.getI64Type();
// 4xi64 view is used for inserting the base pointer.
VectorType payloadI64Ty = VectorType::get(4, i64Ty);
// Initialize payload to zero.
@@ -227,38 +338,6 @@ class CreateNdDescToXeVMPattern
rewriter, loc,
DenseElementsAttr::get(payloadTy, IntegerAttr::get(payloadElemTy, 0)));
- Value baseAddr;
- Value baseShapeW;
- Value baseShapeH;
-
- // Source can be a memref or a pointer (ui64, ui32, i64 or i32).
- SmallVector<OpFoldResult> mixedSizes = op.getMixedSizes();
- SmallVector<OpFoldResult> mixedStrides = op.getMixedStrides();
- // Descriptor shape is expected to be 2D.
- int64_t rank = mixedSizes.size();
- auto sourceTy = source.getType();
- auto sourceMemrefTy = dyn_cast<MemRefType>(sourceTy);
- // If source is a memref, we need to extract the aligned pointer as index.
- // Pointer type is passed as i32 or i64 by type converter.
- if (sourceMemrefTy) {
- if (!sourceMemrefTy.hasRank()) {
- return rewriter.notifyMatchFailure(op, "Expected ranked Memref.");
- }
- // Access adaptor after failure check to avoid rolling back generated code
- // for materialization cast.
- baseAddr = adaptor.getSource();
- } else {
- baseAddr = adaptor.getSource();
- if (baseAddr.getType() != i64Ty) {
- // Pointer type may be i32. Cast to i64 if needed.
- baseAddr = arith::ExtUIOp::create(rewriter, loc, i64Ty, baseAddr);
- }
- }
- // 1D tensor descriptor is just the base address.
- if (rank == 1) {
- rewriter.replaceOp(op, baseAddr);
- return success();
- }
// Utility for creating offset values from op fold result.
auto createOffset = [&](SmallVector<OpFoldResult> &ofrVec,
unsigned idx) -> Value {
@@ -266,10 +345,15 @@ class CreateNdDescToXeVMPattern
val = getValueOrCreateCastToIndexLike(rewriter, loc, payloadElemTy, val);
return val;
};
- // For ND descriptors, the last 2 dimensions are the 2D tile (H, W).
- // Any leading dimensions are batch dims with associated strides.
- baseShapeW = createOffset(mixedSizes, rank - 1);
- baseShapeH = createOffset(mixedSizes, rank - 2);
+ // The descriptor's innermost 2 dims are the 2D tile (H, W).
+ Value baseShapeW = createOffset(mixedSizes, rank - 1);
+ // Height of the flattened plane: every leading (batch) plane is stacked
+ // into the surface, so the boundary check covers an out-of-range batch.
+ // For rank 2 this is just size[0].
+ Value baseShapeH = createOffset(mixedSizes, rank - 2);
+ for (int64_t d = 0; d < rank - 2; ++d)
+ baseShapeH = arith::MulIOp::create(rewriter, loc, baseShapeH,
+ createOffset(mixedSizes, d));
// Pitch is the stride of dim rank-2 (the row stride of the 2D tile).
Value basePitch = createOffset(mixedStrides, rank - 2);
// Populate payload.
@@ -288,6 +372,26 @@ class CreateNdDescToXeVMPattern
payload =
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
+ // Leading (batch) strides go into the spare payload slots as a number of
+ // rows; the load/store/prefetch lowering turns the batch offsets into a row
+ // offset with them. Row units keep them independent of the element type, so
+ // an element-type repack cannot put them out of step with the pitch.
+ for (int64_t d = 0; d < rank - 2; ++d) {
+ std::optional<int64_t> leading = getConstantIntValue(mixedStrides[d]);
+ std::optional<int64_t> pitch =
+ getConstantIntValue(mixedStrides[rank - 2]);
+ Value leadingRowStride;
+ if (leading && pitch && *pitch != 0) {
+ leadingRowStride = arith::ConstantIntOp::create(
+ rewriter, loc, payloadElemTy, *leading / *pitch);
+ } else {
+ leadingRowStride = arith::DivUIOp::create(
+ rewriter, loc, createOffset(mixedStrides, d), basePitch);
+ }
+ payload = vector::InsertOp::create(
+ rewriter, loc, leadingRowStride, payload,
+ static_cast<int>(NdTdescOffset::LeadingStride0) + d);
+ }
rewriter.replaceOp(op, payload);
return success();
}
@@ -313,6 +417,14 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
if (opOffsetsSize != tileRank)
return rewriter.notifyMatchFailure(
op, "Expected offset rank to match descriptor rank.");
+ if (tileRank > 2 && llvm::any_of(tdescTy.getShape().drop_back(2),
+ [](int64_t d) { return d != 1; }))
+ return rewriter.notifyMatchFailure(
+ op, "Expected leading (batch) descriptor dims to be unit.");
+ if (tileRank > maxNdTdescRank)
+ return rewriter.notifyMatchFailure(
+ op, "Expected descriptor rank <= " + std::to_string(maxNdTdescRank) +
+ ".");
auto elemType = tdescTy.getElementType();
auto elemBitSize = elemType.getIntOrFloatBitWidth();
bool isSubByte = elemBitSize < 8;
@@ -387,9 +499,6 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
Value basePitch = vector::ExtractOp::create(
rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BasePitch));
- // For rank > 2, leading (batch) dim offsets should be 0 after unrolling
- // (batch is baked into the base pointer via memref.subview during
- // blocking). Use only the last 2 offsets for the 2D block operation.
Value offsetW = getValueOrCreateConstantIntOp(rewriter, loc,
mixedOffsets[tileRank - 1]);
offsetW = getValueOrCreateCastToIndexLike(rewriter, loc,
@@ -398,6 +507,23 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
mixedOffsets[tileRank - 2]);
offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), offsetH);
+ // Turn the leading (batch) offsets into a row offset into the flattened
+ // plane, using the row-unit batch strides encoded at create time:
+ // offsetH += sum_d offset[d] * leadingRowStride[d]
+ // The base pointer stays at the true base, so an out-of-range batch index
+ // is caught by the HW boundary check instead of moving the surface to
+ // unmapped memory.
+ for (int64_t d = 0; d < tileRank - 2; ++d) {
+ Value off =
+ getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[d]);
+ off = getValueOrCreateCastToIndexLike(rewriter, loc,
+ rewriter.getI32Type(), off);
+ Value rowStride = vector::ExtractOp::create(
+ rewriter, loc, tdesc,
+ static_cast<int>(NdTdescOffset::LeadingStride0) + d);
+ Value term = arith::MulIOp::create(rewriter, loc, off, rowStride);
+ offsetH = arith::AddIOp::create(rewriter, loc, offsetH, term);
+ }
// Convert base pointer (i64) to LLVM pointer type.
Value basePtrLLVM =
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtr);
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 1bd5951c9f7f1..686e4342215b0 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -212,9 +212,6 @@ LogicalResult CreateMemDescOp::verify() {
void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
Type tdesc, TypedValue<MemRefType> source) {
- [[maybe_unused]] auto ty = source.getType();
- assert(ty.hasStaticShape() && "expecting a memref with static shape");
-
build(builder, state, tdesc, source, ValueRange({}) /* empty dynamic shape */,
ValueRange({}) /* empty dynamic strides */,
DenseI64ArrayAttr({}) /* empty const shape*/,
@@ -260,8 +257,8 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
}
LogicalResult CreateNdDescOp::verify() {
- size_t rank = getMixedSizes().size();
- bool invalidRank = rank != getMixedStrides().size();
+ auto srcMemrefTy = dyn_cast<MemRefType>(getSourceType());
+ size_t rank = srcMemrefTy ? srcMemrefTy.getRank() : getMixedSizes().size();
bool invalidElemTy = false;
// Memory space of created TensorDesc should match with the source.
@@ -280,18 +277,23 @@ LogicalResult CreateNdDescOp::verify() {
if (auto memrefTy = dyn_cast<MemRefType>(getSourceType()))
invalidElemTy |= memrefTy.getElementType() != getElementType();
+ bool hasExplicitShapeStrides =
+ !getShape().empty() || !getStrides().empty() ||
+ (getConstShapeAttr() && !getConstShapeAttr().empty()) ||
+ (getConstStridesAttr() && !getConstStridesAttr().empty());
+
if (llvm::isa<IntegerType>(getSourceType())) {
// strides and shape must present for integer source.
if (getMixedStrides().empty() || getMixedSizes().empty())
return emitOpError("expecting strides and shape to be present for "
"integer source.");
+ if (getMixedSizes().size() != getMixedStrides().size())
+ return emitOpError("Expecting the rank of shape and strides to match.");
+ } else if (srcMemrefTy && hasExplicitShapeStrides) {
+ return emitOpError("shape and strides should not be specified for a memref "
+ "source; they are inferred from the memref.");
}
- if (invalidRank)
- return emitOpError(
- "Expecting the rank of shape, strides, and source (if source "
- "is a memref) should match with each other.");
-
// check result TensorDesc rank
if (getType().getRank() > (int64_t)rank)
return emitOpError("Expecting the TensorDesc rank is not greater than the "
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
index 590d1804167a5..a3aaf01f356d5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
@@ -137,13 +137,16 @@ class OptimizeCreateNdDescOp : public OpRewritePattern<xegpu::CreateNdDescOp> {
tdescType.getBoundaryCheck(), tdescType.getMemorySpace(),
tdescType.getLayout());
- // The memory region is unchanged; pass through the existing shape/strides.
- // The general builder recognizes the static-memref case and drops the
- // redundant attributes.
- auto newOp = xegpu::CreateNdDescOp::create(
- rewriter, op.getLoc(), newTdescType, source, op.getMixedSizes(),
- op.getMixedStrides());
- rewriter.replaceOp(op, newOp.getResult());
+ Value newOp;
+ if (isa<MemRefType>(source.getType()))
+ newOp =
+ xegpu::CreateNdDescOp::create(rewriter, op.getLoc(), newTdescType,
+ cast<TypedValue<MemRefType>>(source));
+ else
+ newOp = xegpu::CreateNdDescOp::create(rewriter, op.getLoc(), newTdescType,
+ source, op.getMixedSizes(),
+ op.getMixedStrides());
+ rewriter.replaceOp(op, newOp);
return success();
}
};
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 1bf04f1f095a8..c64defc471de8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -50,10 +50,7 @@ getMaybeLaneData(xegpu::TensorDescType tdescType) {
auto layout = tdescType.getLayoutAttr();
if (!layout)
return std::nullopt;
- auto laneData = layout.getEffectiveLaneDataAsInt();
- if (laneData.size() != 2)
- return std::nullopt;
- return laneData;
+ return xegpu::getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt());
}
/// Get the 2D lane layout from a tensor desc type if it exists.
@@ -62,10 +59,8 @@ getMaybeLaneLayout(xegpu::TensorDescType tdescType) {
auto layout = tdescType.getLayoutAttr();
if (!layout)
return std::nullopt;
- auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
- if (laneLayout.size() != 2)
- return std::nullopt;
- return laneLayout;
+ return xegpu::getInner2DIfUnitLeadingDims(
+ layout.getEffectiveLaneLayoutAsInt());
}
/// A layout can be optimized if its lane layout is transposed (lane[0] != 1 &&
@@ -139,15 +134,22 @@ tryOptimize(xegpu::TensorDescType tdescType,
if (counts.size() != 1 || counts[0] != 1)
return tdescType;
int arrayLen = counts[0];
- int supportedHeight =
- xegpu::getLargestDivisor(static_cast<int>(requiredShape[0]), heights);
- int supportedWidth =
- xegpu::getLargestDivisor(static_cast<int>(requiredShape[1]), widths);
+ // The transpose acts on the innermost 2 dims; any leading dims are unit
+ // batch.
+ int64_t rank = requiredShape.size();
+ int supportedHeight = xegpu::getLargestDivisor(
+ static_cast<int>(requiredShape[rank - 2]), heights);
+ int supportedWidth = xegpu::getLargestDivisor(
+ static_cast<int>(requiredShape[rank - 1]), widths);
// If no supported height or width found, return the original type.
if (supportedHeight == -1 || supportedWidth == -1)
return tdescType;
- SmallVector<int64_t> supportedShape = {supportedHeight, supportedWidth};
+ // Preserve leading (unit) batch dims; only the inner 2 dims are reshaped.
+ SmallVector<int64_t> supportedShape(requiredShape.begin(),
+ requiredShape.end() - 2);
+ supportedShape.push_back(supportedHeight);
+ supportedShape.push_back(supportedWidth);
auto ctx = tdescType.getContext();
auto origLayout = tdescType.getLayoutAttr();
auto laneLayoutI64 = origLayout.getEffectiveLaneLayoutAsInt();
@@ -156,7 +158,9 @@ tryOptimize(xegpu::TensorDescType tdescType,
xegpu::LayoutAttr newLayout = xegpu::LayoutAttr::get(
ctx, /*lane_layout=*/DenseI32ArrayAttr::get(ctx, laneLayoutI32),
- /*lane_data=*/DenseI32ArrayAttr::get(ctx, {1, 1}),
+ /*lane_data=*/
+ DenseI32ArrayAttr::get(ctx,
+ SmallVector<int32_t>(laneLayoutI32.size(), 1)),
/*order=*/origLayout.getOrder());
// Array length can not be larger than 1 for transpose case.
@@ -201,6 +205,7 @@ static Value generateLoads(ConversionPatternRewriter &rewriter,
xegpu::LoadNdOp origLoadOp) {
Location loc = data.getLoc();
assert(offsets.size() >= 2 && "Expecting at least 2 offsets for 2D LoadNdOp");
+ int64_t rank = data.getType().getRank();
Value offsetDim0 = convertToValue(rewriter, loc, offsets[offsets.size() - 2]);
Value offsetDim1 = convertToValue(rewriter, loc, offsets[offsets.size() - 1]);
SmallVector<int64_t> supportedShape(newTensorDesc.getType().getShape());
@@ -209,10 +214,15 @@ static Value generateLoads(ConversionPatternRewriter &rewriter,
auto shapeRatio = computeShapeRatio(data.getType().getShape(),
supportedShape)
.value(); // `ratio` must be defined if we reach here.
- for (int64_t h = 0; h < shapeRatio[0]; ++h) {
- for (int64_t w = 0; w < shapeRatio[1]; ++w) {
- int64_t localOffsetDim0 = h * supportedShape[0];
- int64_t localOffsetDim1 = w * supportedShape[1];
+ // The loop below only walks the last 2 entries of `shapeRatio`; leading
+ // (batch) dims are unit, so their ratio is 1 and their offsets pass through
+ // unchanged.
+ int64_t suppDim0 = supportedShape[rank - 2];
+ int64_t suppDim1 = supportedShape[rank - 1];
+ for (int64_t h = 0; h < shapeRatio[rank - 2]; ++h) {
+ for (int64_t w = 0; w < shapeRatio[rank - 1]; ++w) {
+ int64_t localOffsetDim0 = h * suppDim0;
+ int64_t localOffsetDim1 = w * suppDim1;
Value loadOffsetX = arith::AddIOp::create(
rewriter, loc, offsetDim0,
arith::ConstantIndexOp::create(rewriter, loc, localOffsetDim0)
@@ -221,21 +231,28 @@ static Value generateLoads(ConversionPatternRewriter &rewriter,
rewriter, loc, offsetDim1,
arith::ConstantIndexOp::create(rewriter, loc, localOffsetDim1)
.getResult());
+ // Keep the leading (batch) offsets; replace only the inner 2.
+ SmallVector<OpFoldResult> loadOffsets(offsets.begin(), offsets.end());
+ loadOffsets[loadOffsets.size() - 2] = loadOffsetX;
+ loadOffsets[loadOffsets.size() - 1] = loadOffsetY;
auto loadOp = xegpu::LoadNdOp::create(
rewriter, loc,
VectorType::get(supportedShape, data.getType().getElementType()),
- newTensorDesc, ArrayRef<OpFoldResult>{loadOffsetX, loadOffsetY},
- origLoadOp.getPackedAttr(), origLoadOp.getTransposeAttr(),
- origLoadOp.getL1HintAttr(), origLoadOp.getL2HintAttr(),
- origLoadOp.getL3HintAttr(), origLoadOp.getLayoutAttr());
+ newTensorDesc, loadOffsets, origLoadOp.getPackedAttr(),
+ origLoadOp.getTransposeAttr(), origLoadOp.getL1HintAttr(),
+ origLoadOp.getL2HintAttr(), origLoadOp.getL3HintAttr(),
+ origLoadOp.getLayoutAttr());
// Set the layout for the loadOp.
auto layoutAttr = newTensorDesc.getType().getLayoutAttr();
loadOp.setAnchorLayout(layoutAttr);
- // Insert the loaded block into the right position in data.
+ // Insert the loaded block into the right position in data (leading dims
+ // at 0, inner 2 dims at the local tile offset).
+ SmallVector<int64_t> insertPos(rank, 0);
+ insertPos[rank - 2] = localOffsetDim0;
+ insertPos[rank - 1] = localOffsetDim1;
+ SmallVector<int64_t> insertStrides(rank, 1);
auto insertOp = vector::InsertStridedSliceOp::create(
- rewriter, loc, loadOp.getResult(), data,
- ArrayRef<int64_t>{localOffsetDim0, localOffsetDim1},
- ArrayRef<int64_t>{1, 1});
+ rewriter, loc, loadOp.getResult(), data, insertPos, insertStrides);
// InsertOp must have the same layout as newTensorDesc.
xegpu::setTemporaryLayout(insertOp->getOpResult(0), layoutAttr);
data = insertOp.getResult();
@@ -269,48 +286,68 @@ class XeGPUCreateNdDescOpPattern final
auto convertType = tryOptimize(tdescTy, targetuArch);
if (convertType == tdescTy)
return failure();
- auto strides = createNdOp.getMixedStrides();
- auto maybeConstInnerStride = getConstantIntValue(strides.back());
+ Location loc = createNdOp.getLoc();
+ Value source = createNdOp.getSource();
+ auto memrefType = dyn_cast<MemRefType>(source.getType());
+
+ bool dynamicMemref =
+ memrefType && !xegpu::hasStaticShapeAndStrides(memrefType);
+ SmallVector<OpFoldResult> mixedSizes;
+ SmallVector<OpFoldResult> mixedStrides;
+ memref::ExtractStridedMetadataOp meta;
+ if (dynamicMemref) {
+ meta = memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
+ mixedSizes = meta.getConstifiedMixedSizes();
+ mixedStrides = meta.getConstifiedMixedStrides();
+ } else {
+ mixedSizes = createNdOp.getMixedSizes();
+ mixedStrides = createNdOp.getMixedStrides();
+ }
+
+ auto maybeConstInnerStride = getConstantIntValue(mixedStrides.back());
// Only row-major memrefs are expected for now.
if (!maybeConstInnerStride || *maybeConstInnerStride != 1)
return rewriter.notifyMatchFailure(
createNdOp, "Expecting row-major memref for transpose optimization.");
- Value source = createNdOp.getSource();
auto optionalLaneData = getMaybeLaneData(tdescTy);
assert(optionalLaneData && "Expected 2D lane data");
auto laneData = optionalLaneData.value();
int64_t innerLaneData = laneData[1];
- auto memrefType = dyn_cast<MemRefType>(source.getType());
// Inner dimension of the shape must be adjusted based on innerLaneData.
- SmallVector<OpFoldResult> modifiedShape(createNdOp.getMixedSizes());
+ SmallVector<OpFoldResult> modifiedShape(mixedSizes);
modifiedShape.back() = divideByConstant(
- rewriter, createNdOp.getLoc(),
- convertToValue(rewriter, createNdOp.getLoc(), modifiedShape.back()),
+ rewriter, loc, convertToValue(rewriter, loc, modifiedShape.back()),
innerLaneData);
- // Similarly, second to last stride must be adjusted.
- assert(strides.size() >= 2 &&
+ // Repacking to a wider element rescales every stride but the innermost.
+ assert(mixedStrides.size() >= 2 &&
"Expected at least 2 strides for CreateNdDescOp");
- SmallVector<OpFoldResult> modifiedStrides(strides);
- modifiedStrides[modifiedStrides.size() - 2] = divideByConstant(
- rewriter, createNdOp.getLoc(),
- convertToValue(rewriter, createNdOp.getLoc(),
- modifiedStrides[modifiedStrides.size() - 2]),
- innerLaneData);
-
- // If the source is a static memref, we need to extract the pointer to
- // base address.
- if (memrefType && memrefType.hasStaticShape()) {
- auto extractOp = memref::ExtractAlignedPointerAsIndexOp::create(
- rewriter, createNdOp.getLoc(), source);
- source = arith::IndexCastOp::create(rewriter, createNdOp.getLoc(),
- rewriter.getI64Type(),
- extractOp.getResult())
- .getResult();
+ SmallVector<OpFoldResult> modifiedStrides(mixedStrides);
+ for (size_t i = 0; i + 1 < modifiedStrides.size(); ++i)
+ modifiedStrides[i] = divideByConstant(
+ rewriter, loc, convertToValue(rewriter, loc, modifiedStrides[i]),
+ innerLaneData);
+
+ if (memrefType) {
+ Value baseIdx;
+ if (dynamicMemref) {
+ // Base = aligned base pointer + structural offset (in bytes).
+ Value alignedPtr = memref::ExtractAlignedPointerAsIndexOp::create(
+ rewriter, loc, meta.getBaseBuffer());
+ Value elemBytes = arith::ConstantIndexOp::create(
+ rewriter, loc, memrefType.getElementTypeBitWidth() / 8);
+ Value offBytes =
+ arith::MulIOp::create(rewriter, loc, meta.getOffset(), elemBytes);
+ baseIdx = arith::AddIOp::create(rewriter, loc, alignedPtr, offBytes);
+ } else {
+ baseIdx = memref::ExtractAlignedPointerAsIndexOp::create(rewriter, loc,
+ source);
+ }
+ source = arith::IndexCastOp::create(rewriter, loc, rewriter.getI64Type(),
+ baseIdx);
}
// Create a new CreateNdDescOp with the modified shape and converted type.
auto newCreateNdDescOp = xegpu::CreateNdDescOp::create(
- rewriter, createNdOp.getLoc(), convertType, source, modifiedShape,
- modifiedStrides);
+ rewriter, loc, convertType, source, modifiedShape, modifiedStrides);
rewriter.replaceOp(createNdOp, newCreateNdDescOp.getResult());
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 14061ab24ffe0..5914e23f1f11a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -120,54 +120,6 @@ struct UnrollPattern : public OpRewritePattern<SourceOp> {
return SmallVector<Value>();
}
- /// Helper for the rank > 2 case shared by Load/Store/PrefetchNd unroll
- /// patterns. The matching CreateNdDesc unroll pattern produces one tdesc
- /// per batch tile (the batch offset is baked into its base pointer via
- /// memref.subview), so here we only need to iterate the inner 2D offsets
- /// for each batch tdesc.
- ///
- /// Packs `srcTdesc` into one tdesc per batch tile, then iterates the inner
- /// 2D tile offsets for each batch tdesc and invokes `createOp` with
- /// (batchTdesc, fullOffsets), where fullOffsets is `batchRank` zeros
- /// followed by the inner offsets. Returns the values produced by createOp,
- /// flattened across (batch, inner) iteration order.
- SmallVector<Value> unrollNdBatch(
- Value srcTdesc, xegpu::TensorDescType tdescTy,
- ArrayRef<int64_t> targetShape, ArrayRef<OpFoldResult> mixedOffsets,
- int64_t batchRank,
- llvm::function_ref<Value(Value, SmallVector<OpFoldResult>)> createOp,
- Location loc, PatternRewriter &rewriter) const {
- ArrayRef<int64_t> shape = tdescTy.getShape();
- SmallVector<int64_t> innerShape(shape.begin() + batchRank, shape.end());
- SmallVector<int64_t> innerTarget(targetShape.begin() + batchRank,
- targetShape.end());
-
- SmallVector<Type> batchTdescTypes = getUnrolledTypes(tdescTy, targetShape);
- SmallVector<Value> batchTdescs =
- pack(srcTdesc, batchTdescTypes, targetShape, loc, rewriter);
-
- auto innerTdescTy = xegpu::TensorDescType::get(
- tdescTy.getContext(), innerShape, tdescTy.getElementType(),
- tdescTy.getEncoding(), /*layout=*/nullptr);
-
- SmallVector<OpFoldResult> innerOffsets(mixedOffsets.begin() + batchRank,
- mixedOffsets.end());
-
- SmallVector<Value> newOps;
- for (Value batchTdesc : batchTdescs) {
- auto wrappedCreate = [&](SmallVector<OpFoldResult> offsets) -> Value {
- SmallVector<OpFoldResult> fullOffsets(batchRank,
- rewriter.getIndexAttr(0));
- fullOffsets.append(offsets.begin(), offsets.end());
- return createOp(batchTdesc, fullOffsets);
- };
- auto perBatch = unrollByTile(innerOffsets, innerTdescTy, innerTarget,
- wrappedCreate, loc, rewriter);
- newOps.append(perBatch.begin(), perBatch.end());
- }
- return newOps;
- }
-
/// Helper to pack operands for DPAS-like operations with early return if
/// no unrolling is needed.
SmallVector<Value> packOperandForDpas(Value operand,
@@ -248,63 +200,27 @@ struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
if (!targetShape)
return failure();
- int64_t rank = tdescTy.getRank();
- int64_t batchRank = rank - 2;
-
- // For rank <= 2 or non-memref source: existing single-tdesc behavior.
- if (batchRank <= 0 || !isa<MemRefType>(op.getSourceType())) {
- SmallVector<Value> newOps;
- auto newTdescTy = getUnrolledTypes(tdescTy, *targetShape)[0];
- auto newOp = xegpu::CreateNdDescOp::create(
- rewriter, loc, newTdescTy, op.getSource(), op.getMixedSizes(),
- op.getMixedStrides());
- newOps.push_back(newOp);
- Value castOp = unpack(newOps, tdescTy, *targetShape, loc, rewriter);
- rewriter.replaceOp(op, castOp);
- return success();
- }
-
- // For rank > 2 with memref source: create one tdesc per batch tile via
- // memref.subview. Each subview slices the batch dimensions, so the
- // resulting tdesc has the batch offset baked into its base pointer.
- // The inner dimensions remain full-size for reuse across multiple
- // load/store operations with
diff erent offsets.
- ArrayRef<int64_t> shape = tdescTy.getShape();
- SmallVector<int64_t> batchBlockSize(targetShape->begin(),
- targetShape->begin() + batchRank);
- batchBlockSize.append(shape.begin() + batchRank, shape.end());
-
- auto newTdescTy =
- cast<xegpu::TensorDescType>(getUnrolledTypes(tdescTy, *targetShape)[0]);
+ // Keep the high-D source; only the tile shape shrinks.
+ Value src = op.getSource();
+ auto makeCreateNd = [&](Type tdesc) -> Value {
+ auto ndTy = cast<xegpu::TensorDescType>(tdesc);
+ if (isa<MemRefType>(src.getType()))
+ return xegpu::CreateNdDescOp::create(rewriter, loc, ndTy,
+ cast<TypedValue<MemRefType>>(src));
+ return xegpu::CreateNdDescOp::create(
+ rewriter, loc, ndTy, src, op.getMixedSizes(), op.getMixedStrides());
+ };
+ SmallVector<Type> newTdescTys = getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> newOps;
- for (SmallVector<int64_t> batchOffsets :
- StaticTileOffsetRange(shape, batchBlockSize)) {
- // Build memref.subview operands. The subview slices contiguously along
- // each batch dimension (no gaps), so the subview's element stride is 1
- // for every dim. This is unrelated to the source memref's strides, which
- // describe the layout of the original buffer and are propagated by the
- // SubViewOp builder onto the resulting memref type.
- SmallVector<OpFoldResult> subviewOffsets;
- for (int64_t off : batchOffsets)
- subviewOffsets.push_back(rewriter.getIndexAttr(off));
-
- SmallVector<OpFoldResult> subviewSizes;
- for (int64_t d : batchBlockSize)
- subviewSizes.push_back(rewriter.getIndexAttr(d));
-
- SmallVector<OpFoldResult> subviewStrides(rank, rewriter.getIndexAttr(1));
-
- auto subview = memref::SubViewOp::create(rewriter, loc, op.getSource(),
- subviewOffsets, subviewSizes,
- subviewStrides);
-
- auto newOp = xegpu::CreateNdDescOp::create(
- rewriter, loc, newTdescTy,
- cast<TypedValue<MemRefType>>(subview.getResult()));
- newOps.push_back(newOp);
+ if (tdescTy.getRank() <= 2) {
+ // 2D: one tdesc, broadcast across tiles by pack/unpack.
+ newOps.push_back(makeCreateNd(newTdescTys[0]));
+ } else {
+ // >2D: one tdesc per tile, so the source count matches the pack count.
+ for (Type t : newTdescTys)
+ newOps.push_back(makeCreateNd(t));
}
-
Value castOp = unpack(newOps, tdescTy, *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
@@ -326,37 +242,20 @@ struct UnrollPrefetchNdOp : public UnrollPattern<xegpu::PrefetchNdOp> {
if (layout)
layout = layout.dropInstData();
- int64_t rank = tdescTy.getRank();
- int64_t batchRank = rank - 2;
-
- if (batchRank <= 0) {
- SmallVector<Type> convertedTdescTypes =
- getUnrolledTypes(tdescTy, *targetShape);
- SmallVector<Value> convertedTdesc = pack(
- op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
-
- auto createPrefetch = [&](SmallVector<OpFoldResult> offsets) -> Value {
- xegpu::PrefetchNdOp::create(rewriter, loc, convertedTdesc[0], offsets,
- op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), layout);
- return nullptr;
- };
- unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape, createPrefetch,
- loc, rewriter);
- } else {
- // Rank > 2: batch tdescs cover [batchTarget..., innerShape...].
- // Each batch tdesc is reused for multiple inner prefetches via offsets.
- auto createPrefetch =
- [&](Value tdesc, SmallVector<OpFoldResult> fullOffsets) -> Value {
- xegpu::PrefetchNdOp::create(rewriter, loc, tdesc, fullOffsets,
- op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), layout);
- return nullptr;
- };
- this->unrollNdBatch(op.getTensorDesc(), tdescTy, *targetShape,
- op.getMixedOffsets(), batchRank, createPrefetch, loc,
- rewriter);
- }
+ // Batch (leading) dims unroll to unit tiles; one tdesc serves all.
+ SmallVector<Type> convertedTdescTypes =
+ getUnrolledTypes(tdescTy, *targetShape);
+ SmallVector<Value> convertedTdesc = pack(
+ op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
+
+ auto createPrefetch = [&](SmallVector<OpFoldResult> offsets) -> Value {
+ xegpu::PrefetchNdOp::create(rewriter, loc, convertedTdesc[0], offsets,
+ op.getL1HintAttr(), op.getL2HintAttr(),
+ op.getL3HintAttr(), layout);
+ return nullptr;
+ };
+ unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape, createPrefetch,
+ loc, rewriter);
rewriter.eraseOp(op);
return success();
@@ -383,39 +282,22 @@ struct UnrollLoadNdOp : public UnrollPattern<xegpu::LoadNdOp> {
Type elemTy = tdescTy.getElementType();
VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);
- int64_t rank = tdescTy.getRank();
- int64_t batchRank = rank - 2;
SmallVector<Value> newOps;
- if (batchRank <= 0) {
- // Rank <= 2: original behavior with single tdesc.
- SmallVector<Type> convertedTdescTypes =
- getUnrolledTypes(tdescTy, *targetShape);
- SmallVector<Value> convertedTdescs = pack(
- op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
-
- auto createLoad = [&](SmallVector<OpFoldResult> offsets) -> Value {
- return xegpu::LoadNdOp::create(
- rewriter, loc, newValueTy, convertedTdescs[0], offsets,
- op.getPackedAttr(), op.getTransposeAttr(), op.getL1HintAttr(),
- op.getL2HintAttr(), op.getL3HintAttr(), layout);
- };
- newOps = unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape,
- createLoad, loc, rewriter);
- } else {
- // Rank > 2: batch tdescs cover [batchTarget..., innerShape...].
- // Each batch tdesc is reused for multiple inner loads via offsets.
- auto createLoad = [&](Value tdesc,
- SmallVector<OpFoldResult> fullOffsets) -> Value {
- return xegpu::LoadNdOp::create(
- rewriter, loc, newValueTy, tdesc, fullOffsets, op.getPackedAttr(),
- op.getTransposeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), layout);
- };
- newOps = this->unrollNdBatch(op.getTensorDesc(), tdescTy, *targetShape,
- op.getMixedOffsets(), batchRank, createLoad,
- loc, rewriter);
- }
+ // Batch (leading) dims unroll to unit tiles; one tdesc serves all.
+ SmallVector<Type> convertedTdescTypes =
+ getUnrolledTypes(tdescTy, *targetShape);
+ SmallVector<Value> convertedTdescs = pack(
+ op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
+
+ auto createLoad = [&](SmallVector<OpFoldResult> offsets) -> Value {
+ return xegpu::LoadNdOp::create(
+ rewriter, loc, newValueTy, convertedTdescs[0], offsets,
+ op.getPackedAttr(), op.getTransposeAttr(), op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(), layout);
+ };
+ newOps = unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape,
+ createLoad, loc, rewriter);
Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
rewriter.replaceOp(op, castOp);
@@ -445,42 +327,24 @@ struct UnrollStoreNdOp : public UnrollPattern<xegpu::StoreNdOp> {
SmallVector<Value> convertedValues =
pack(op.getValue(), convertedValTypes, *targetShape, loc, rewriter);
- int64_t rank = tdescTy.getRank();
- int64_t batchRank = rank - 2;
size_t valueIndex = 0;
- if (batchRank <= 0) {
- SmallVector<Type> convertedTdescTypes =
- getUnrolledTypes(tdescTy, *targetShape);
- SmallVector<Value> convertedTdescs = pack(
- op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
-
- auto createStore = [&](SmallVector<OpFoldResult> offsets) {
- xegpu::StoreNdOp::create(rewriter, loc, convertedValues[valueIndex++],
- convertedTdescs[0], offsets,
- op.getL1HintAttr(), op.getL2HintAttr(),
- op.getL3HintAttr(), layout);
- return (Value) nullptr;
- };
- unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape, createStore,
- loc, rewriter);
- } else {
- // Rank > 2: batch tdescs cover [batchTarget..., innerShape...].
- // Each batch tdesc is reused for multiple inner stores via offsets.
- // valueIndex advances across (batch, inner) iterations in the same
- // order unrollNdBatch invokes the callback, so it stays in sync with
- // the pre-packed convertedValues.
- auto createStore = [&](Value tdesc,
- SmallVector<OpFoldResult> fullOffsets) -> Value {
- xegpu::StoreNdOp::create(
- rewriter, loc, convertedValues[valueIndex++], tdesc, fullOffsets,
- op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), layout);
- return nullptr;
- };
- this->unrollNdBatch(op.getTensorDesc(), tdescTy, *targetShape,
- op.getMixedOffsets(), batchRank, createStore, loc,
- rewriter);
- }
+ // Batch (leading) dims unroll to unit tiles like any other dim. valueIndex
+ // advances in unrollByTile's tile order, staying in sync with the packed
+ // values.
+ SmallVector<Type> convertedTdescTypes =
+ getUnrolledTypes(tdescTy, *targetShape);
+ SmallVector<Value> convertedTdescs = pack(
+ op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
+
+ auto createStore = [&](SmallVector<OpFoldResult> offsets) {
+ xegpu::StoreNdOp::create(rewriter, loc, convertedValues[valueIndex++],
+ convertedTdescs[0], offsets, op.getL1HintAttr(),
+ op.getL2HintAttr(), op.getL3HintAttr(), layout);
+ return (Value) nullptr;
+ };
+ unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape, createStore, loc,
+ rewriter);
rewriter.eraseOp(op);
return success();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index a24e9b2fd7e0f..0e8a386fb08b6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -203,10 +203,14 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
xegpu::TensorDescType::get(ctx, sgShape, elemTy, tdescTy.getEncoding(),
layout.dropSgLayoutAndData());
+ Value src = op.getSource();
SmallVector<Value> newCreateNdOps(count);
- std::generate(newCreateNdOps.begin(), newCreateNdOps.end(), [&]() {
- return xegpu::CreateNdDescOp::create(rewriter, loc, newTdescTy,
- op.getSource(), op.getMixedSizes(),
+ std::generate(newCreateNdOps.begin(), newCreateNdOps.end(), [&]() -> Value {
+ if (isa<MemRefType>(src.getType()))
+ return xegpu::CreateNdDescOp::create(rewriter, loc, newTdescTy,
+ cast<TypedValue<MemRefType>>(src));
+ return xegpu::CreateNdDescOp::create(rewriter, loc, newTdescTy, src,
+ op.getMixedSizes(),
op.getMixedStrides());
});
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 07dcb438ca553..76269cf193d13 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -773,13 +773,21 @@ template int
xegpu::getLargestDivisor<unsigned>(unsigned dim, ArrayRef<unsigned> candidates,
ArrayRef<unsigned> candidateMultiples);
+std::optional<SmallVector<int64_t>>
+xegpu::getInner2DIfUnitLeadingDims(ArrayRef<int64_t> vals) {
+ if (vals.size() < 2)
+ return std::nullopt;
+ if (llvm::any_of(vals.drop_back(2), [](int64_t v) { return v != 1; }))
+ return std::nullopt;
+ return SmallVector<int64_t>(vals.take_back(2));
+}
+
bool xegpu::requirePacked(const xegpu::DistributeLayoutAttr layout) {
if (!layout)
return false;
- auto laneData = layout.getEffectiveLaneDataAsInt();
- if (laneData.size() != 2)
- return false;
- return laneData[0] != 1;
+ auto laneData =
+ getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt());
+ return laneData && (*laneData)[0] != 1;
}
bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
@@ -790,10 +798,19 @@ bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
return false;
if (!layout)
return false;
- auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
- if (laneLayout.size() != 2)
+ auto laneLayout =
+ getInner2DIfUnitLeadingDims(layout.getEffectiveLaneLayoutAsInt());
+ return laneLayout && (*laneLayout)[0] == uArch->getSubgroupSize() &&
+ (*laneLayout)[1] == 1;
+}
+
+bool xegpu::hasStaticShapeAndStrides(MemRefType type) {
+ if (!type.hasStaticShape())
return false;
- return laneLayout[0] == uArch->getSubgroupSize() && laneLayout[1] == 1;
+ SmallVector<int64_t> strides;
+ int64_t offset;
+ return succeeded(type.getStridesAndOffset(strides, offset)) &&
+ llvm::none_of(strides, ShapedType::isDynamic);
}
// Check if dst shape is an expansion of src shape by inserting unit dimensions.
diff --git a/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
index c77efa03f3483..d037fd30dc3c0 100644
--- a/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir
@@ -9,18 +9,8 @@ func.func @load_1D_vector(%source: memref<8x16x32xf32>, %offset: index) -> vecto
// CHECK-LABEL: @load_1D_vector(
// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
-// CHECK: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]
-// CHECK: %[[BASE_BUFFER:.+]], %[[OFFSET1:.+]], %[[SIZES:.+]], %[[STRIDES:.+]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// CHECK-SAME: : memref<32xf32, strided<[1], offset: ?>> -> memref<f32>, index, index, index
-// CHECK: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// CHECK-SAME: : memref<f32> -> index
-// CHECK: %[[MUL:.+]] = arith.muli %[[OFFSET1]], %[[ELEM_BYTES]] : index
-// CHECK: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// CHECK: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [32],
-// CHECK-SAME: strides : [1] : i64 -> !xegpu.tensor_desc<8xf32,
-// CHECK-SAME: boundary_check = false
+// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0] [1, 1, 32] [1, 1, 1] : memref<8x16x32xf32> to memref<32xf32, strided<[1], offset: ?>>
+// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<32xf32, strided<[1], offset: ?>> -> !xegpu.tensor_desc<8xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]]]{{.*}}-> vector<8xf32>
// CHECK: return %[[VEC]]
@@ -36,16 +26,8 @@ func.func @load_2D_vector(%source: memref<8x16x32xf32>,
// CHECK-LABEL: @load_2D_vector(
// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
-// CHECK: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]
-// CHECK: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// CHECK-SAME: : memref<f32> -> index
-// CHECK: %[[MUL:.+]] = arith.muli %[[OFF1]], %[[ELEM_BYTES]] : index
-// CHECK: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// CHECK: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [16, 32],
-// CHECK-SAME: strides : [32, 1] : i64 -> !xegpu.tensor_desc<8x16xf32>
+// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0] [1, 16, 32] [1, 1, 1] : memref<8x16x32xf32> to memref<16x32xf32, strided<[32, 1], offset: ?>>
+// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<8x16xf32>
// CHECK: return %[[VEC]]
@@ -61,15 +43,9 @@ func.func @load_dynamic_source(%source: memref<?x?x?xf32>,
// CHECK-LABEL: @load_dynamic_source(
// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,
// CHECK-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// CHECK: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]
-// CHECK: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// CHECK: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]] : memref<f32> -> index
-// CHECK: %[[MUL:.+]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// CHECK: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// CHECK: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [%[[SIZES]]#0, %[[SIZES]]#1],
-// CHECK-SAME: strides : [%[[STRIDES]]#0, 1] : i64 -> !xegpu.tensor_desc<8x16xf32>
+// CHECK: %{{.*}}, %{{.*}}, %[[SIZES:.+]]:3, %{{.+}}:3 = memref.extract_strided_metadata %[[SRC]]
+// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0] [1, %[[SIZES]]#1, %[[SIZES]]#2] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<?x?xf32, strided<[?, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF1]], %[[OFF2]]]{{.*}}-> vector<8x16xf32>
// CHECK: return %[[VEC]]
diff --git a/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
index 8ff2e6ee7d13c..dc9cabe0dbb18 100644
--- a/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir
@@ -11,18 +11,8 @@ func.func @store_1D_vector(%vec: vector<8xf32>,
// CHECK-SAME: %[[VEC:.+]]: vector<8xf32>,
// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
-// CHECK: %[[ELEM_BYTES:.*]] = arith.constant 4 : index
-// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]
-// CHECK: %[[BASE_BUFFER:.+]], %[[OFFSET1:.+]], %[[SIZES:.+]], %[[STRIDES:.+]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// CHECK-SAME: : memref<32xf32, strided<[1], offset: ?>> -> memref<f32>, index, index, index
-// CHECK: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// CHECK-SAME: : memref<f32> -> index
-// CHECK: %[[MUL:.+]] = arith.muli %[[OFFSET1]], %[[ELEM_BYTES]] : index
-// CHECK: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// CHECK: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [32],
-// CHECK-SAME: strides : [1] : i64 -> !xegpu.tensor_desc<8xf32,
-// CHECK-SAME: boundary_check = false
+// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0] [1, 1, 32] [1, 1, 1] : memref<8x16x32xf32> to memref<32xf32, strided<[1], offset: ?>>
+// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<32xf32, strided<[1], offset: ?>> -> !xegpu.tensor_desc<8xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]]] : vector<8xf32>
// -----
@@ -38,16 +28,8 @@ func.func @store_2D_vector(%vec: vector<8x16xf32>,
// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>,
// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
-// CHECK: %[[ELEM_BYTES:.*]] = arith.constant 4 : index
-// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]
-// CHECK: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// CHECK-SAME: : memref<f32> -> index
-// CHECK: %[[MUL:.+]] = arith.muli %[[OFF1]], %[[ELEM_BYTES]] : index
-// CHECK: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// CHECK: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [16, 32],
-// CHECK-SAME: strides : [32, 1] : i64 -> !xegpu.tensor_desc<8x16xf32>
+// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0] [1, 16, 32] [1, 1, 1] : memref<8x16x32xf32> to memref<16x32xf32, strided<[32, 1], offset: ?>>
+// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>
// -----
@@ -63,15 +45,9 @@ func.func @store_dynamic_source(%vec: vector<8x16xf32>,
// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>,
// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,
// CHECK-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// CHECK: %[[ELEM_BYTES:.*]] = arith.constant 4 : index
-// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]
-// CHECK: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// CHECK: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]] : memref<f32> -> index
-// CHECK: %[[MUL:.+]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// CHECK: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// CHECK: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [%[[SIZES]]#0, %[[SIZES]]#1],
-// CHECK-SAME: strides : [%[[STRIDES]]#0, 1] : i64 -> !xegpu.tensor
+// CHECK: %{{.*}}, %{{.*}}, %[[SIZES:.+]]:3, %{{.+}}:3 = memref.extract_strided_metadata %[[SRC]]
+// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0] [1, %[[SIZES]]#1, %[[SIZES]]#2] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<?x?xf32, strided<[?, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF1]], %[[OFF2]]] : vector<8x16xf32>
// -----
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 4cc4a0db5b63c..1464aa1ace477 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -35,17 +35,8 @@ gpu.func @load_2D_vector(%source: memref<8x16x32xf32>,
// LOAD-ND-LABEL: @load_2D_vector(
// LOAD-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// LOAD-ND-SAME: %[[OFFSET:.+]]: index
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// LOAD-ND-SAME: : memref<f32> -> index
-// LOAD-ND: %[[MUL:.*]] = arith.muli %[[OFF1]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.*]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [16, 32],
-// LOAD-ND-SAME: strides : [32, 1] : i64 -> !xegpu.tensor_desc<8x16xf32,
-// LOAD-ND-SAME: boundary_check = false
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0] [1, 16, 32] [1, 1, 1] : memref<8x16x32xf32> to memref<16x32xf32, strided<[32, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<8x16xf32>
// LOAD-ND: return %[[VEC]]
@@ -137,17 +128,8 @@ gpu.func @load_transpose_3d_memref(%source: memref<32x64x128xf32>,
// LOAD-ND-LABEL: @load_transpose_3d_memref(
// LOAD-ND-SAME: %[[SRC:.+]]: memref<32x64x128xf32>,
// LOAD-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index) -> vector<8x16xf32> {
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// LOAD-ND-SAME: : memref<f32> -> index
-// LOAD-ND: %[[MUL:.*]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.*]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [64, 128],
-// LOAD-ND-SAME: strides : [128, 1] : i64 -> !xegpu.tensor_desc<16x8xf32,
-// LOAD-ND-SAME: #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0] [1, 64, 128] [1, 1, 1] : memref<32x64x128xf32> to memref<64x128xf32, strided<[128, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<64x128xf32, strided<[128, 1], offset: ?>> -> !xegpu.tensor_desc<16x8xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF1]], %[[OFF2]]]
// LOAD-ND-SAME: : !xegpu.tensor_desc<16x8xf32, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<16x8xf32>
// LOAD-ND: %[[VEC_TRANSPOSED:.+]] = vector.transpose %[[VEC]], [1, 0] : vector<16x8xf32> to vector<8x16xf32>
@@ -201,16 +183,9 @@ gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
// LOAD-ND-LABEL: @load_dynamic_source(
// LOAD-ND-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,
// LOAD-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]] : memref<f32> -> index
-// LOAD-ND: %[[MUL:.*]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.*]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [%[[SIZES]]#0, %[[SIZES]]#1],
-// LOAD-ND-SAME: strides : [%[[STRIDES]]#0, 1] : i64 -> !xegpu.tensor_desc<8x16xf32,
-// LOAD-ND-SAME: #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %{{.+}}, %{{.+}}, %[[SIZES:.+]]:3, %{{.+}}:3 = memref.extract_strided_metadata %[[SRC]]
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0] [1, %[[SIZES]]#1, %[[SIZES]]#2] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<?x?xf32, strided<[?, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF1]], %[[OFF2]]]{{.*}}-> vector<8x16xf32>
// LOAD-ND: return %[[VEC]]
@@ -244,15 +219,8 @@ gpu.func @load_dynamic_source2(%source: memref<?x8x16xf32>,
// LOAD-ND-LABEL: @load_dynamic_source2(
// LOAD-ND-SAME: %[[SRC:.+]]: memref<?x8x16xf32>,
// LOAD-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// LOAD-ND: %[[MUL:.*]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.*]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [8, 16], strides : [16, 1] :
-// LOAD-ND-SAME: i64 -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0] [1, 8, 16] [1, 1, 1] : memref<?x8x16xf32> to memref<8x16xf32, strided<[16, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<8x16xf32, strided<[16, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%{{.*}}, %{{.*}}] : !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<8x16xf32>
// LOAD-ND: return %[[VEC]] : vector<8x16xf32>
@@ -283,8 +251,8 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
// LOAD-ND-LABEL: @load_dynamic_source3(
// LOAD-ND-SAME: %[[SRC:.+]]: memref<?x?x?x?x?xf32>
// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:4, %[[STRIDES:.*]]:4 = memref.extract_strided_metadata %[[SUBVIEW]]
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc
+// LOAD-ND-SAME: to memref<?x?x?x?xf32, strided<[?, ?, ?, 1], offset: ?>>
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SUBVIEW]]
// LOAD-ND-SAME: -> !xegpu.tensor_desc<2x4x8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]]
// LOAD-ND-SAME: -> vector<2x4x8x16xf32>
@@ -307,6 +275,36 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
// LOAD-GATHER: return %[[VEC]]
}
+// -----
+// Equal vector and memref rank: the whole memref stays the create_nd source.
+gpu.module @xevm_module {
+gpu.func @load_high_dim_dyn(%source: memref<?x?x8x16xf16>,
+ %i: index, %j: index, %k: index, %l: index) -> vector<2x4x8x16xf16> {
+ %pad = arith.constant 0.0 : f16
+ %0 = vector.transfer_read %source[%i, %j, %k, %l], %pad
+ {in_bounds = [true, true, true, true]}
+ : memref<?x?x8x16xf16>, vector<2x4x8x16xf16>
+ gpu.return %0 : vector<2x4x8x16xf16>
+}
+
+// LOAD-ND-LABEL: @load_high_dim_dyn(
+// LOAD-ND-SAME: %[[SRC:.+]]: memref<?x?x8x16xf16>,
+// LOAD-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index
+// LOAD-ND-NOT: memref.subview
+// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<?x?x8x16xf16> -> !xegpu.tensor_desc<2x4x8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF0]], %[[OFF1]], %[[OFF2]], %[[OFF3]]]{{.*}}-> vector<2x4x8x16xf16>
+// LOAD-ND: return %[[VEC]]
+
+// LOAD-GATHER-LABEL: @load_high_dim_dyn(
+// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<?x?x8x16xf16>
+// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<2x4x8x16xi1>
+// LOAD-GATHER: memref.extract_strided_metadata %[[SRC]]
+// LOAD-GATHER: %[[PTR:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x8x16xf16> -> index
+// LOAD-GATHER: %[[PTR_I:.+]] = arith.index_cast %[[PTR]] : index to i64
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[PTR_I]]{{\[}}%{{.+}}{{\]}}, %[[CST]] : i64, vector<2x4x8x16xindex>, vector<2x4x8x16xi1> -> vector<2x4x8x16xf16>
+// LOAD-GATHER: return %[[VEC]]
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_high_dim_vector(%source: memref<16x32x64xf32>,
@@ -530,15 +528,8 @@ gpu.func @load_from_subview_2D(%source: memref<4096x4096xf16>, %off1: index, %of
// LOAD-ND-LABEL: @load_from_subview_2D(
// LOAD-ND-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
// LOAD-ND-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 2 : index
// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[SUBVIEW]]
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// LOAD-ND: %[[MUL:.*]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.*]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [256, 256], strides : [4096, 1] :
-// LOAD-ND-SAME: i64 -> !xegpu.tensor_desc<8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF2]], %[[OFF2]]]{{.*}}-> vector<8x16xf16>
// LOAD-ND: return %[[VEC]]
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index dad2740f5c0ea..c6432378b496a 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -38,17 +38,8 @@ gpu.func @store_2D_vector(%vec: vector<8x16xf32>,
// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16xf32>,
// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]
-// STORE-ND: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// STORE-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// STORE-ND-SAME: : memref<f32> -> index
-// STORE-ND: %[[MUL:.+]] = arith.muli %[[OFF1]], %[[ELEM_BYTES]] : index
-// STORE-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// STORE-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [16, 32],
-// STORE-ND-SAME: strides : [32, 1] : i64 -> !xegpu.tensor_desc<8x16xf32,
-// STORE-ND-SAME: boundary_check = false
+// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0] [1, 16, 32] [1, 1, 1] : memref<8x16x32xf32> to memref<16x32xf32, strided<[32, 1], offset: ?>>
+// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>
// STORE-SCATTER-LABEL: @store_2D_vector(
@@ -80,15 +71,9 @@ gpu.func @store_dynamic_source(%vec: vector<8x16xf32>,
// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16xf32>,
// STORE-ND-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,
// STORE-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// STORE-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]
-// STORE-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]
-// STORE-ND: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]] : memref<f32> -> index
-// STORE-ND: %[[MUL:.+]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// STORE-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// STORE-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [%[[SIZES]]#0, %[[SIZES]]#1],
-// STORE-ND-SAME: strides : [%[[STRIDES]]#0, 1] : i64 -> !xegpu.tensor
+// STORE-ND: %{{.*}}, %{{.*}}, %[[SIZES:.+]]:3, %{{.+}}:3 = memref.extract_strided_metadata %[[SRC]]
+// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0] [1, %[[SIZES]]#1, %[[SIZES]]#2] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]] : memref<?x?xf32, strided<[?, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF1]], %[[OFF2]]] : vector<8x16xf32>
// STORE-SCATTER-LABEL: @store_dynamic_source(
@@ -106,6 +91,33 @@ gpu.func @store_dynamic_source(%vec: vector<8x16xf32>,
// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>
}
+// -----
+// Equal vector and memref rank: the whole memref stays the create_nd source.
+gpu.module @xevm_module {
+gpu.func @store_high_dim_dyn(%vec: vector<1x1x8x16xf16>, %source: memref<?x?x8x16xf16>,
+ %i: index, %j: index, %k: index, %l: index) {
+ vector.transfer_write %vec, %source[%i, %j, %k, %l]
+ {in_bounds = [true, true, true, true]}
+ : vector<1x1x8x16xf16>, memref<?x?x8x16xf16>
+ gpu.return
+}
+
+// STORE-ND-LABEL: @store_high_dim_dyn(
+// STORE-ND-SAME: %[[VEC:.+]]: vector<1x1x8x16xf16>, %[[SRC:.+]]: memref<?x?x8x16xf16>,
+// STORE-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index
+// STORE-ND-NOT: memref.subview
+// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[SRC]] : memref<?x?x8x16xf16> -> !xegpu.tensor_desc<1x1x8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF0]], %[[OFF1]], %[[OFF2]], %[[OFF3]]] : vector<1x1x8x16xf16>
+
+// STORE-SCATTER-LABEL: @store_high_dim_dyn(
+// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<1x1x8x16xf16>, %[[SRC:.+]]: memref<?x?x8x16xf16>
+// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<1x1x8x16xi1>
+// STORE-SCATTER: memref.extract_strided_metadata %[[SRC]]
+// STORE-SCATTER: %[[PTR:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x8x16xf16> -> index
+// STORE-SCATTER: %[[PTR_I:.+]] = arith.index_cast %[[PTR]] : index to i64
+// STORE-SCATTER: xegpu.store %[[VEC]], %[[PTR_I]]{{\[}}%{{.+}}{{\]}}, %[[CST]] : vector<1x1x8x16xf16>, i64, vector<1x1x8x16xindex>, vector<1x1x8x16xi1>
+}
+
// -----
gpu.module @xevm_module {
gpu.func @store_out_of_bounds(%vec: vector<8x16xf32>,
diff --git a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
index 34654126ce8d2..939adc951ac38 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -44,23 +44,40 @@ gpu.module @create_nd_tdesc {
// CHECK: %[[VAR19:.*]] = vector.insert %[[PITCH2]], %[[VAR18]] [4] : i32 into vector<8xi32>
%src_tdesc = xegpu.create_nd_tdesc %srcce : memref<16x32xf32> -> !xegpu.tensor_desc<8x16xf32>
- // CHECK: %[[C1:.*]] = arith.constant 1 : index
- %c1 = arith.constant 1 : index
- // CHECK: %[[C64:.*]] = arith.constant 64 : index
- %size_x = arith.constant 64 : index
- // CHECK: %[[C16:.*]] = arith.constant 16 : index
- %BLOCK_DMODEL = arith.constant 16 : index
+ // A dynamic memref uses the bare form; shape/strides come from it.
+ // CHECK: %{{.*}}, %{{.*}}, %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[DYN]] : memref<?x?xf16>
// CHECK: %[[CST_3:.*]] = arith.constant dense<0> : vector<8xi32>
- // CHECK: %[[SHAPE_W3:.*]] = arith.index_cast %[[C16]] : index to i32
- // CHECK: %[[SHAPE_H3:.*]] = arith.index_cast %[[C64]] : index to i32
- // CHECK: %[[PITCH3:.*]] = arith.index_cast %[[C16]] : index to i32
+ // CHECK: %[[SHAPE_W3:.*]] = arith.index_cast %[[SIZES]]#1 : index to i32
+ // CHECK: %[[SHAPE_H3:.*]] = arith.index_cast %[[SIZES]]#0 : index to i32
+ // CHECK: %[[PITCH3:.*]] = arith.index_cast %[[STRIDES]]#0 : index to i32
// CHECK: %[[VAR25:.*]] = vector.bitcast %[[CST_3]] : vector<8xi32> to vector<4xi64>
- // CHECK: %[[VAR26:.*]] = vector.insert %[[DYN_ADDR_OFFSET:.*]], %[[VAR25]] [0] : i64 into vector<4xi64>
+ // CHECK: %[[VAR26:.*]] = vector.insert %{{.*}}, %[[VAR25]] [0] : i64 into vector<4xi64>
// CHECK: %[[VAR27:.*]] = vector.bitcast %[[VAR26]] : vector<4xi64> to vector<8xi32>
// CHECK: %[[VAR28:.*]] = vector.insert %[[SHAPE_W3]], %[[VAR27]] [2] : i32 into vector<8xi32>
// CHECK: %[[VAR29:.*]] = vector.insert %[[SHAPE_H3]], %[[VAR28]] [3] : i32 into vector<8xi32>
// CHECK: %[[VAR30:.*]] = vector.insert %[[PITCH3]], %[[VAR29]] [4] : i32 into vector<8xi32>
- %dyn_tdesc = xegpu.create_nd_tdesc %dyn, shape: [%size_x, %BLOCK_DMODEL], strides: [%BLOCK_DMODEL, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<16x16xf16>
+ %dyn_tdesc = xegpu.create_nd_tdesc %dyn : memref<?x?xf16> -> !xegpu.tensor_desc<16x16xf16>
gpu.return
}
+
+ // Batched (>2D): base_height spans all planes; slot 5 = batch row stride.
+ // CHECK-LABEL: gpu.func @create_nd_tdesc_batch_dyn(
+ // CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf16>
+ gpu.func @create_nd_tdesc_batch_dyn(%src: memref<?x?x?xf16>) -> vector<8xi32> {
+ // CHECK: %{{.+}}, %{{.+}}, %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
+ // CHECK: %[[W:.+]] = arith.index_cast %[[SIZES]]#2 : index to i32
+ // CHECK: %[[H:.+]] = arith.index_cast %[[SIZES]]#1 : index to i32
+ // CHECK: %[[BATCH:.+]] = arith.index_cast %[[SIZES]]#0 : index to i32
+ // CHECK: %[[FLAT_H:.+]] = arith.muli %[[H]], %[[BATCH]] : i32
+ // CHECK: %[[PITCH:.+]] = arith.index_cast %[[STRIDES]]#1 : index to i32
+ // CHECK: %[[P2:.+]] = vector.insert %[[W]], %{{.+}} [2] : i32 into vector<8xi32>
+ // CHECK: %[[P3:.+]] = vector.insert %[[FLAT_H]], %[[P2]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[P4:.+]] = vector.insert %[[PITCH]], %[[P3]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[LS0:.+]] = arith.index_cast %[[STRIDES]]#0 : index to i32
+ // CHECK: %[[ROWS0:.+]] = arith.divui %[[LS0]], %[[PITCH]] : i32
+ // CHECK: vector.insert %[[ROWS0]], %[[P4]] [5] : i32 into vector<8xi32>
+ %t = xegpu.create_nd_tdesc %src : memref<?x?x?xf16> -> !xegpu.tensor_desc<1x8x16xf16>
+ %c = builtin.unrealized_conversion_cast %t : !xegpu.tensor_desc<1x8x16xf16> to vector<8xi32>
+ gpu.return %c : vector<8xi32>
+ }
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index 93bda1dfc0dca..dfaa9521786a9 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -54,3 +54,53 @@ gpu.module @test_kernel {
gpu.return
}
}
+
+// -----
+
+// A non-unit leading dim cannot be lowered to a 2D-block op.
+
+gpu.module @test_kernel {
+ gpu.func @load_nd_non_unit_batch(%src: memref<4x8x16xf32>, %z: index) kernel {
+ %c0 = arith.constant 0 : index
+ %t = xegpu.create_nd_tdesc %src : memref<4x8x16xf32> -> !xegpu.tensor_desc<2x8x16xf32>
+ // expected-error at +1 {{failed to legalize operation 'xegpu.load_nd' that was explicitly marked illegal}}
+ %v = xegpu.load_nd %t[%z, %c0, %c0] : !xegpu.tensor_desc<2x8x16xf32> -> vector<16xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+// The payload has room for only 3 leading strides, so rank > 5 is rejected.
+
+gpu.module @test_kernel {
+ gpu.func @create_nd_tdesc_rank_too_large(%src: memref<2x2x2x2x8x16xf32>) kernel {
+ // expected-error at +1 {{failed to legalize operation 'xegpu.create_nd_tdesc' that was explicitly marked illegal}}
+ %t = xegpu.create_nd_tdesc %src : memref<2x2x2x2x8x16xf32> -> !xegpu.tensor_desc<1x1x1x1x8x16xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+// A memref with gaps between planes has no exact flattened-plane view.
+
+gpu.module @test_kernel {
+ gpu.func @create_nd_tdesc_plane_gap(%src: memref<2x8x16xf32, strided<[200, 16, 1]>>) kernel {
+ // expected-error at +1 {{failed to legalize operation 'xegpu.create_nd_tdesc' that was explicitly marked illegal}}
+ %t = xegpu.create_nd_tdesc %src : memref<2x8x16xf32, strided<[200, 16, 1]>> -> !xegpu.tensor_desc<1x8x16xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+// Same check applies to an integer source, whose strides are explicit.
+
+gpu.module @test_kernel {
+ gpu.func @create_nd_tdesc_plane_gap_ptr(%ptr: i64) kernel {
+ // expected-error at +1 {{failed to legalize operation 'xegpu.create_nd_tdesc' that was explicitly marked illegal}}
+ %t = xegpu.create_nd_tdesc %ptr, shape: [2, 8, 16], strides: [200, 16, 1] : i64 -> !xegpu.tensor_desc<1x8x16xf32>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
index d45fa79bb2e63..408950c0d8f0f 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
@@ -91,4 +91,46 @@ gpu.module @load_store_check {
vector.store %loaded, %dstte[%c0, %c0] : memref<32x16xi8>, vector<32xi8>
gpu.return
}
+
+ // Batched (>2D): the batch offset becomes a row offset; base ptr unchanged.
+ // CHECK-LABEL: gpu.func @load_store_batch_dyn(
+ // CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>, %[[DST:.+]]: memref<?x?x?xf32>, %[[Z:.+]]: index
+ gpu.func @load_store_batch_dyn(%src: memref<?x?x?xf32>, %dst: memref<?x?x?xf32>,
+ %z: index) kernel {
+ %c0 = arith.constant 0 : index
+ // CHECK-DAG: %[[C4I32:.+]] = arith.constant 4 : i32
+
+ // In-plane offset is 0, so canonicalize folds the add away.
+ // CHECK: %{{.+}}, %{{.+}}, %[[LSZ:.+]]:3, %[[LSTR:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
+ // CHECK: %[[LH:.+]] = arith.index_cast %[[LSZ]]#1 : index to i32
+ // CHECK: %[[LBATCH:.+]] = arith.index_cast %[[LSZ]]#0 : index to i32
+ // CHECK: %[[LFLAT_H:.+]] = arith.muli %[[LH]], %[[LBATCH]] : i32
+ // CHECK: %[[LPITCH:.+]] = arith.index_cast %[[LSTR]]#1 : index to i32
+ // CHECK: %[[LSTRIDE:.+]] = arith.index_cast %[[LSTR]]#0 : index to i32
+ // CHECK: %[[LROWS:.+]] = arith.divui %[[LSTRIDE]], %[[LPITCH]] : i32
+ // CHECK: %[[LBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
+ // CHECK: %[[LZ:.+]] = arith.index_cast %[[Z]] : index to i32
+ // CHECK: %[[LY:.+]] = arith.muli %[[LZ]], %[[LROWS]] : i32
+ // CHECK: %[[LPTR:.+]] = llvm.inttoptr %[[LBASE]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.blockload2d %[[LPTR]], %{{.+}}, %[[LFLAT_H]], %{{.+}}, %{{.+}}, %[[LY]]{{.*}}tile_height = 8 : i32, tile_width = 16 : i32
+ %st = xegpu.create_nd_tdesc %src : memref<?x?x?xf32> -> !xegpu.tensor_desc<1x8x16xf32>
+ %v = xegpu.load_nd %st[%z, %c0, %c0] : !xegpu.tensor_desc<1x8x16xf32> -> vector<8xf32>
+
+ // Store: same handling.
+ // CHECK: %{{.+}}, %{{.+}}, %[[SSZ:.+]]:3, %[[SSTR:.+]]:3 = memref.extract_strided_metadata %[[DST]]
+ // CHECK: %[[SH:.+]] = arith.index_cast %[[SSZ]]#1 : index to i32
+ // CHECK: %[[SBATCH:.+]] = arith.index_cast %[[SSZ]]#0 : index to i32
+ // CHECK: %[[SFLAT_H:.+]] = arith.muli %[[SH]], %[[SBATCH]] : i32
+ // CHECK: %[[SPITCH:.+]] = arith.index_cast %[[SSTR]]#1 : index to i32
+ // CHECK: %[[SSTRIDE:.+]] = arith.index_cast %[[SSTR]]#0 : index to i32
+ // CHECK: %[[SROWS:.+]] = arith.divui %[[SSTRIDE]], %[[SPITCH]] : i32
+ // CHECK: %[[SBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
+ // CHECK: %[[SZ:.+]] = arith.index_cast %[[Z]] : index to i32
+ // CHECK: %[[SY:.+]] = arith.muli %[[SZ]], %[[SROWS]] : i32
+ // CHECK: %[[SPTR:.+]] = llvm.inttoptr %[[SBASE]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.blockstore2d %[[SPTR]], %{{.+}}, %[[SFLAT_H]], %{{.+}}, %{{.+}}, %[[SY]]{{.*}}tile_height = 8 : i32, tile_width = 16 : i32
+ %dt = xegpu.create_nd_tdesc %dst : memref<?x?x?xf32> -> !xegpu.tensor_desc<1x8x16xf32>
+ xegpu.store_nd %v, %dt[%z, %c0, %c0] : vector<8xf32>, !xegpu.tensor_desc<1x8x16xf32>
+ gpu.return
+ }
}
diff --git a/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir b/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
index 3582f8995af75..693ba46dbf74b 100644
--- a/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
@@ -203,11 +203,10 @@ gpu.module @test {
// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf16>, %[[H:.*]]: index, %[[W:.*]]: index)
func.func @test_dynamic_memref_source(%arg0: memref<?x?xf16>, %h: index, %w: index) -> vector<16x16xf16> {
%c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]], shape : [%[[H]], %[[W]]], strides : [%[[W]], %{{.*}}]
- // CHECK-SAME: memref<?x?xf16> -> !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>>
- %tdesc = xegpu.create_nd_tdesc %arg0, shape : [%h, %w], strides : [%w, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<32x32xf16>
+ // A dynamic memref uses the bare form; its shape/strides come from the memref.
+ // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<?x?xf16> -> !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>>
+ %tdesc = xegpu.create_nd_tdesc %arg0 : memref<?x?xf16> -> !xegpu.tensor_desc<32x32xf16>
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][%{{.*}}, %{{.*}}]
// CHECK-SAME: -> vector<64x16xf16>
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index 2648303a3bb55..932cf6ff9b201 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -9,6 +9,18 @@ func.func @create_nd_tdesc_1(%src: memref<24xf32>) {
// -----
+// Explicit shape/strides on a memref source is deprecated and rejected.
+func.func @create_nd_tdesc_memref_explicit_shape(%src: memref<?x?xf16>,
+ %h: index, %w: index) {
+ %c1 = arith.constant 1 : index
+ // expected-error at +1 {{shape and strides should not be specified for a memref source}}
+ %1 = xegpu.create_nd_tdesc %src, shape: [%h, %w], strides: [%w, %c1]
+ : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
+ return
+}
+
+// -----
+
func.func @create_nd_tdesc_2(%src: memref<24x32xf32>) {
// expected-error at +1 {{TensorDesc should have the same element type with the source if it is a memref}}
%1 = xegpu.create_nd_tdesc %src : memref<24x32xf32> -> !xegpu.tensor_desc<8x16xf16>
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index fcafe17c39746..5e795017e2122 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -81,19 +81,17 @@ gpu.func @test_create_nd_tdesc_8(%src: ui64, %w : index, %h : index, %x : index,
// CHECK-LABEL: func @test_create_nd_tdesc_9({{.*}})
gpu.func @test_create_nd_tdesc_9(%src: memref<?x?xf16>, %w : index, %h : index, %x : index, %y : index) {
-
- %c1 = arith.constant 1 : index
- // CHECK: %[[REG:.*]] = xegpu.create_nd_tdesc %arg0, shape : [%arg2, %arg1], strides : [%arg1, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
- %1 = xegpu.create_nd_tdesc %src , shape:[%h, %w], strides:[%w, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
+ // A dynamic-shape memref uses the bare form; shape/strides come from it.
+ // CHECK: %[[REG:.*]] = xegpu.create_nd_tdesc %arg0 : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %1 = xegpu.create_nd_tdesc %src : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
gpu.return
}
// CHECK-LABEL: func @test_create_nd_tdesc_10({{.*}})
gpu.func @test_create_nd_tdesc_10(%src: memref<?x?xf16>, %w : index, %h : index, %x : index, %y : index) {
- %c1 = arith.constant 1 : index
- // CHECK: %[[REG:.*]] = xegpu.create_nd_tdesc %arg0, shape : [%arg2, %arg1], strides : [%arg1, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
- %2 = xegpu.create_nd_tdesc %src, shape:[%h, %w], strides:[%w, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
+ // CHECK: %[[REG:.*]] = xegpu.create_nd_tdesc %arg0 : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
+ %2 = xegpu.create_nd_tdesc %src : memref<?x?xf16> -> !xegpu.tensor_desc<8x16xf16>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index 3bcdadaca3f0e..00c810edd605a 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -504,3 +504,32 @@ gpu.module @xevm_test {
gpu.return
}
}
+
+// -----
+// Transpose optimization on a >2D descriptor with unit leading dims.
+// CHECK-LABEL: gpu.func @transpose_4d(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<?x?x64x64xf16>) -> vector<1x1x16x16xf16> {
+// CHECK-DAG: %[[C2048:.*]] = arith.constant 2048 : index
+// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %{{.+}}, %{{.+}}, %[[SIZES:.+]]:4, %[[STRIDES:.+]]:4 = memref.extract_strided_metadata %[[ARG0]]
+// CHECK: %[[LSTRIDE:.*]] = arith.shrui %[[STRIDES]]#0, %[[C1]] : index
+// CHECK: %[[PTR:.*]] = memref.extract_aligned_pointer_as_index %{{.+}} : memref<f16> -> index
+// CHECK: %[[T0:.*]] = arith.index_cast %[[PTR]] : index to i64
+// CHECK: %[[BDESC:.*]] = xegpu.create_nd_tdesc %[[T0]], shape : [%[[SIZES]]#0, %[[SIZES]]#1, 64, %[[C32]]], strides : [%[[LSTRIDE]], %[[C2048]], %[[C32]], 1] : i64
+// CHECK-SAME: -> !xegpu.tensor_desc<1x1x16x8xi32, #xegpu.layout<lane_layout = [1, 1, 16, 1], lane_data = [1, 1, 1, 1], order = [2, 3, 1, 0]>>
+// CHECK: %[[B:.*]] = xegpu.load_nd %[[BDESC]][%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}]{{.*}} -> vector<1x1x16x8xi32>
+// CHECK: %[[BITCAST:.*]] = vector.bitcast %[[B]] : vector<1x1x16x8xi32> to vector<1x1x16x16xf16>
+// CHECK: vector.transpose %[[BITCAST]], [0, 1, 3, 2] : vector<1x1x16x16xf16> to vector<1x1x16x16xf16>
+#b4 = #xegpu.layout<lane_layout = [1, 1, 16, 1], lane_data = [1, 1, 1, 2], order = [2, 3, 1, 0]>
+#bt4 = #xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 2, 1]>
+gpu.module @xevm_module {
+gpu.func @transpose_4d(%arg0: memref<?x?x64x64xf16>) -> vector<1x1x16x16xf16> {
+ %c0 = arith.constant 0 : index
+ %c32 = arith.constant 32 : index
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<?x?x64x64xf16> -> !xegpu.tensor_desc<1x1x16x16xf16, #b4>
+ %1 = xegpu.load_nd %0[%c0, %c0, %c0, %c32] { result_layout = #b4 } : !xegpu.tensor_desc<1x1x16x16xf16, #b4> -> vector<1x1x16x16xf16>
+ %2 = vector.transpose %1, [0, 1, 3, 2] { layout_result_0 = #bt4 } : vector<1x1x16x16xf16> to vector<1x1x16x16xf16>
+ gpu.return %2 : vector<1x1x16x16xf16>
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 6cde69e3cde6e..0f3cea16797f1 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -60,6 +60,32 @@ gpu.func @load_nd_transpose() {
gpu.return
}
+// >2D load with unit leading dims: packed set from the inner lane_data.
+// CHECK-LABEL: gpu.func @load_nd_packed_4d
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]] <{packed}> : !xegpu.tensor_desc<1x1x16x16xf16> -> vector<16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<16xf16> to vector<1x1x16x1xf16>
+gpu.func @load_nd_packed_4d() {
+ %c0 = arith.constant 0 : index
+ %0 = "some_op"() : () -> !xegpu.tensor_desc<1x1x16x16xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0, %c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 2, 1]>}
+ : !xegpu.tensor_desc<1x1x16x16xf16> -> vector<1x1x16x16xf16>
+ gpu.return
+}
+
+// >2D load with unit leading dims: transposed on the inner 2D tile.
+// CHECK-LABEL: gpu.func @load_nd_transpose_4d
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]] <{transpose = array<i64: 1, 0>}> : !xegpu.tensor_desc<1x1x16x8xf32> -> vector<8xf32>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<8xf32> to vector<1x1x1x8xf32>
+gpu.func @load_nd_transpose_4d() {
+ %c0 = arith.constant 0 : index
+ %0 = "some_op"() : () -> !xegpu.tensor_desc<1x1x16x8xf32>
+ %1 = xegpu.load_nd %0[%c0, %c0, %c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 1, 16, 1], lane_data = [1, 1, 1, 1]>}
+ : !xegpu.tensor_desc<1x1x16x8xf32> -> vector<1x1x16x8xf32>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @load_nd_array_length
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}}[%[[C0]], %[[C0]]] : !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<array_length = 2 : i64>> -> vector<64xf16>
diff --git a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
index 62f17457b007b..f23ae46905652 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -360,5 +360,34 @@ gpu.module @test {
gpu.return %0 : vector<32xf32>
}
+//-----
+ // Unrolling a >2D nd desc keeps the whole memref as create_nd source.
+ // CHECK-LABEL: gpu.func @load_store_nd_3d
+ // CHECK-SAME: [[arg0:%.+]]: memref<4x8x16xf32>, [[z:%.+]]: index
+ // CHECK-NOT: memref.subview
+ // CHECK: [[t:%.+]] = xegpu.create_nd_tdesc [[arg0]] : memref<4x8x16xf32> -> !xegpu.tensor_desc<1x8x16xf32>
+ // CHECK: xegpu.load_nd [[t]]{{\[}}[[z]], {{.*}}] : !xegpu.tensor_desc<1x8x16xf32> -> vector<1x8x16xf32>
+ // CHECK: [[z1:%.+]] = arith.addi [[z]], {{%.+}}
+ // CHECK: xegpu.load_nd [[t]]{{\[}}[[z1]], {{.*}}]
+ // CHECK: [[z2:%.+]] = arith.addi [[z]], {{%.+}}
+ // CHECK: xegpu.load_nd [[t]]{{\[}}[[z2]], {{.*}}]
+ // CHECK: [[z3:%.+]] = arith.addi [[z]], {{%.+}}
+ // CHECK: xegpu.load_nd [[t]]{{\[}}[[z3]], {{.*}}]
+ // CHECK: xegpu.store_nd {{%.+}}, [[t]]{{\[}}[[z]], {{.*}}]
+ // CHECK-COUNT-3: xegpu.store_nd {{%.+}}, [[t]]
+ // CHECK-NOT: memref.subview
+ gpu.func @load_store_nd_3d(%src: memref<4x8x16xf32>, %z: index) {
+ %c0 = arith.constant 0 : index
+ %t = xegpu.create_nd_tdesc %src : memref<4x8x16xf32>
+ -> !xegpu.tensor_desc<4x8x16xf32, #xegpu.layout<inst_data = [1, 8, 16]>>
+ %v = xegpu.load_nd %t[%z, %c0, %c0]
+ : !xegpu.tensor_desc<4x8x16xf32, #xegpu.layout<inst_data = [1, 8, 16]>>
+ -> vector<4x8x16xf32>
+ xegpu.store_nd %v, %t[%z, %c0, %c0]
+ : vector<4x8x16xf32>,
+ !xegpu.tensor_desc<4x8x16xf32, #xegpu.layout<inst_data = [1, 8, 16]>>
+ gpu.return
+ }
+
}
More information about the Mlir-commits
mailing list