[Mlir-commits] [mlir] [mlir][xegpu] Lower dynamic high-D nd load/store via base-pointer fold (PR #215711)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Aug 11 20:22:13 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
Reworks lowering of >2D (batched) `xegpu.create_nd_tdesc`/`load_nd`/`store_nd`/ `prefetch_nd` (batched GEMM, rank-4 flash-attention) to keep the full high-D memref as the descriptor source and fold the leading (batch) offsets into the base pointer, instead of slicing a per-batch `memref.subview` during blocking.
For example, a batched load keeps the whole (possibly dynamic) memref and carries the batch index on the offset:
```mlir
%t = xegpu.create_nd_tdesc %src : memref<?x?x?xf32> -> !xegpu.tensor_desc<1x8x16xf32>
%v = xegpu.load_nd %t[%z, %c0, %c0] : !xegpu.tensor_desc<1x8x16xf32> -> vector<8xf32>
```
lowers so the 2D-block surface stays the innermost 8x16 matrix and %z is folded into the base pointer (base += %z * stride0 * elemBytes):
```mlir
%_, %_, %sizes:3, %strides:3 = memref.extract_strided_metadata %src
%base = vector.extract %payload[0] // base pointer (i64)
%zi = arith.index_cast %z : index to i64
%s0 = arith.extui %strides#<!-- -->0 : i32 to i64 // leading (batch) stride
%boff = arith.muli %zi, %s0
%off = arith.muli %boff, %c4_i64 // * elemBytes
%addr = arith.addi %base, %off // folded base pointer
%ptr = llvm.inttoptr %addr : i64 to !llvm.ptr<1>
%v = xevm.blockload2d %ptr, %w, %h, %pitch, %c0, %c0
{tile_height = 8, tile_width = 16} // surface = innermost 8x16
```
This supersedes the subview approach (#<!-- -->201725), which could not produce a valid base for a **dynamic-shape** memref, and avoids collapsing all leading dims into one tall 2D surface that overflows the HW 2D-block surface limit for large attention.
---
Patch is 94.21 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/215711.diff
20 Files Affected:
- (modified) mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (+9-42)
- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+82-11)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp (+23-5)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUArrayLengthOptimization.cpp (+15-7)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp (+121-51)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp (+71-198)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+10-3)
- (modified) mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp (+17-4)
- (modified) mlir/test/Conversion/VectorToXeGPU/load-to-xegpu.mlir (+7-31)
- (modified) mlir/test/Conversion/VectorToXeGPU/store-to-xegpu.mlir (+7-31)
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir (+45-51)
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir (+35-20)
- (modified) mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir (+34-11)
- (modified) mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir (+48)
- (modified) mlir/test/Dialect/XeGPU/array-len-op-unit.mlir (+3-4)
- (modified) mlir/test/Dialect/XeGPU/invalid.mlir (+13)
- (modified) mlir/test/Dialect/XeGPU/ops.mlir (+6-7)
- (modified) mlir/test/Dialect/XeGPU/peephole-optimize.mlir (+37)
- (modified) mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir (+29)
- (modified) mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir (+34)
``````````diff
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..269360cafe0d6 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -51,10 +51,13 @@ 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); 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) {
@@ -232,28 +235,52 @@ 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);
+
+ // 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);
@@ -288,6 +315,21 @@ 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);
+ }
+ }
rewriter.replaceOp(op, payload);
return success();
}
@@ -387,9 +429,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 +439,36 @@ 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);
+ }
// 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..0124b700ac1dc 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...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/215711
More information about the Mlir-commits
mailing list