[Mlir-commits] [mlir] [MLIR][XeGPU] Add unrolling/blocking support for 3D+ batched operations (PR #201725)
Charitha Saumya
llvmlistbot at llvm.org
Fri Jun 5 15:49:28 PDT 2026
================
@@ -195,16 +195,59 @@ 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]);
+
SmallVector<Value> newOps;
+ for (SmallVector<int64_t> batchOffsets :
+ StaticTileOffsetRange(shape, batchBlockSize)) {
+ SmallVector<OpFoldResult> svOffsets;
----------------
charithaintc wrote:
what is `svOffsets` mean here? consider spelling out what sv means
https://github.com/llvm/llvm-project/pull/201725
More information about the Mlir-commits
mailing list