[Mlir-commits] [mlir] [MLIR][XeGPU] Add unrolling/blocking support for 3D+ batched operations (PR #201725)

[Artem Kroviakov]

================
@@ -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;
+      for (int64_t off : batchOffsets)
+        svOffsets.push_back(rewriter.getIndexAttr(off));
+
+      SmallVector<OpFoldResult> svSizes;
+      for (int64_t d : batchBlockSize)
+        svSizes.push_back(rewriter.getIndexAttr(d));
+
+      SmallVector<OpFoldResult> svStrides(rank, rewriter.getIndexAttr(1));
----------------
akroviakov wrote:

+1, the lowering makes use of strides for the basePitch, so this line may need some explanation.

https://github.com/llvm/llvm-project/pull/201725