[Mlir-commits] [mlir] [mlir][xegpu] Lower dynamic high-D nd load/store via base-pointer fold (PR #215711)
Jianhui Li
llvmlistbot at llvm.org
Sat Aug 22 12:28:08 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/215711
>From 2ac466b9f32759fe8645c799402f35c4359ee777 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 11 Aug 2026 21:32:12 +0000
Subject: [PATCH 1/8] [mlir][xegpu] Lower dynamic high-D nd load/store via
base-pointer fold
Replace the per-batch memref.subview scheme for >2D xegpu.create_nd_tdesc /
load_nd / store_nd / prefetch_nd with a subview-free lowering that keeps the
original (possibly dynamic) high-D memref as the descriptor source and folds
the leading (batch) offsets into the base pointer at XeVM lowering time. The
2D-block surface stays the innermost matrix, so the HW surface-size limits are
unaffected, and a dynamic-shape source never needs a per-batch subview (which
could not guarantee a valid base). This fixes flash-attention rank-4 kernels
with dynamic z,h batch offsets (previously the batch offset was silently
dropped, so only head (0,0) was correct).
- VectorToXeGPU: keep the memref as the create_nd source (drop the
i64-pointer path); the high-D offsets ride the load/store.
- XeGPUToXeVM: encode the leading-dim element strides into spare payload
slots at create time and fold `base += sum_d offset[d]*stride[d]*elemBytes`
at load/store time (memref and integer/pointer sources); recover a bare
dynamic memref's shape/strides via extract_strided_metadata.
- XeGPUUnroll / WgToSgDistribute / ArrayLengthOptimization: rebuild create_nd
from the full memref (bare-memref builder, no subview); batch is unrolled
as unit-leading tiles carried on the offsets.
- create_nd_tdesc op: single-arg builder + verifier accept a bare dynamic
memref; specifying explicit shape/strides for a memref source is now
rejected (deprecated) since the memref is authoritative.
- Generalize the transpose optimization (PeepHoleOptimizer, requireTranspose)
and the VNNI pack decision (requirePacked) to >2D descriptors with unit
leading dims by operating on the innermost 2 dims; the 2D path is
unchanged. Convert all non-innermost strides (not just the pitch) to the
repacked element unit so the batch fold uses consistent units.
Updates the affected VectorToXeGPU / XeGPUToXeVM / XeGPU dialect tests.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 51 +---
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 112 ++++++-
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 28 +-
.../XeGPUArrayLengthOptimization.cpp | 22 +-
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 176 +++++++----
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 277 +++++-------------
.../Transforms/XeGPUWgToSgDistribute.cpp | 13 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 21 +-
.../VectorToXeGPU/load-to-xegpu.mlir | 38 +--
.../VectorToXeGPU/store-to-xegpu.mlir | 38 +--
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 63 +---
.../transfer-write-to-xegpu.mlir | 25 +-
.../XeGPUToXeVM/create_nd_tdesc.mlir | 19 +-
.../test/Dialect/XeGPU/array-len-op-unit.mlir | 7 +-
mlir/test/Dialect/XeGPU/ops.mlir | 13 +-
15 files changed, 427 insertions(+), 476 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 8a45836426931..083c63198733b 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -130,50 +130,17 @@ static xegpu::CreateNdDescOp createNdDescriptor(PatternRewriter &rewriter,
Location loc,
xegpu::TensorDescType descType,
TypedValue<MemRefType> src) {
- MemRefType srcTy = src.getType();
+ [[maybe_unused]] 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;
+ // Keep the memref itself as the source (static or dynamic) rather than
+ // collapsing it to an i64 base address. The memref value carries the base
+ // pointer *and* the (possibly dynamic) offset/shape/strides, which the XeGPU
+ // -> XeVM lowering recovers via memref metadata. In particular, for a >2D
+ // descriptor the leading (batch) offsets stay on the load/store and are
+ // folded into the base pointer at lowering time, so a dynamic-shape source
+ // never needs a per-batch subview (which could not guarantee a valid base).
+ return xegpu::CreateNdDescOp::create(rewriter, loc, descType, src);
}
// Adjusts the strides of a memref according to a given permutation map for
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 78d99cf88b768..e30951fdb9e1d 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -51,10 +51,15 @@ 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, // Element strides of the leading (batch) dims of a >2D
+ LeadingStride1 = 6, // descriptor (i32). The load/store/prefetch lowering
+ LeadingStride2 = 7, // uses these to fold the batch offsets into the base
+ // pointer, keeping the 2D-block surface at the innermost
+ // matrix. Left at 0 for 2D descriptors.
};
static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
@@ -232,28 +237,61 @@ class CreateNdDescToXeVMPattern
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);
+
+ // Shape and strides. For a memref source they are recovered from the memref
+ // value itself (a dynamic memref carries them at runtime and does not
+ // attach them as op operands); otherwise they come from the op's explicit
+ // operands.
+ SmallVector<OpFoldResult> mixedSizes;
+ SmallVector<OpFoldResult> mixedStrides;
+ memref::ExtractStridedMetadataOp srcMeta;
+ // Lazily materialize memref metadata (base buffer, offset, sizes, strides).
+ auto getSrcMeta = [&]() -> memref::ExtractStridedMetadataOp {
+ if (!srcMeta)
+ srcMeta =
+ memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
+ return srcMeta;
+ };
// 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.");
}
+ SmallVector<int64_t> staticStrides;
+ int64_t staticOffset;
+ if (failed(
+ sourceMemrefTy.getStridesAndOffset(staticStrides, staticOffset)))
+ return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
+ // A memref is authoritative for its own shape/strides (explicit
+ // shape/strides on a memref create_nd are deprecated). A fully static
+ // memref yields constants directly; a dynamic one recovers its dynamic
+ // dims from runtime metadata (getConstified* keeps static dims as attrs).
+ bool allStatic = sourceMemrefTy.hasStaticShape() &&
+ llvm::none_of(staticStrides, ShapedType::isDynamic);
+ if (allStatic) {
+ mixedSizes = op.getMixedSizes();
+ mixedStrides = op.getMixedStrides();
+ } else {
+ mixedSizes = getSrcMeta().getConstifiedMixedSizes();
+ mixedStrides = getSrcMeta().getConstifiedMixedStrides();
+ }
// Access adaptor after failure check to avoid rolling back generated code
// for materialization cast.
baseAddr = adaptor.getSource();
} else {
+ mixedSizes = op.getMixedSizes();
+ mixedStrides = op.getMixedStrides();
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);
}
}
+ // Descriptor shape rank.
+ int64_t rank = mixedSizes.size();
// 1D tensor descriptor is just the base address.
if (rank == 1) {
rewriter.replaceOp(op, baseAddr);
@@ -288,6 +326,27 @@ class CreateNdDescToXeVMPattern
payload =
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
+ // For a >2D descriptor, encode the leading (batch) dim element strides into
+ // the spare payload slots. The load/store/prefetch lowering folds the batch
+ // offsets into the base pointer using these, so the batch position stays
+ // out of the 2D-block surface (which remains the innermost matrix computed
+ // above). This applies to both memref sources and integer (pointer) sources
+ // with explicit shape/strides -- e.g. a batched descriptor whose source was
+ // rewritten to an i64 base by the transpose peephole; in that case the
+ // strides come from the op's explicit operands. 2D descriptors leave these
+ // slots at 0 and the load/store path skips reading them (2D path
+ // unchanged).
+ if (rank > 2) {
+ if (rank - 2 > 3)
+ return rewriter.notifyMatchFailure(
+ op, "Batched nd descriptor supports at most 3 leading dims.");
+ for (int64_t d = 0; d < rank - 2; ++d) {
+ Value leadingStride = createOffset(mixedStrides, d);
+ payload = vector::InsertOp::create(
+ rewriter, loc, leadingStride, payload,
+ static_cast<int>(NdTdescOffset::LeadingStride0) + d);
+ }
+ }
rewriter.replaceOp(op, payload);
return success();
}
@@ -387,9 +446,8 @@ 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.
+ // x = innermost offset; y = second-to-last offset. These address the 2D
+ // tile within the innermost-matrix surface.
Value offsetW = getValueOrCreateConstantIntOp(rewriter, loc,
mixedOffsets[tileRank - 1]);
offsetW = getValueOrCreateCastToIndexLike(rewriter, loc,
@@ -398,6 +456,38 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
mixedOffsets[tileRank - 2]);
offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), offsetH);
+ // For a >2D descriptor, fold the leading (batch) offsets into the base
+ // pointer: basePtr += (sum_d offset[d] * leadingStride[d]) * elemBytes.
+ // The batch element strides were encoded into the payload at create time
+ // (see CreateNdDescToXeVMPattern). This keeps the 2D-block surface at the
+ // innermost matrix (so the HW surface-size limits are unaffected) instead
+ // of baking the batch into a (possibly out-of-range) collapsed surface.
+ if (tileRank > 2) {
+ Type i64Ty = rewriter.getI64Type();
+ Value batchElemOffset;
+ for (int64_t d = 0; d < tileRank - 2; ++d) {
+ Value off =
+ getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[d]);
+ off = getValueOrCreateCastToIndexLike(rewriter, loc, i64Ty, off);
+ Value strideI32 = vector::ExtractOp::create(
+ rewriter, loc, tdesc,
+ static_cast<int>(NdTdescOffset::LeadingStride0) + d);
+ Value stride =
+ arith::ExtUIOp::create(rewriter, loc, i64Ty, strideI32);
+ Value term = arith::MulIOp::create(rewriter, loc, off, stride);
+ batchElemOffset =
+ batchElemOffset
+ ? arith::AddIOp::create(rewriter, loc, batchElemOffset, term)
+ .getResult()
+ : term;
+ }
+ Value elemByteSizeI64 =
+ arith::ConstantIntOp::create(rewriter, loc, i64Ty, elemBitSize / 8);
+ Value batchByteOffset = arith::MulIOp::create(
+ rewriter, loc, batchElemOffset, elemByteSizeI64);
+ basePtr =
+ arith::AddIOp::create(rewriter, loc, basePtr, batchByteOffset);
+ }
// 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..9d1624d054bbc 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -212,9 +212,9 @@ 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");
-
+ // A dynamic-shape memref is allowed: its (possibly dynamic) shape/strides are
+ // carried by the memref value itself and recovered at lowering time via
+ // memref metadata, so no explicit shape/stride operands are attached here.
build(builder, state, tdesc, source, ValueRange({}) /* empty dynamic shape */,
ValueRange({}) /* empty dynamic strides */,
DenseI64ArrayAttr({}) /* empty const shape*/,
@@ -260,8 +260,13 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
}
LogicalResult CreateNdDescOp::verify() {
- size_t rank = getMixedSizes().size();
- bool invalidRank = rank != getMixedStrides().size();
+ // getMixedSizes()/getMixedStrides() cannot represent a *bare dynamic* memref
+ // (a dynamic dim has no SSA operand and no const attr), so derive the rank
+ // from the memref type for a memref source; only an integer source needs the
+ // shape/stride operands (checked below).
+ auto srcMemrefTy = dyn_cast<MemRefType>(getSourceType());
+ size_t rank = srcMemrefTy ? srcMemrefTy.getRank() : getMixedSizes().size();
+ bool invalidRank = false;
bool invalidElemTy = false;
// Memory space of created TensorDesc should match with the source.
@@ -280,11 +285,24 @@ 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.");
+ invalidRank = getMixedSizes().size() != getMixedStrides().size();
+ } else if (srcMemrefTy && hasExplicitShapeStrides) {
+ // Deprecated: a memref source carries its own shape/strides (recovered at
+ // lowering time, including for dynamic dims), so specifying them
+ // explicitly is redundant and no longer supported. Use the bare form:
+ // xegpu.create_nd_tdesc %memref : memref<...> -> !xegpu.tensor_desc<...>
+ return emitOpError("shape and strides should not be specified for a memref "
+ "source; they are inferred from the memref.");
}
if (invalidRank)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
index 590d1804167a5..ffca61128e283 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
@@ -137,13 +137,21 @@ 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());
+ // The memory region is unchanged. A memref source carries its own
+ // (possibly dynamic) shape/strides, so use the bare-memref builder rather
+ // than getMixedSizes/getMixedStrides (which cannot represent a dynamic
+ // memref dim, and for which explicit shape/strides on a memref are no
+ // longer allowed); an integer source still needs its explicit operands.
+ 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..4b7ee133006e5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -44,28 +44,39 @@ using namespace mlir;
namespace {
-/// Get the 2D lane data from a tensor desc type if it exists.
+/// Get the 2D lane data from a tensor desc type if it exists. The transpose
+/// optimization acts on the innermost 2 dims; a >2D descriptor is accepted only
+/// when its leading (batch) dims are unit, in which case the inner 2 dims are
+/// returned.
static std::optional<SmallVector<int64_t>>
getMaybeLaneData(xegpu::TensorDescType tdescType) {
auto layout = tdescType.getLayoutAttr();
if (!layout)
return std::nullopt;
auto laneData = layout.getEffectiveLaneDataAsInt();
- if (laneData.size() != 2)
+ if (laneData.size() < 2)
return std::nullopt;
- return laneData;
+ for (int64_t d : ArrayRef<int64_t>(laneData).drop_back(2))
+ if (d != 1)
+ return std::nullopt;
+ return SmallVector<int64_t>(laneData.end() - 2, laneData.end());
}
-/// Get the 2D lane layout from a tensor desc type if it exists.
+/// Get the 2D lane layout from a tensor desc type if it exists. As with
+/// getMaybeLaneData, a >2D descriptor with unit leading dims yields its
+/// inner 2.
static std::optional<SmallVector<int64_t>>
getMaybeLaneLayout(xegpu::TensorDescType tdescType) {
auto layout = tdescType.getLayoutAttr();
if (!layout)
return std::nullopt;
auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
- if (laneLayout.size() != 2)
+ if (laneLayout.size() < 2)
return std::nullopt;
- return laneLayout;
+ for (int64_t d : ArrayRef<int64_t>(laneLayout).drop_back(2))
+ if (d != 1)
+ return std::nullopt;
+ return SmallVector<int64_t>(laneLayout.end() - 2, laneLayout.end());
}
/// A layout can be optimized if its lane layout is transposed (lane[0] != 1 &&
@@ -139,15 +150,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 +174,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 +221,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 +230,14 @@ 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 transpose reshapes only the innermost 2 dims; any leading (batch) dims
+ // are unit, so their offsets pass through unchanged and they are not tiled.
+ 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 +246,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 +301,90 @@ 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());
+
+ // A fully static memref (static shape and strides) can use the op
+ // accessors directly. A dynamic memref cannot -- getMixedSizes/getMixed
+ // Strides can't represent a dynamic dim that has no SSA operand -- so its
+ // shape/strides are recovered from runtime metadata.
+ auto isStaticMemref = [](MemRefType mt) {
+ if (!mt.hasStaticShape())
+ return false;
+ SmallVector<int64_t> st;
+ int64_t off;
+ return succeeded(mt.getStridesAndOffset(st, off)) &&
+ llvm::none_of(st, ShapedType::isDynamic);
+ };
+ bool dynamicMemref = memrefType && !isStaticMemref(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 narrow elements into wider ones (f16 -> i32) reinterprets the
+ // innermost dim, so every stride *except* the innermost (which stays 1)
+ // counts the wider element and must be divided by innerLaneData. This
+ // includes the pitch (second-to-last) and, for a >2D descriptor, the
+ // leading (batch) strides -- the latter are read back by the XeVM load/
+ // store fold, which multiplies them by the *repacked* element byte size,
+ // so they must already be in repacked-element units. (For a 2D descriptor
+ // this is exactly the second-to-last stride, matching prior behavior.)
+ 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);
+
+ // The repacked shape/strides differ from the memref's own layout, so a
+ // memref source must be lowered to an i64 base pointer -- specifying
+ // explicit shape/strides on a memref create_nd is not allowed.
+ 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..6f2d34a952eeb 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,39 @@ 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 different 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 original (possibly high-D) source and only shrink the
+ // descriptor's tile shape. Leading (batch) dims are unrolled to unit tiles
+ // like any other dim; the batch position rides on the load/store/prefetch
+ // offsets (see unrollByTile) and is folded into the base pointer at XeVM
+ // lowering. We therefore never slice the source with a memref.subview,
+ // which for a dynamic-shape source could not guarantee a valid base
+ // address.
+ Value src = op.getSource();
+ auto makeCreateNd = [&](Type tdesc) -> Value {
+ auto ndTy = cast<xegpu::TensorDescType>(tdesc);
+ if (isa<MemRefType>(src.getType()))
+ // A memref source carries its own (possibly dynamic) shape/strides; use
+ // the bare-memref builder rather than getMixedSizes/getMixedStrides,
+ // which cannot represent a dynamic memref dim without an SSA operand.
+ 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) {
+ // For 2D descriptors a single tdesc is reused for every unrolled tile via
+ // offsets; the pack/unpack resolution broadcasts the one source.
+ newOps.push_back(makeCreateNd(newTdescTys[0]));
+ } else {
+ // For >2D (batched) descriptors the unpack expands the leading dims, so
+ // the source count must match the consumer's pack count. Emit one
+ // (identical, full-source) tdesc per tile; the redundant ones fold away.
+ 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 +254,22 @@ 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 are unrolled to unit tiles by unrollByTile like any
+ // other dim; their offsets ride on the prefetch and are folded into the
+ // base pointer at XeVM lowering, so no per-batch tdesc/subview is needed.
+ 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 +296,24 @@ 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 are unrolled to unit tiles by unrollByTile like any
+ // other dim; their offsets ride on the load and are folded into the base
+ // pointer at XeVM lowering, so a single tdesc is reused for all tiles.
+ 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 +343,25 @@ 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 are unrolled to unit tiles by unrollByTile like any
+ // other dim; their offsets ride on the store and are folded into the base
+ // pointer at XeVM lowering. valueIndex advances in the same tile order
+ // unrollByTile iterates, so it stays in sync with the pre-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 afea358dafe29..92eb62ce4744c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -203,10 +203,17 @@ 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 {
+ // A memref source carries its own (possibly dynamic) shape/strides; use
+ // the bare-memref builder rather than getMixedSizes/getMixedStrides,
+ // which cannot represent a dynamic memref dim without an SSA operand.
+ 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..7f32e2c37b425 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -777,9 +777,15 @@ bool xegpu::requirePacked(const xegpu::DistributeLayoutAttr layout) {
if (!layout)
return false;
auto laneData = layout.getEffectiveLaneDataAsInt();
- if (laneData.size() != 2)
+ // Packing (VNNI) applies to the innermost 2 dims. A >2D layout is accepted
+ // only when its leading (batch) dims are unit; packed iff lane_data[rank-2]
+ // (the second-to-last, "col") dim is not 1.
+ if (laneData.size() < 2)
return false;
- return laneData[0] != 1;
+ for (int64_t d : ArrayRef<int64_t>(laneData).drop_back(2))
+ if (d != 1)
+ return false;
+ return laneData[laneData.size() - 2] != 1;
}
bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
@@ -791,9 +797,16 @@ bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
if (!layout)
return false;
auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
- if (laneLayout.size() != 2)
+ // The transpose acts on the innermost 2 dims. A >2D layout is accepted only
+ // when its leading (batch) dims are unit; the inner 2 dims must be the
+ // transposed [subgroupSize, 1] form.
+ if (laneLayout.size() < 2)
return false;
- return laneLayout[0] == uArch->getSubgroupSize() && laneLayout[1] == 1;
+ for (int64_t d : ArrayRef<int64_t>(laneLayout).drop_back(2))
+ if (d != 1)
+ return false;
+ return laneLayout[laneLayout.size() - 2] == uArch->getSubgroupSize() &&
+ laneLayout[laneLayout.size() - 1] == 1;
}
// 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..b252c2a7a5cde 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>
@@ -530,15 +498,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..e262bbdc4073c 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(
diff --git a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
index 34654126ce8d2..78c9919aebb6c 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -44,23 +44,20 @@ 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: its shape/strides are recovered
+ // from the memref via extract_strided_metadata (no explicit operands).
// 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: %{{.*}}, %{{.*}}, %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[DYN]] : memref<?x?xf16>
+ // 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
}
}
diff --git a/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir b/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
index 340a10a99cb88..8397b956fd165 100644
--- a/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/array-len-op-unit.mlir
@@ -201,11 +201,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/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index f733491bfb7ee..38f7f09cfc3d6 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -81,19 +81,18 @@ 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; its shape/strides are inferred
+ // from the memref (explicit shape/strides for a memref source is deprecated).
+ // 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
}
>From 23c40e0bf03cc62b0fad7a8d9323ce48cfa0a2ee Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 11 Aug 2026 23:52:28 +0000
Subject: [PATCH 2/8] [mlir][xegpu] Add tests for batched nd descriptor
lowering + deprecation
- create_nd_tdesc_batch.mlir: check that a >2D (batched) create_nd encodes the
innermost matrix as the 2D-block surface and the leading (batch) dim element
strides into the spare payload slots (slot 5..), which the load/store fold
reads to fold batch offsets into the base pointer.
- invalid.mlir: check that specifying explicit shape/strides for a memref
source is now rejected (a memref is authoritative for its own shape/strides).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPUToXeVM/create_nd_tdesc_batch.mlir | 29 +++++++++++++++++++
mlir/test/Dialect/XeGPU/invalid.mlir | 13 +++++++++
2 files changed, 42 insertions(+)
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir
diff --git a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir
new file mode 100644
index 0000000000000..432f589060535
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
+
+// A >2D (batched) create_nd descriptor keeps the innermost matrix as the
+// 2D-block surface (base_shape_w/h and pitch come from the innermost two memref
+// dims) and encodes the leading (batch) dim element strides into the spare
+// payload slots (5..). The matching load/store lowering reads those strides to
+// fold the batch offsets into the base pointer, so the batch position stays out
+// of the surface. Here memref<4x64x128xf16> has strides [8192, 128, 1]:
+// base_shape_w = 128 (size[2]), base_shape_h = 64 (size[1]),
+// base_pitch = 128 (stride[1]), leading stride slot 5 = 8192 (stride[0]).
+gpu.module @create_nd_batch {
+ // CHECK-LABEL: gpu.func @create_nd_3d
+ gpu.func @create_nd_3d(%src: memref<4x64x128xf16>) -> vector<8xi32> {
+ // CHECK: %[[W:.+]] = arith.trunci %{{.*}} : i64 to i32
+ // CHECK: %[[H:.+]] = arith.trunci %{{.*}} : i64 to i32
+ // CHECK: %[[PITCH:.+]] = arith.trunci %{{.*}} : i64 to i32
+ // CHECK: %[[P0:.+]] = vector.insert %{{.*}}, %{{.*}} [0] : i64 into vector<4xi64>
+ // CHECK: %[[P1:.+]] = vector.bitcast %[[P0]] : vector<4xi64> to vector<8xi32>
+ // CHECK: %[[P2:.+]] = vector.insert %[[W]], %[[P1]] [2] : i32 into vector<8xi32>
+ // CHECK: %[[P3:.+]] = vector.insert %[[H]], %[[P2]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[P4:.+]] = vector.insert %[[PITCH]], %[[P3]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[C8192:.+]] = arith.constant 8192 : i64
+ // CHECK: %[[LS0:.+]] = arith.trunci %[[C8192]] : i64 to i32
+ // CHECK: %{{.+}} = vector.insert %[[LS0]], %[[P4]] [5] : i32 into vector<8xi32>
+ %t = xegpu.create_nd_tdesc %src : memref<4x64x128xf16> -> !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/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index dc68f5136b4a9..c8b39e43ef0c3 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -9,6 +9,19 @@ func.func @create_nd_tdesc_1(%src: memref<24xf32>) {
// -----
+// A memref source is authoritative for its own shape/strides, so specifying
+// them explicitly is deprecated and rejected. Use the bare form instead.
+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>
>From a6956cf45e3d6f24d75e6dbcaba8919b014ba51e Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 12 Aug 2026 02:19:51 +0000
Subject: [PATCH 3/8] [mlir][xegpu] Add lit tests for dynamic high-D nd
lowering
Positive tests for the base-pointer-fold lowering of >2D (batched)
create_nd_tdesc/load_nd/store_nd that keeps the full high-D memref as the
create_nd source (no per-batch subview) and folds leading offsets into the
base pointer. The tests are folded into the existing per-pass test files and
use dynamic-shape sources wherever the pass allows:
- VectorToXeGPU (transfer-read-to-xegpu.mlir @load_high_dim_dyn,
transfer-write-to-xegpu.mlir @store_high_dim_dyn): an equal-rank high-D
transfer on a dynamic memref keeps the whole memref as the create_nd source
-- no rank-collapsing subview -- on the nd path.
- XeGPU peephole (peephole-optimize.mlir @transpose_4d): the transpose
optimization fires on a >2D descriptor with unit leading dims over a dynamic
memref, repacking the innermost 2D tile f16->i32, preserving leading dims,
and dividing every non-innermost stride by the pack factor (the dynamic
leading stride via arith.shrui).
- SgToLaneDistribute (sg-to-lane-distribute-unit.mlir @load_nd_packed_4d /
@load_nd_transpose_4d): the packed (VNNI) and transpose attributes are set
from the innermost lane layout when leading dims are unit.
- XeGPU unroll (xegpu-unroll-patterns.mlir @load_store_nd_3d): a 3D descriptor
unrolls to N full-source descriptors (no subview) with the batch index
carried on each load/store offset (the unrolled batch dim is static by
construction).
- XeGPUToXeVM (create_nd_tdesc.mlir @create_nd_tdesc_batch_dyn,
loadstore_nd.mlir @load_store_batch_dyn): a batched create_nd on a dynamic
memref encodes the recovered leading stride into the payload, and the
matching load/store folds the batch offset into the base pointer while the
2D-block surface stays the innermost matrix.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 33 +++++++++++++
.../transfer-write-to-xegpu.mlir | 30 ++++++++++++
.../XeGPUToXeVM/create_nd_tdesc.mlir | 26 ++++++++++
.../XeGPUToXeVM/create_nd_tdesc_batch.mlir | 29 -----------
.../Conversion/XeGPUToXeVM/loadstore_nd.mlir | 48 +++++++++++++++++++
.../test/Dialect/XeGPU/peephole-optimize.mlir | 37 ++++++++++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 29 +++++++++++
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 34 +++++++++++++
8 files changed, 237 insertions(+), 29 deletions(-)
delete mode 100644 mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index b252c2a7a5cde..14bec1bbaf231 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -275,6 +275,39 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
// LOAD-GATHER: return %[[VEC]]
}
+// -----
+// A dynamic high-D read whose vector rank equals the memref rank keeps the
+// *whole* memref as the create_nd source on the nd path: no rank-collapsing
+// memref.subview is emitted and the leading (batch) offsets stay on the
+// load_nd (they are folded into the base pointer later, at XeVM lowering).
+gpu.module @xevm_module {
+gpu.func @load_high_dim_dyn(%source: memref<?x?x8x16xf16>,
+ %i: index, %j: index, %k: index, %l: index) -> vector<1x1x8x16xf16> {
+ %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<1x1x8x16xf16>
+ gpu.return %0 : vector<1x1x8x16xf16>
+}
+
+// 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<1x1x8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
+// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF0]], %[[OFF1]], %[[OFF2]], %[[OFF3]]]{{.*}}-> vector<1x1x8x16xf16>
+// 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<1x1x8x16xi1>
+// 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<1x1x8x16xindex>, vector<1x1x8x16xi1> -> vector<1x1x8x16xf16>
+// LOAD-GATHER: return %[[VEC]]
+}
+
// -----
gpu.module @xevm_module {
gpu.func @load_high_dim_vector(%source: memref<16x32x64xf32>,
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index e262bbdc4073c..71c586f0ca54d 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -91,6 +91,36 @@ gpu.func @store_dynamic_source(%vec: vector<8x16xf32>,
// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>
}
+// -----
+// A dynamic high-D write whose vector rank equals the memref rank keeps the
+// *whole* memref as the create_nd source on the nd path: no rank-collapsing
+// memref.subview is emitted and the leading (batch) offsets stay on the
+// store_nd (they are folded into the base pointer later, at XeVM lowering).
+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 78c9919aebb6c..66990487d2805 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -60,4 +60,30 @@ gpu.module @create_nd_tdesc {
%dyn_tdesc = xegpu.create_nd_tdesc %dyn : memref<?x?xf16> -> !xegpu.tensor_desc<16x16xf16>
gpu.return
}
+
+ // A >2D (batched) create_nd descriptor keeps the innermost matrix as the
+ // 2D-block surface (base_shape_w/h and pitch come from the innermost two
+ // memref dims) and encodes the leading (batch) dim element strides into the
+ // spare payload slots (5..). For a dynamic source the sizes/strides are
+ // recovered via extract_strided_metadata: base_shape_w = sizes#2,
+ // base_shape_h = sizes#1, base_pitch = strides#1, and leading stride slot 5 =
+ // strides#0. The matching load/store lowering reads those strides to fold the
+ // batch offsets into the base pointer, so the batch position stays out of the
+ // surface.
+ // 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: %[[PITCH:.+]] = arith.index_cast %[[STRIDES]]#1 : index to i32
+ // CHECK: %[[P2:.+]] = vector.insert %[[W]], %{{.+}} [2] : i32 into vector<8xi32>
+ // CHECK: %[[P3:.+]] = vector.insert %[[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: vector.insert %[[LS0]], %[[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/create_nd_tdesc_batch.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir
deleted file mode 100644
index 432f589060535..0000000000000
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc_batch.mlir
+++ /dev/null
@@ -1,29 +0,0 @@
-// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
-
-// A >2D (batched) create_nd descriptor keeps the innermost matrix as the
-// 2D-block surface (base_shape_w/h and pitch come from the innermost two memref
-// dims) and encodes the leading (batch) dim element strides into the spare
-// payload slots (5..). The matching load/store lowering reads those strides to
-// fold the batch offsets into the base pointer, so the batch position stays out
-// of the surface. Here memref<4x64x128xf16> has strides [8192, 128, 1]:
-// base_shape_w = 128 (size[2]), base_shape_h = 64 (size[1]),
-// base_pitch = 128 (stride[1]), leading stride slot 5 = 8192 (stride[0]).
-gpu.module @create_nd_batch {
- // CHECK-LABEL: gpu.func @create_nd_3d
- gpu.func @create_nd_3d(%src: memref<4x64x128xf16>) -> vector<8xi32> {
- // CHECK: %[[W:.+]] = arith.trunci %{{.*}} : i64 to i32
- // CHECK: %[[H:.+]] = arith.trunci %{{.*}} : i64 to i32
- // CHECK: %[[PITCH:.+]] = arith.trunci %{{.*}} : i64 to i32
- // CHECK: %[[P0:.+]] = vector.insert %{{.*}}, %{{.*}} [0] : i64 into vector<4xi64>
- // CHECK: %[[P1:.+]] = vector.bitcast %[[P0]] : vector<4xi64> to vector<8xi32>
- // CHECK: %[[P2:.+]] = vector.insert %[[W]], %[[P1]] [2] : i32 into vector<8xi32>
- // CHECK: %[[P3:.+]] = vector.insert %[[H]], %[[P2]] [3] : i32 into vector<8xi32>
- // CHECK: %[[P4:.+]] = vector.insert %[[PITCH]], %[[P3]] [4] : i32 into vector<8xi32>
- // CHECK: %[[C8192:.+]] = arith.constant 8192 : i64
- // CHECK: %[[LS0:.+]] = arith.trunci %[[C8192]] : i64 to i32
- // CHECK: %{{.+}} = vector.insert %[[LS0]], %[[P4]] [5] : i32 into vector<8xi32>
- %t = xegpu.create_nd_tdesc %src : memref<4x64x128xf16> -> !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/loadstore_nd.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
index d45fa79bb2e63..1fde43119d8cf 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
@@ -91,4 +91,52 @@ gpu.module @load_store_check {
vector.store %loaded, %dstte[%c0, %c0] : memref<32x16xi8>, vector<32xi8>
gpu.return
}
+
+ // A >2D (batched) load_nd / store_nd folds the leading (batch) offsets into
+ // the base pointer -- base += (sum_d offset[d] * leadingStride[d]) *
+ // elemByteSize -- while the 2D-block surface stays the innermost matrix. For
+ // a dynamic-shape source the sizes/strides are recovered via
+ // extract_strided_metadata: the leading (batch) element stride is carried in
+ // payload slot 5, folding batch index %z contributes %z * stride#0 * 4 bytes
+ // to the base, and the surface is the innermost sizes#2 x sizes#1 matrix.
+ // 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
+ // CHECK-DAG: %[[C4I64:.+]] = arith.constant 4 : i64
+
+ // Load: recover shape/strides, encode the leading stride into slot 5, and
+ // fold the batch index into the base pointer.
+ // CHECK: %{{.+}}, %{{.+}}, %[[LSZ:.+]]:3, %[[LSTR:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
+ // CHECK: %[[LSTRIDE:.+]] = arith.index_cast %[[LSTR]]#0 : index to i32
+ // CHECK: vector.insert %[[LSTRIDE]], %{{.+}} [5] : i32 into vector<8xi32>
+ // CHECK: %[[LBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
+ // CHECK: %[[LZ:.+]] = arith.index_cast %[[Z]] : index to i64
+ // CHECK: %[[LSE:.+]] = arith.extui %[[LSTRIDE]] : i32 to i64
+ // CHECK: %[[LMUL:.+]] = arith.muli %[[LZ]], %[[LSE]] : i64
+ // CHECK: %[[LOFF:.+]] = arith.muli %[[LMUL]], %[[C4I64]] : i64
+ // CHECK: %[[LADDR:.+]] = arith.addi %[[LBASE]], %[[LOFF]] : i64
+ // CHECK: %[[LPTR:.+]] = llvm.inttoptr %[[LADDR]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.blockload2d %[[LPTR]]{{.*}}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 batch fold on the base.
+ // CHECK: %{{.+}}, %{{.+}}, %[[SSZ:.+]]:3, %[[SSTR:.+]]:3 = memref.extract_strided_metadata %[[DST]]
+ // CHECK: %[[SSTRIDE:.+]] = arith.index_cast %[[SSTR]]#0 : index to i32
+ // CHECK: vector.insert %[[SSTRIDE]], %{{.+}} [5] : i32 into vector<8xi32>
+ // CHECK: %[[SBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
+ // CHECK: %[[SZ:.+]] = arith.index_cast %[[Z]] : index to i64
+ // CHECK: %[[SSE:.+]] = arith.extui %[[SSTRIDE]] : i32 to i64
+ // CHECK: %[[SMUL:.+]] = arith.muli %[[SZ]], %[[SSE]] : i64
+ // CHECK: %[[SOFF:.+]] = arith.muli %[[SMUL]], %[[C4I64]] : i64
+ // CHECK: %[[SADDR:.+]] = arith.addi %[[SBASE]], %[[SOFF]] : i64
+ // CHECK: %[[SPTR:.+]] = llvm.inttoptr %[[SADDR]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.blockstore2d %[[SPTR]]{{.*}}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/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index fa5ff9ca56d34..9ccbb0de861ad 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -504,3 +504,40 @@ gpu.module @xevm_test {
gpu.return
}
}
+
+// -----
+// The transpose optimization also fires for a >2D descriptor whose leading
+// (batch) dims are unit: it operates on the innermost two dims, repacking
+// f16->i32 for the HW transpose load. The leading unit dims are preserved on
+// the descriptor/result, lane_data is rank-matched to all ones, and every
+// non-innermost stride is divided by the pack factor (2) so the later XeVM
+// batch fold uses repacked-element units. For a dynamic-shape source the
+// sizes/strides are recovered via extract_strided_metadata; the dynamic
+// leading (batch) stride is divided by the pack factor with arith.shrui, while
+// the static inner strides fold to constants (4096 -> 2048, 64 -> 32).
+// 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 a7fe2d9534351..e95e6788ab5e6 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,35 @@ gpu.func @load_nd_transpose() {
gpu.return
}
+// A >2D load with unit leading dims is treated as its innermost 2D tile: the
+// packed attribute is still set when the inner lane_data marks a packed
+// (VNNI) B operand.
+// 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
+}
+
+// A >2D load with unit leading dims is likewise transposed on its innermost 2D
+// tile when the inner lane layout requires it.
+// 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 49fdf1cbee174..dafc4d7339a27 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -360,5 +360,39 @@ gpu.module @test {
gpu.return %0 : vector<32xf32>
}
+//-----
+ // Unrolling a >2D (batched) nd descriptor keeps the whole memref as the
+ // create_nd source -- no memref.subview -- and yields a single unit-leading
+ // tdesc that is reused across the batch tiles. The leading (batch) offset
+ // stays on each unrolled load/store and is incremented per tile (it is *not*
+ // zeroed and baked into a subview base). inst_data [1, 8, 16] unrolls the
+ // size-4 batch dim into 4 tiles.
+ // 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
+ }
+
}
>From 5a5b993f0fb6ba5f7649cdf9e960bb93bd998e71 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 12 Aug 2026 03:14:12 +0000
Subject: [PATCH 4/8] [mlir][xegpu] Trim comments and simplify dynamic high-D
nd lowering
Review cleanup of the batched/high-D nd lowering: reduce verbose in-function
comments to the genuinely non-obvious motivation (no memref.subview for a
dynamic source, dividing all non-innermost strides on repack, 1-vs-N unrolled
tdescs, the base-pointer batch-fold formula), and drop mechanical narration.
Also replace the lazy getSrcMeta lambda in CreateNdDescToXeVMPattern with a
plain local -- it was only used on the dynamic-memref path -- matching the
metadata-recovery pattern used elsewhere. No functional change.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 55 ++++++-------------
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 18 +++---
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 42 ++++++--------
3 files changed, 42 insertions(+), 73 deletions(-)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index e30951fdb9e1d..269360cafe0d6 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -56,10 +56,8 @@ enum class NdTdescOffset : uint32_t {
BaseShapeH = 3, // Base shape height (i32)
BasePitch = 4, // Base pitch/stride of dim rank-2 (i32)
LeadingStride0 = 5, // Element strides of the leading (batch) dims of a >2D
- LeadingStride1 = 6, // descriptor (i32). The load/store/prefetch lowering
- LeadingStride2 = 7, // uses these to fold the batch offsets into the base
- // pointer, keeping the 2D-block surface at the innermost
- // matrix. Left at 0 for 2D descriptors.
+ LeadingStride1 = 6, // descriptor (i32); folded into the base pointer by the
+ LeadingStride2 = 7, // load/store lowering. Left at 0 for 2D descriptors.
};
static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
@@ -240,20 +238,11 @@ class CreateNdDescToXeVMPattern
auto sourceTy = source.getType();
auto sourceMemrefTy = dyn_cast<MemRefType>(sourceTy);
- // Shape and strides. For a memref source they are recovered from the memref
- // value itself (a dynamic memref carries them at runtime and does not
- // attach them as op operands); otherwise they come from the op's explicit
- // operands.
+ // For a memref source, shape/strides come from the memref (dynamic dims are
+ // recovered from runtime metadata); an integer source carries them as
+ // explicit op operands.
SmallVector<OpFoldResult> mixedSizes;
SmallVector<OpFoldResult> mixedStrides;
- memref::ExtractStridedMetadataOp srcMeta;
- // Lazily materialize memref metadata (base buffer, offset, sizes, strides).
- auto getSrcMeta = [&]() -> memref::ExtractStridedMetadataOp {
- if (!srcMeta)
- srcMeta =
- memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
- return srcMeta;
- };
// 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) {
@@ -265,18 +254,18 @@ class CreateNdDescToXeVMPattern
if (failed(
sourceMemrefTy.getStridesAndOffset(staticStrides, staticOffset)))
return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
- // A memref is authoritative for its own shape/strides (explicit
- // shape/strides on a memref create_nd are deprecated). A fully static
- // memref yields constants directly; a dynamic one recovers its dynamic
- // dims from runtime metadata (getConstified* keeps static dims as attrs).
+ // A fully static memref yields constants directly; a dynamic one recovers
+ // its dynamic dims from runtime metadata.
bool allStatic = sourceMemrefTy.hasStaticShape() &&
llvm::none_of(staticStrides, ShapedType::isDynamic);
if (allStatic) {
mixedSizes = op.getMixedSizes();
mixedStrides = op.getMixedStrides();
} else {
- mixedSizes = getSrcMeta().getConstifiedMixedSizes();
- mixedStrides = getSrcMeta().getConstifiedMixedStrides();
+ auto srcMeta =
+ memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
+ mixedSizes = srcMeta.getConstifiedMixedSizes();
+ mixedStrides = srcMeta.getConstifiedMixedStrides();
}
// Access adaptor after failure check to avoid rolling back generated code
// for materialization cast.
@@ -327,15 +316,9 @@ class CreateNdDescToXeVMPattern
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
// For a >2D descriptor, encode the leading (batch) dim element strides into
- // the spare payload slots. The load/store/prefetch lowering folds the batch
- // offsets into the base pointer using these, so the batch position stays
- // out of the 2D-block surface (which remains the innermost matrix computed
- // above). This applies to both memref sources and integer (pointer) sources
- // with explicit shape/strides -- e.g. a batched descriptor whose source was
- // rewritten to an i64 base by the transpose peephole; in that case the
- // strides come from the op's explicit operands. 2D descriptors leave these
- // slots at 0 and the load/store path skips reading them (2D path
- // unchanged).
+ // the spare payload slots; the load/store/prefetch lowering folds the batch
+ // offsets into the base pointer with these, keeping the 2D-block surface at
+ // the innermost matrix. 2D descriptors leave these slots at 0.
if (rank > 2) {
if (rank - 2 > 3)
return rewriter.notifyMatchFailure(
@@ -456,12 +439,10 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
mixedOffsets[tileRank - 2]);
offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), offsetH);
- // For a >2D descriptor, fold the leading (batch) offsets into the base
- // pointer: basePtr += (sum_d offset[d] * leadingStride[d]) * elemBytes.
- // The batch element strides were encoded into the payload at create time
- // (see CreateNdDescToXeVMPattern). This keeps the 2D-block surface at the
- // innermost matrix (so the HW surface-size limits are unaffected) instead
- // of baking the batch into a (possibly out-of-range) collapsed surface.
+ // Fold the leading (batch) offsets into the base pointer:
+ // basePtr += (sum_d offset[d] * leadingStride[d]) * elemBytes
+ // using the batch strides encoded at create time. This keeps the 2D-block
+ // surface at the innermost matrix, avoiding the HW surface-size limits.
if (tileRank > 2) {
Type i64Ty = rewriter.getI64Type();
Value batchElemOffset;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 4b7ee133006e5..0124b700ac1dc 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -305,10 +305,9 @@ class XeGPUCreateNdDescOpPattern final
Value source = createNdOp.getSource();
auto memrefType = dyn_cast<MemRefType>(source.getType());
- // A fully static memref (static shape and strides) can use the op
- // accessors directly. A dynamic memref cannot -- getMixedSizes/getMixed
- // Strides can't represent a dynamic dim that has no SSA operand -- so its
- // shape/strides are recovered from runtime metadata.
+ // A dynamic memref's shape/strides are recovered from runtime metadata
+ // (getMixedSizes/getMixedStrides can't represent a dynamic dim); a static
+ // one uses the op accessors directly.
auto isStaticMemref = [](MemRefType mt) {
if (!mt.hasStaticShape())
return false;
@@ -346,13 +345,10 @@ class XeGPUCreateNdDescOpPattern final
rewriter, loc, convertToValue(rewriter, loc, modifiedShape.back()),
innerLaneData);
// Repacking narrow elements into wider ones (f16 -> i32) reinterprets the
- // innermost dim, so every stride *except* the innermost (which stays 1)
- // counts the wider element and must be divided by innerLaneData. This
- // includes the pitch (second-to-last) and, for a >2D descriptor, the
- // leading (batch) strides -- the latter are read back by the XeVM load/
- // store fold, which multiplies them by the *repacked* element byte size,
- // so they must already be in repacked-element units. (For a 2D descriptor
- // this is exactly the second-to-last stride, matching prior behavior.)
+ // innermost dim, so every stride except the innermost (which stays 1) must
+ // be divided by innerLaneData to count the wider element -- the pitch and,
+ // for a >2D descriptor, the leading (batch) strides (the XeVM fold reads
+ // the latter back in repacked-element units).
assert(mixedStrides.size() >= 2 &&
"Expected at least 2 strides for CreateNdDescOp");
SmallVector<OpFoldResult> modifiedStrides(mixedStrides);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 6f2d34a952eeb..f14df6de90c20 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -200,20 +200,17 @@ struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
if (!targetShape)
return failure();
- // Keep the original (possibly high-D) source and only shrink the
- // descriptor's tile shape. Leading (batch) dims are unrolled to unit tiles
- // like any other dim; the batch position rides on the load/store/prefetch
- // offsets (see unrollByTile) and is folded into the base pointer at XeVM
- // lowering. We therefore never slice the source with a memref.subview,
- // which for a dynamic-shape source could not guarantee a valid base
- // address.
+ // Keep the original (possibly high-D) source and only shrink the tile
+ // shape; batch dims unroll to unit tiles and their offsets ride on the
+ // load/store/prefetch, folded into the base pointer at XeVM lowering. We
+ // never slice the source with a memref.subview -- for a dynamic-shape
+ // source it could not guarantee a valid base address.
Value src = op.getSource();
auto makeCreateNd = [&](Type tdesc) -> Value {
auto ndTy = cast<xegpu::TensorDescType>(tdesc);
+ // A memref source uses the bare-memref builder: a dynamic dim has no SSA
+ // operand for getMixedSizes/getMixedStrides.
if (isa<MemRefType>(src.getType()))
- // A memref source carries its own (possibly dynamic) shape/strides; use
- // the bare-memref builder rather than getMixedSizes/getMixedStrides,
- // which cannot represent a dynamic memref dim without an SSA operand.
return xegpu::CreateNdDescOp::create(rewriter, loc, ndTy,
cast<TypedValue<MemRefType>>(src));
return xegpu::CreateNdDescOp::create(
@@ -223,13 +220,11 @@ struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
SmallVector<Type> newTdescTys = getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> newOps;
if (tdescTy.getRank() <= 2) {
- // For 2D descriptors a single tdesc is reused for every unrolled tile via
- // offsets; the pack/unpack resolution broadcasts the one source.
+ // 2D: one tdesc, broadcast across tiles by pack/unpack.
newOps.push_back(makeCreateNd(newTdescTys[0]));
} else {
- // For >2D (batched) descriptors the unpack expands the leading dims, so
- // the source count must match the consumer's pack count. Emit one
- // (identical, full-source) tdesc per tile; the redundant ones fold away.
+ // >2D: unpack expands the leading dims, so the source count must match
+ // the pack count. Emit one (identical) tdesc per tile; CSE folds them.
for (Type t : newTdescTys)
newOps.push_back(makeCreateNd(t));
}
@@ -254,9 +249,8 @@ struct UnrollPrefetchNdOp : public UnrollPattern<xegpu::PrefetchNdOp> {
if (layout)
layout = layout.dropInstData();
- // Batch (leading) dims are unrolled to unit tiles by unrollByTile like any
- // other dim; their offsets ride on the prefetch and are folded into the
- // base pointer at XeVM lowering, so no per-batch tdesc/subview is needed.
+ // Batch (leading) dims unroll to unit tiles like any other dim; a single
+ // tdesc serves all tiles (offsets are folded into the base at XeVM).
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdesc = pack(
@@ -298,9 +292,8 @@ struct UnrollLoadNdOp : public UnrollPattern<xegpu::LoadNdOp> {
SmallVector<Value> newOps;
- // Batch (leading) dims are unrolled to unit tiles by unrollByTile like any
- // other dim; their offsets ride on the load and are folded into the base
- // pointer at XeVM lowering, so a single tdesc is reused for all tiles.
+ // Batch (leading) dims unroll to unit tiles like any other dim; a single
+ // tdesc serves all tiles (offsets are folded into the base at XeVM).
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdescs = pack(
@@ -345,10 +338,9 @@ struct UnrollStoreNdOp : public UnrollPattern<xegpu::StoreNdOp> {
size_t valueIndex = 0;
- // Batch (leading) dims are unrolled to unit tiles by unrollByTile like any
- // other dim; their offsets ride on the store and are folded into the base
- // pointer at XeVM lowering. valueIndex advances in the same tile order
- // unrollByTile iterates, so it stays in sync with the pre-packed values.
+ // 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(
>From e7bfc651b24979fdc97408b93c3536cdf5de3102 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 12 Aug 2026 17:07:56 +0000
Subject: [PATCH 5/8] [mlir][xegpu] Update create_nd_tdesc docs for the memref
deprecation
Document that a memref source (static or dynamic) infers its shape/strides from
the memref and must not pass them explicitly; only a pointer (uint64_t) source
supplies shape/strides. Update the op description example to the bare dynamic-
memref form, and relabel the pointer example accordingly.
Also switch the @load_high_dim_dyn VectorToXeGPU test to non-unit (2x4) batch
dims so the equal-rank no-subview lowering is unambiguous, and drop a stale
comment tail in createNdDescriptor.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 30 +++++++++----------
.../VectorToXeGPU/VectorToXeGPU.cpp | 5 +---
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 19 ++++++------
3 files changed, 26 insertions(+), 28 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 49b98922cee4c..e6bb8311e84ad 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/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 083c63198733b..6367d10b665a9 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -136,10 +136,7 @@ static xegpu::CreateNdDescOp createNdDescriptor(PatternRewriter &rewriter,
// Keep the memref itself as the source (static or dynamic) rather than
// collapsing it to an i64 base address. The memref value carries the base
// pointer *and* the (possibly dynamic) offset/shape/strides, which the XeGPU
- // -> XeVM lowering recovers via memref metadata. In particular, for a >2D
- // descriptor the leading (batch) offsets stay on the load/store and are
- // folded into the base pointer at lowering time, so a dynamic-shape source
- // never needs a per-batch subview (which could not guarantee a valid base).
+ // -> XeVM lowering recovers via memref metadata.
return xegpu::CreateNdDescOp::create(rewriter, loc, descType, src);
}
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 14bec1bbaf231..875086f78750e 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -278,33 +278,34 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
// -----
// A dynamic high-D read whose vector rank equals the memref rank keeps the
// *whole* memref as the create_nd source on the nd path: no rank-collapsing
-// memref.subview is emitted and the leading (batch) offsets stay on the
-// load_nd (they are folded into the base pointer later, at XeVM lowering).
+// memref.subview is emitted and the leading (batch) offsets (here the non-unit
+// 2x4 dims) stay on the load_nd. They are folded into the base pointer later,
+// at XeVM lowering, once blocking has unrolled them to unit tiles.
gpu.module @xevm_module {
gpu.func @load_high_dim_dyn(%source: memref<?x?x8x16xf16>,
- %i: index, %j: index, %k: index, %l: index) -> vector<1x1x8x16xf16> {
+ %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<1x1x8x16xf16>
- gpu.return %0 : vector<1x1x8x16xf16>
+ : 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<1x1x8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF0]], %[[OFF1]], %[[OFF2]], %[[OFF3]]]{{.*}}-> vector<1x1x8x16xf16>
+// 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<1x1x8x16xi1>
+// 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<1x1x8x16xindex>, vector<1x1x8x16xi1> -> vector<1x1x8x16xf16>
+// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[PTR_I]]{{\[}}%{{.+}}{{\]}}, %[[CST]] : i64, vector<2x4x8x16xindex>, vector<2x4x8x16xi1> -> vector<2x4x8x16xf16>
// LOAD-GATHER: return %[[VEC]]
}
>From 2360676b1cd219e864a73e6e837bb7656360f70e Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 22 Aug 2026 05:21:11 +0000
Subject: [PATCH 6/8] [MLIR][XeGPU] Key nd descriptor lowering on the
descriptor rank
CreateNdDescToXeVMPattern derived the rank from the source shape, but the
TensorDescType converter keys the lowered payload on the descriptor rank.
A descriptor whose rank differed from its source's therefore built a
vector<8xi32> payload where an index was expected, or folded a leading
stride taken from the wrong source dim into the base pointer, both with
no diagnostic.
- Take the rank from the descriptor type, and require it to match the
source rank (subview the source instead).
- Hoist every failure check above the first op creation, so a bail-out
leaves no dangling IR or materialization cast to roll back.
- Reject descriptors with non-unit leading dims -- a 2D-block op
transfers the innermost 2 dims only, so tile_height/tile_width would
silently ignore them -- and ranks exceeding the 3 spare payload slots.
- Factor the duplicated inner-2D-layout and static-shape/strides
predicates into xegpu::getInner2DIfUnitLeadingDims and
xegpu::hasStaticShapeAndStrides.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 13 +-
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 157 +++++++++---------
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 34 +---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 44 ++---
.../XeGPUToXeVM/create_nd_tdesc.mlir | 2 +-
.../XeGPUToXeVM/failed_conversion.mlir | 27 +++
6 files changed, 151 insertions(+), 126 deletions(-)
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/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 269360cafe0d6..97c9291387142 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -60,6 +60,10 @@ enum class NdTdescOffset : uint32_t {
LeadingStride2 = 7, // load/store 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:
@@ -217,12 +221,67 @@ 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) +
+ ").");
+
+ 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();
+ }
+
+ // getMixedSizes/getMixedStrides cannot represent a dynamic memref dim, so
+ // recover those from runtime metadata.
+ 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.
@@ -230,62 +289,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).
- auto sourceTy = source.getType();
- auto sourceMemrefTy = dyn_cast<MemRefType>(sourceTy);
-
- // For a memref source, shape/strides come from the memref (dynamic dims are
- // recovered from runtime metadata); an integer source carries them as
- // explicit op operands.
- SmallVector<OpFoldResult> mixedSizes;
- SmallVector<OpFoldResult> mixedStrides;
- // 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.");
- }
- SmallVector<int64_t> staticStrides;
- int64_t staticOffset;
- if (failed(
- sourceMemrefTy.getStridesAndOffset(staticStrides, staticOffset)))
- return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
- // A fully static memref yields constants directly; a dynamic one recovers
- // its dynamic dims from runtime metadata.
- bool allStatic = sourceMemrefTy.hasStaticShape() &&
- llvm::none_of(staticStrides, ShapedType::isDynamic);
- if (allStatic) {
- mixedSizes = op.getMixedSizes();
- mixedStrides = op.getMixedStrides();
- } else {
- auto srcMeta =
- memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
- mixedSizes = srcMeta.getConstifiedMixedSizes();
- mixedStrides = srcMeta.getConstifiedMixedStrides();
- }
- // Access adaptor after failure check to avoid rolling back generated code
- // for materialization cast.
- baseAddr = adaptor.getSource();
- } else {
- mixedSizes = op.getMixedSizes();
- mixedStrides = op.getMixedStrides();
- 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);
- }
- }
- // Descriptor shape rank.
- int64_t rank = mixedSizes.size();
- // 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 {
@@ -293,10 +296,9 @@ 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);
+ Value baseShapeH = createOffset(mixedSizes, rank - 2);
// Pitch is the stride of dim rank-2 (the row stride of the 2D tile).
Value basePitch = createOffset(mixedStrides, rank - 2);
// Populate payload.
@@ -315,20 +317,13 @@ class CreateNdDescToXeVMPattern
payload =
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
- // For a >2D descriptor, encode the leading (batch) dim element strides into
- // the spare payload slots; the load/store/prefetch lowering folds the batch
- // offsets into the base pointer with these, keeping the 2D-block surface at
- // the innermost matrix. 2D descriptors leave these slots at 0.
- if (rank > 2) {
- if (rank - 2 > 3)
- return rewriter.notifyMatchFailure(
- op, "Batched nd descriptor supports at most 3 leading dims.");
- for (int64_t d = 0; d < rank - 2; ++d) {
- Value leadingStride = createOffset(mixedStrides, d);
- payload = vector::InsertOp::create(
- rewriter, loc, leadingStride, payload,
- static_cast<int>(NdTdescOffset::LeadingStride0) + d);
- }
+ // Leading (batch) strides go into the spare payload slots; the load/store/
+ // prefetch lowering folds the batch offsets into the base pointer with it.
+ for (int64_t d = 0; d < rank - 2; ++d) {
+ Value leadingStride = createOffset(mixedStrides, d);
+ payload = vector::InsertOp::create(
+ rewriter, loc, leadingStride, payload,
+ static_cast<int>(NdTdescOffset::LeadingStride0) + d);
}
rewriter.replaceOp(op, payload);
return success();
@@ -355,6 +350,16 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
if (opOffsetsSize != tileRank)
return rewriter.notifyMatchFailure(
op, "Expected offset rank to match descriptor rank.");
+ // A 2D-block op transfers the innermost 2 dims only, so a leading dim must
+ // be a unit tile whose offset can be folded into the base pointer.
+ 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;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 0124b700ac1dc..7e609be3de4a3 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -53,13 +53,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;
- for (int64_t d : ArrayRef<int64_t>(laneData).drop_back(2))
- if (d != 1)
- return std::nullopt;
- return SmallVector<int64_t>(laneData.end() - 2, laneData.end());
+ return xegpu::getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt());
}
/// Get the 2D lane layout from a tensor desc type if it exists. As with
@@ -70,13 +64,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;
- for (int64_t d : ArrayRef<int64_t>(laneLayout).drop_back(2))
- if (d != 1)
- return std::nullopt;
- return SmallVector<int64_t>(laneLayout.end() - 2, laneLayout.end());
+ return xegpu::getInner2DIfUnitLeadingDims(
+ layout.getEffectiveLaneLayoutAsInt());
}
/// A layout can be optimized if its lane layout is transposed (lane[0] != 1 &&
@@ -305,19 +294,10 @@ class XeGPUCreateNdDescOpPattern final
Value source = createNdOp.getSource();
auto memrefType = dyn_cast<MemRefType>(source.getType());
- // A dynamic memref's shape/strides are recovered from runtime metadata
- // (getMixedSizes/getMixedStrides can't represent a dynamic dim); a static
- // one uses the op accessors directly.
- auto isStaticMemref = [](MemRefType mt) {
- if (!mt.hasStaticShape())
- return false;
- SmallVector<int64_t> st;
- int64_t off;
- return succeeded(mt.getStridesAndOffset(st, off)) &&
- llvm::none_of(st, ShapedType::isDynamic);
- };
- bool dynamicMemref = memrefType && !isStaticMemref(memrefType);
-
+ // getMixedSizes/getMixedStrides cannot represent a dynamic memref dim, so
+ // recover those from runtime metadata.
+ bool dynamicMemref =
+ memrefType && !xegpu::hasStaticShapeAndStrides(memrefType);
SmallVector<OpFoldResult> mixedSizes;
SmallVector<OpFoldResult> mixedStrides;
memref::ExtractStridedMetadataOp meta;
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 7f32e2c37b425..76269cf193d13 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -773,19 +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();
- // Packing (VNNI) applies to the innermost 2 dims. A >2D layout is accepted
- // only when its leading (batch) dims are unit; packed iff lane_data[rank-2]
- // (the second-to-last, "col") dim is not 1.
- if (laneData.size() < 2)
- return false;
- for (int64_t d : ArrayRef<int64_t>(laneData).drop_back(2))
- if (d != 1)
- return false;
- return laneData[laneData.size() - 2] != 1;
+ auto laneData =
+ getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt());
+ return laneData && (*laneData)[0] != 1;
}
bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
@@ -796,17 +798,19 @@ bool xegpu::requireTranspose(const xegpu::DistributeLayoutAttr layout,
return false;
if (!layout)
return false;
- auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
- // The transpose acts on the innermost 2 dims. A >2D layout is accepted only
- // when its leading (batch) dims are unit; the inner 2 dims must be the
- // transposed [subgroupSize, 1] form.
- 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;
- for (int64_t d : ArrayRef<int64_t>(laneLayout).drop_back(2))
- if (d != 1)
- return false;
- return laneLayout[laneLayout.size() - 2] == uArch->getSubgroupSize() &&
- laneLayout[laneLayout.size() - 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/XeGPUToXeVM/create_nd_tdesc.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
index 66990487d2805..0f1da6047ef59 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -46,8 +46,8 @@ gpu.module @create_nd_tdesc {
// A dynamic memref uses the bare form: its shape/strides are recovered
// from the memref via extract_strided_metadata (no explicit operands).
- // CHECK: %[[CST_3:.*]] = arith.constant dense<0> : vector<8xi32>
// 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 %[[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
diff --git a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
index cabc65aa0e41d..9b06e80d35eda 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -28,3 +28,30 @@ gpu.module @test_kernel {
gpu.return
}
}
+
+// -----
+
+// A 2D-block op transfers the innermost 2 dims only, so a non-unit leading dim
+// must be rejected rather than lowered to a load of the wrong shape.
+
+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
+ }
+}
>From a5884261cd219a17c70210f70d06c1aed84d6867 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 22 Aug 2026 18:42:23 +0000
Subject: [PATCH 7/8] [mlir][xegpu] Carry batch offsets in offset_h instead of
the base pointer
A >2D nd descriptor was lowered by keeping the 2D-block surface at one
innermost matrix and folding the leading (batch) offsets into the base
pointer. That aims the surface at unmapped memory when a batch index is out
of range, which the HW boundary check cannot detect.
View the source as a single flattened 2D plane instead: base_height becomes
the product of all dims but the innermost, so the surface covers every
leading (batch) plane, and a batch position becomes a row offset into it.
base_ptr stays at the true base, so an out-of-range batch index lands past
base_height where the boundary check handles it.
The leading strides are now encoded as row counts (stride[d] / stride[R-2])
rather than element counts. Being a ratio, they are dimensionless and cannot
drift out of step with the pitch under element-type repacking (e.g. the
f16 -> i32 transpose repack), which removes a class of unit-mismatch bug.
This also drops the i64 extend and the elemByteSize multiply from the fold.
The flattened-plane view has three limitations, documented above
CreateNdDescToXeVMPattern: each leading stride must be a whole number of
rows (rejected when statically known otherwise, for both memref and integer
sources); base_height grows to the product of the leading dims and can
exceed the HW surface height; and plane boundaries are invisible to the
boundary check, so a tile whose rows run past size[R-2] reads the next
plane instead of zero padding.
Also trims stale and redundant comments across the XeGPU passes and tests,
including several that still described the base-pointer fold.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 6 +-
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 136 ++++++++++++------
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 11 --
.../XeGPUArrayLengthOptimization.cpp | 5 -
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 25 +---
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 17 +--
.../Transforms/XeGPUWgToSgDistribute.cpp | 3 -
.../VectorToXeGPU/transfer-read-to-xegpu.mlir | 6 +-
.../transfer-write-to-xegpu.mlir | 5 +-
.../XeGPUToXeVM/create_nd_tdesc.mlir | 20 +--
.../XeGPUToXeVM/failed_conversion.mlir | 27 +++-
.../Conversion/XeGPUToXeVM/loadstore_nd.mlir | 48 +++----
mlir/test/Dialect/XeGPU/invalid.mlir | 3 +-
mlir/test/Dialect/XeGPU/ops.mlir | 3 +-
.../test/Dialect/XeGPU/peephole-optimize.mlir | 10 +-
.../XeGPU/sg-to-lane-distribute-unit.mlir | 7 +-
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 7 +-
17 files changed, 167 insertions(+), 172 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 6367d10b665a9..1f6efdc091a75 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -133,10 +133,8 @@ static xegpu::CreateNdDescOp createNdDescriptor(PatternRewriter &rewriter,
[[maybe_unused]] MemRefType srcTy = src.getType();
assert(srcTy.isStrided() && "Expected strided memref type");
- // Keep the memref itself as the source (static or dynamic) rather than
- // collapsing it to an i64 base address. The memref value carries the base
- // pointer *and* the (possibly dynamic) offset/shape/strides, which the XeGPU
- // -> XeVM lowering recovers via memref metadata.
+ // Keep the memref as the source instead of collapsing it to an i64 base: it
+ // carries the (possibly dynamic) offset/shape/strides the lowering recovers.
return xegpu::CreateNdDescOp::create(rewriter, loc, descType, src);
}
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 5a2c67183a8a0..f2121d5178cf9 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -55,9 +55,9 @@ enum class NdTdescOffset : uint32_t {
BaseShapeW = 2, // Base shape width (i32)
BaseShapeH = 3, // Base shape height (i32)
BasePitch = 4, // Base pitch/stride of dim rank-2 (i32)
- LeadingStride0 = 5, // Element strides of the leading (batch) dims of a >2D
- LeadingStride1 = 6, // descriptor (i32); folded into the base pointer by the
- LeadingStride2 = 7, // load/store lowering. Left at 0 for 2D descriptors.
+ 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.
@@ -212,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;
@@ -247,6 +272,34 @@ class CreateNdDescToXeVMPattern
std::to_string(maxNdTdescLeadingDims) +
" leading dims (rank <= " + std::to_string(maxNdTdescRank) +
").");
+ // The flattened-plane view needs every leading stride to be a whole number
+ // of rows. Static strides are checked here; dynamic ones are assumed to
+ // divide evenly (limitation 1 above).
+ 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();
@@ -264,8 +317,6 @@ class CreateNdDescToXeVMPattern
return success();
}
- // getMixedSizes/getMixedStrides cannot represent a dynamic memref dim, so
- // recover those from runtime metadata.
SmallVector<OpFoldResult> mixedSizes;
SmallVector<OpFoldResult> mixedStrides;
if (memrefTy && !xegpu::hasStaticShapeAndStrides(memrefTy)) {
@@ -298,7 +349,13 @@ class CreateNdDescToXeVMPattern
};
// 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.
@@ -317,12 +374,24 @@ class CreateNdDescToXeVMPattern
payload =
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
- // Leading (batch) strides go into the spare payload slots; the load/store/
- // prefetch lowering folds the batch offsets into the base pointer with it.
+ // 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) {
- Value leadingStride = createOffset(mixedStrides, 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, leadingStride, payload,
+ rewriter, loc, leadingRowStride, payload,
static_cast<int>(NdTdescOffset::LeadingStride0) + d);
}
rewriter.replaceOp(op, payload);
@@ -350,8 +419,6 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
if (opOffsetsSize != tileRank)
return rewriter.notifyMatchFailure(
op, "Expected offset rank to match descriptor rank.");
- // A 2D-block op transfers the innermost 2 dims only, so a leading dim must
- // be a unit tile whose offset can be folded into the base pointer.
if (tileRank > 2 && llvm::any_of(tdescTy.getShape().drop_back(2),
[](int64_t d) { return d != 1; }))
return rewriter.notifyMatchFailure(
@@ -434,8 +501,6 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
Value basePitch = vector::ExtractOp::create(
rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BasePitch));
- // x = innermost offset; y = second-to-last offset. These address the 2D
- // tile within the innermost-matrix surface.
Value offsetW = getValueOrCreateConstantIntOp(rewriter, loc,
mixedOffsets[tileRank - 1]);
offsetW = getValueOrCreateCastToIndexLike(rewriter, loc,
@@ -444,35 +509,22 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
mixedOffsets[tileRank - 2]);
offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
rewriter.getI32Type(), offsetH);
- // Fold the leading (batch) offsets into the base pointer:
- // basePtr += (sum_d offset[d] * leadingStride[d]) * elemBytes
- // using the batch strides encoded at create time. This keeps the 2D-block
- // surface at the innermost matrix, avoiding the HW surface-size limits.
- if (tileRank > 2) {
- Type i64Ty = rewriter.getI64Type();
- Value batchElemOffset;
- for (int64_t d = 0; d < tileRank - 2; ++d) {
- Value off =
- getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[d]);
- off = getValueOrCreateCastToIndexLike(rewriter, loc, i64Ty, off);
- Value strideI32 = vector::ExtractOp::create(
- rewriter, loc, tdesc,
- static_cast<int>(NdTdescOffset::LeadingStride0) + d);
- Value stride =
- arith::ExtUIOp::create(rewriter, loc, i64Ty, strideI32);
- Value term = arith::MulIOp::create(rewriter, loc, off, stride);
- batchElemOffset =
- batchElemOffset
- ? arith::AddIOp::create(rewriter, loc, batchElemOffset, term)
- .getResult()
- : term;
- }
- Value elemByteSizeI64 =
- arith::ConstantIntOp::create(rewriter, loc, i64Ty, elemBitSize / 8);
- Value batchByteOffset = arith::MulIOp::create(
- rewriter, loc, batchElemOffset, elemByteSizeI64);
- basePtr =
- arith::AddIOp::create(rewriter, loc, basePtr, batchByteOffset);
+ // 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 =
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 9d1624d054bbc..a906b8cbe8802 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) {
- // A dynamic-shape memref is allowed: its (possibly dynamic) shape/strides are
- // carried by the memref value itself and recovered at lowering time via
- // memref metadata, so no explicit shape/stride operands are attached here.
build(builder, state, tdesc, source, ValueRange({}) /* empty dynamic shape */,
ValueRange({}) /* empty dynamic strides */,
DenseI64ArrayAttr({}) /* empty const shape*/,
@@ -260,10 +257,6 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
}
LogicalResult CreateNdDescOp::verify() {
- // getMixedSizes()/getMixedStrides() cannot represent a *bare dynamic* memref
- // (a dynamic dim has no SSA operand and no const attr), so derive the rank
- // from the memref type for a memref source; only an integer source needs the
- // shape/stride operands (checked below).
auto srcMemrefTy = dyn_cast<MemRefType>(getSourceType());
size_t rank = srcMemrefTy ? srcMemrefTy.getRank() : getMixedSizes().size();
bool invalidRank = false;
@@ -297,10 +290,6 @@ LogicalResult CreateNdDescOp::verify() {
"integer source.");
invalidRank = getMixedSizes().size() != getMixedStrides().size();
} else if (srcMemrefTy && hasExplicitShapeStrides) {
- // Deprecated: a memref source carries its own shape/strides (recovered at
- // lowering time, including for dynamic dims), so specifying them
- // explicitly is redundant and no longer supported. Use the bare form:
- // xegpu.create_nd_tdesc %memref : memref<...> -> !xegpu.tensor_desc<...>
return emitOpError("shape and strides should not be specified for a memref "
"source; they are inferred from the memref.");
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
index ffca61128e283..a3aaf01f356d5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp
@@ -137,11 +137,6 @@ class OptimizeCreateNdDescOp : public OpRewritePattern<xegpu::CreateNdDescOp> {
tdescType.getBoundaryCheck(), tdescType.getMemorySpace(),
tdescType.getLayout());
- // The memory region is unchanged. A memref source carries its own
- // (possibly dynamic) shape/strides, so use the bare-memref builder rather
- // than getMixedSizes/getMixedStrides (which cannot represent a dynamic
- // memref dim, and for which explicit shape/strides on a memref are no
- // longer allowed); an integer source still needs its explicit operands.
Value newOp;
if (isa<MemRefType>(source.getType()))
newOp =
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 7e609be3de4a3..c64defc471de8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -44,10 +44,7 @@ using namespace mlir;
namespace {
-/// Get the 2D lane data from a tensor desc type if it exists. The transpose
-/// optimization acts on the innermost 2 dims; a >2D descriptor is accepted only
-/// when its leading (batch) dims are unit, in which case the inner 2 dims are
-/// returned.
+/// Get the 2D lane data from a tensor desc type if it exists.
static std::optional<SmallVector<int64_t>>
getMaybeLaneData(xegpu::TensorDescType tdescType) {
auto layout = tdescType.getLayoutAttr();
@@ -56,9 +53,7 @@ getMaybeLaneData(xegpu::TensorDescType tdescType) {
return xegpu::getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt());
}
-/// Get the 2D lane layout from a tensor desc type if it exists. As with
-/// getMaybeLaneData, a >2D descriptor with unit leading dims yields its
-/// inner 2.
+/// Get the 2D lane layout from a tensor desc type if it exists.
static std::optional<SmallVector<int64_t>>
getMaybeLaneLayout(xegpu::TensorDescType tdescType) {
auto layout = tdescType.getLayoutAttr();
@@ -219,8 +214,9 @@ static Value generateLoads(ConversionPatternRewriter &rewriter,
auto shapeRatio = computeShapeRatio(data.getType().getShape(),
supportedShape)
.value(); // `ratio` must be defined if we reach here.
- // The transpose reshapes only the innermost 2 dims; any leading (batch) dims
- // are unit, so their offsets pass through unchanged and they are not tiled.
+ // 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) {
@@ -294,8 +290,6 @@ class XeGPUCreateNdDescOpPattern final
Value source = createNdOp.getSource();
auto memrefType = dyn_cast<MemRefType>(source.getType());
- // getMixedSizes/getMixedStrides cannot represent a dynamic memref dim, so
- // recover those from runtime metadata.
bool dynamicMemref =
memrefType && !xegpu::hasStaticShapeAndStrides(memrefType);
SmallVector<OpFoldResult> mixedSizes;
@@ -324,11 +318,7 @@ class XeGPUCreateNdDescOpPattern final
modifiedShape.back() = divideByConstant(
rewriter, loc, convertToValue(rewriter, loc, modifiedShape.back()),
innerLaneData);
- // Repacking narrow elements into wider ones (f16 -> i32) reinterprets the
- // innermost dim, so every stride except the innermost (which stays 1) must
- // be divided by innerLaneData to count the wider element -- the pitch and,
- // for a >2D descriptor, the leading (batch) strides (the XeVM fold reads
- // the latter back in repacked-element units).
+ // 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(mixedStrides);
@@ -337,9 +327,6 @@ class XeGPUCreateNdDescOpPattern final
rewriter, loc, convertToValue(rewriter, loc, modifiedStrides[i]),
innerLaneData);
- // The repacked shape/strides differ from the memref's own layout, so a
- // memref source must be lowered to an i64 base pointer -- specifying
- // explicit shape/strides on a memref create_nd is not allowed.
if (memrefType) {
Value baseIdx;
if (dynamicMemref) {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index f14df6de90c20..5914e23f1f11a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -200,16 +200,10 @@ struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
if (!targetShape)
return failure();
- // Keep the original (possibly high-D) source and only shrink the tile
- // shape; batch dims unroll to unit tiles and their offsets ride on the
- // load/store/prefetch, folded into the base pointer at XeVM lowering. We
- // never slice the source with a memref.subview -- for a dynamic-shape
- // source it could not guarantee a valid base address.
+ // 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);
- // A memref source uses the bare-memref builder: a dynamic dim has no SSA
- // operand for getMixedSizes/getMixedStrides.
if (isa<MemRefType>(src.getType()))
return xegpu::CreateNdDescOp::create(rewriter, loc, ndTy,
cast<TypedValue<MemRefType>>(src));
@@ -223,8 +217,7 @@ struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
// 2D: one tdesc, broadcast across tiles by pack/unpack.
newOps.push_back(makeCreateNd(newTdescTys[0]));
} else {
- // >2D: unpack expands the leading dims, so the source count must match
- // the pack count. Emit one (identical) tdesc per tile; CSE folds them.
+ // >2D: one tdesc per tile, so the source count matches the pack count.
for (Type t : newTdescTys)
newOps.push_back(makeCreateNd(t));
}
@@ -249,8 +242,7 @@ struct UnrollPrefetchNdOp : public UnrollPattern<xegpu::PrefetchNdOp> {
if (layout)
layout = layout.dropInstData();
- // Batch (leading) dims unroll to unit tiles like any other dim; a single
- // tdesc serves all tiles (offsets are folded into the base at XeVM).
+ // Batch (leading) dims unroll to unit tiles; one tdesc serves all.
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdesc = pack(
@@ -292,8 +284,7 @@ struct UnrollLoadNdOp : public UnrollPattern<xegpu::LoadNdOp> {
SmallVector<Value> newOps;
- // Batch (leading) dims unroll to unit tiles like any other dim; a single
- // tdesc serves all tiles (offsets are folded into the base at XeVM).
+ // Batch (leading) dims unroll to unit tiles; one tdesc serves all.
SmallVector<Type> convertedTdescTypes =
getUnrolledTypes(tdescTy, *targetShape);
SmallVector<Value> convertedTdescs = pack(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index dbf5c6e3ee1ac..0e8a386fb08b6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -206,9 +206,6 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
Value src = op.getSource();
SmallVector<Value> newCreateNdOps(count);
std::generate(newCreateNdOps.begin(), newCreateNdOps.end(), [&]() -> Value {
- // A memref source carries its own (possibly dynamic) shape/strides; use
- // the bare-memref builder rather than getMixedSizes/getMixedStrides,
- // which cannot represent a dynamic memref dim without an SSA operand.
if (isa<MemRefType>(src.getType()))
return xegpu::CreateNdDescOp::create(rewriter, loc, newTdescTy,
cast<TypedValue<MemRefType>>(src));
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index 875086f78750e..1464aa1ace477 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -276,11 +276,7 @@ gpu.func @load_dynamic_source3(%source: memref<?x?x?x?x?xf32>,
}
// -----
-// A dynamic high-D read whose vector rank equals the memref rank keeps the
-// *whole* memref as the create_nd source on the nd path: no rank-collapsing
-// memref.subview is emitted and the leading (batch) offsets (here the non-unit
-// 2x4 dims) stay on the load_nd. They are folded into the base pointer later,
-// at XeVM lowering, once blocking has unrolled them to unit tiles.
+// 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> {
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index 71c586f0ca54d..c6432378b496a 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -92,10 +92,7 @@ gpu.func @store_dynamic_source(%vec: vector<8x16xf32>,
}
// -----
-// A dynamic high-D write whose vector rank equals the memref rank keeps the
-// *whole* memref as the create_nd source on the nd path: no rank-collapsing
-// memref.subview is emitted and the leading (batch) offsets stay on the
-// store_nd (they are folded into the base pointer later, at XeVM lowering).
+// 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) {
diff --git a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
index 0f1da6047ef59..939adc951ac38 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -44,8 +44,7 @@ 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>
- // A dynamic memref uses the bare form: its shape/strides are recovered
- // from the memref via extract_strided_metadata (no explicit operands).
+ // 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 %[[SIZES]]#1 : index to i32
@@ -61,27 +60,22 @@ gpu.module @create_nd_tdesc {
gpu.return
}
- // A >2D (batched) create_nd descriptor keeps the innermost matrix as the
- // 2D-block surface (base_shape_w/h and pitch come from the innermost two
- // memref dims) and encodes the leading (batch) dim element strides into the
- // spare payload slots (5..). For a dynamic source the sizes/strides are
- // recovered via extract_strided_metadata: base_shape_w = sizes#2,
- // base_shape_h = sizes#1, base_pitch = strides#1, and leading stride slot 5 =
- // strides#0. The matching load/store lowering reads those strides to fold the
- // batch offsets into the base pointer, so the batch position stays out of the
- // surface.
+ // 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 %[[H]], %[[P2]] [3] : 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: vector.insert %[[LS0]], %[[P4]] [5] : i32 into vector<8xi32>
+ // 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 f4eb6f619530f..dfaa9521786a9 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/failed_conversion.mlir
@@ -57,8 +57,7 @@ gpu.module @test_kernel {
// -----
-// A 2D-block op transfers the innermost 2 dims only, so a non-unit leading dim
-// must be rejected rather than lowered to a load of the wrong shape.
+// 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 {
@@ -81,3 +80,27 @@ gpu.module @test_kernel {
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 1fde43119d8cf..408950c0d8f0f 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
@@ -92,49 +92,43 @@ gpu.module @load_store_check {
gpu.return
}
- // A >2D (batched) load_nd / store_nd folds the leading (batch) offsets into
- // the base pointer -- base += (sum_d offset[d] * leadingStride[d]) *
- // elemByteSize -- while the 2D-block surface stays the innermost matrix. For
- // a dynamic-shape source the sizes/strides are recovered via
- // extract_strided_metadata: the leading (batch) element stride is carried in
- // payload slot 5, folding batch index %z contributes %z * stride#0 * 4 bytes
- // to the base, and the surface is the innermost sizes#2 x sizes#1 matrix.
+ // 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
- // CHECK-DAG: %[[C4I64:.+]] = arith.constant 4 : i64
- // Load: recover shape/strides, encode the leading stride into slot 5, and
- // fold the batch index into the base pointer.
+ // 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: vector.insert %[[LSTRIDE]], %{{.+}} [5] : i32 into vector<8xi32>
+ // CHECK: %[[LROWS:.+]] = arith.divui %[[LSTRIDE]], %[[LPITCH]] : i32
// CHECK: %[[LBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
- // CHECK: %[[LZ:.+]] = arith.index_cast %[[Z]] : index to i64
- // CHECK: %[[LSE:.+]] = arith.extui %[[LSTRIDE]] : i32 to i64
- // CHECK: %[[LMUL:.+]] = arith.muli %[[LZ]], %[[LSE]] : i64
- // CHECK: %[[LOFF:.+]] = arith.muli %[[LMUL]], %[[C4I64]] : i64
- // CHECK: %[[LADDR:.+]] = arith.addi %[[LBASE]], %[[LOFF]] : i64
- // CHECK: %[[LPTR:.+]] = llvm.inttoptr %[[LADDR]] : i64 to !llvm.ptr<1>
- // CHECK: xevm.blockload2d %[[LPTR]]{{.*}}tile_height = 8 : i32, tile_width = 16 : i32
+ // 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 batch fold on the base.
+ // 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: vector.insert %[[SSTRIDE]], %{{.+}} [5] : i32 into vector<8xi32>
+ // CHECK: %[[SROWS:.+]] = arith.divui %[[SSTRIDE]], %[[SPITCH]] : i32
// CHECK: %[[SBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
- // CHECK: %[[SZ:.+]] = arith.index_cast %[[Z]] : index to i64
- // CHECK: %[[SSE:.+]] = arith.extui %[[SSTRIDE]] : i32 to i64
- // CHECK: %[[SMUL:.+]] = arith.muli %[[SZ]], %[[SSE]] : i64
- // CHECK: %[[SOFF:.+]] = arith.muli %[[SMUL]], %[[C4I64]] : i64
- // CHECK: %[[SADDR:.+]] = arith.addi %[[SBASE]], %[[SOFF]] : i64
- // CHECK: %[[SPTR:.+]] = llvm.inttoptr %[[SADDR]] : i64 to !llvm.ptr<1>
- // CHECK: xevm.blockstore2d %[[SPTR]]{{.*}}tile_height = 8 : i32, tile_width = 16 : i32
+ // 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/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index 049ebc715ce24..932cf6ff9b201 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -9,8 +9,7 @@ func.func @create_nd_tdesc_1(%src: memref<24xf32>) {
// -----
-// A memref source is authoritative for its own shape/strides, so specifying
-// them explicitly is deprecated and rejected. Use the bare form instead.
+// 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
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index ca215402b4f25..5e795017e2122 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -81,8 +81,7 @@ 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) {
- // A dynamic-shape memref uses the bare form; its shape/strides are inferred
- // from the memref (explicit shape/strides for a memref source is deprecated).
+ // 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>
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index 268c6b150bdb5..00c810edd605a 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -506,15 +506,7 @@ gpu.module @xevm_test {
}
// -----
-// The transpose optimization also fires for a >2D descriptor whose leading
-// (batch) dims are unit: it operates on the innermost two dims, repacking
-// f16->i32 for the HW transpose load. The leading unit dims are preserved on
-// the descriptor/result, lane_data is rank-matched to all ones, and every
-// non-innermost stride is divided by the pack factor (2) so the later XeVM
-// batch fold uses repacked-element units. For a dynamic-shape source the
-// sizes/strides are recovered via extract_strided_metadata; the dynamic
-// leading (batch) stride is divided by the pack factor with arith.shrui, while
-// the static inner strides fold to constants (4096 -> 2048, 64 -> 32).
+// 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
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 4c0beae89af4c..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,9 +60,7 @@ gpu.func @load_nd_transpose() {
gpu.return
}
-// A >2D load with unit leading dims is treated as its innermost 2D tile: the
-// packed attribute is still set when the inner lane_data marks a packed
-// (VNNI) B operand.
+// >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>
@@ -75,8 +73,7 @@ gpu.func @load_nd_packed_4d() {
gpu.return
}
-// A >2D load with unit leading dims is likewise transposed on its innermost 2D
-// tile when the inner lane layout requires it.
+// >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>
diff --git a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
index c52140cc10ffb..f23ae46905652 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -361,12 +361,7 @@ gpu.module @test {
}
//-----
- // Unrolling a >2D (batched) nd descriptor keeps the whole memref as the
- // create_nd source -- no memref.subview -- and yields a single unit-leading
- // tdesc that is reused across the batch tiles. The leading (batch) offset
- // stays on each unrolled load/store and is incremented per tile (it is *not*
- // zeroed and baked into a subview base). inst_data [1, 8, 16] unrolls the
- // size-4 batch dim into 4 tiles.
+ // 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
>From 8d5072dd9f255b9b47ac6fd7ef3fc5f2f46b3643 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 22 Aug 2026 19:26:52 +0000
Subject: [PATCH 8/8] [mlir][xegpu] Review cleanups: inline createNdDescriptor,
simplify verify
Address review feedback on the flattened-plane change:
- Inline the createNdDescriptor helper in VectorToXeGPU. It wrapped a single
CreateNdDescOp::create call, so each of its four call sites now builds the op
directly. The dropped assert(isStrided()) was already guaranteed on the
transfer paths by transferPreconditions, and a non-strided memref reaching
the vector.load/store paths is still rejected by CreateNdDescToXeVMPattern.
- Fold the deferred invalidRank flag in CreateNdDescOp::verify into the
integer-source branch, its only producer, and drop the now-dead
"(if source is a memref)" clause from the message: a memref source cannot
reach that check since explicit shape/strides on a memref are rejected
outright.
- Trim a comment in CreateNdDescToXeVMPattern.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 20 ++++---------------
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 4 +---
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 9 ++-------
3 files changed, 7 insertions(+), 26 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 1f6efdc091a75..9863206f14fe1 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -126,18 +126,6 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
return success();
}
-static xegpu::CreateNdDescOp createNdDescriptor(PatternRewriter &rewriter,
- Location loc,
- xegpu::TensorDescType descType,
- TypedValue<MemRefType> src) {
- [[maybe_unused]] MemRefType srcTy = src.getType();
- assert(srcTy.isStrided() && "Expected strided memref type");
-
- // Keep the memref as the source instead of collapsing it to an i64 base: it
- // carries the (possibly dynamic) offset/shape/strides the lowering recovers.
- return xegpu::CreateNdDescOp::create(rewriter, loc, descType, src);
-}
-
// Adjusts the strides of a memref according to a given permutation map for
// vector operations.
//
@@ -608,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 =
@@ -710,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(
@@ -828,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,
@@ -873,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 f2121d5178cf9..f782be5502a8a 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -272,9 +272,7 @@ class CreateNdDescToXeVMPattern
std::to_string(maxNdTdescLeadingDims) +
" leading dims (rank <= " + std::to_string(maxNdTdescRank) +
").");
- // The flattened-plane view needs every leading stride to be a whole number
- // of rows. Static strides are checked here; dynamic ones are assumed to
- // divide evenly (limitation 1 above).
+ // Limitation 1 above; dynamic strides are assumed to divide evenly.
if (rank > 2) {
SmallVector<std::optional<int64_t>> constStrides(rank, std::nullopt);
if (memrefTy) {
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index a906b8cbe8802..686e4342215b0 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -259,7 +259,6 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
LogicalResult CreateNdDescOp::verify() {
auto srcMemrefTy = dyn_cast<MemRefType>(getSourceType());
size_t rank = srcMemrefTy ? srcMemrefTy.getRank() : getMixedSizes().size();
- bool invalidRank = false;
bool invalidElemTy = false;
// Memory space of created TensorDesc should match with the source.
@@ -288,17 +287,13 @@ LogicalResult CreateNdDescOp::verify() {
if (getMixedStrides().empty() || getMixedSizes().empty())
return emitOpError("expecting strides and shape to be present for "
"integer source.");
- invalidRank = getMixedSizes().size() != getMixedStrides().size();
+ 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 "
More information about the Mlir-commits
mailing list