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

[Jianhui Li]

================
@@ -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));
----------------
Jianhui-Li wrote:

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.

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