[Mlir-commits] [mlir] [mlir][xegpu] Lower dynamic high-D nd load/store via base-pointer fold (PR #215711)

[Sang Ik Lee]

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

This assumes memrefRank == tileRank — input IR is expected to subview the memref down
to the tile's rank. Either the verifier should check that, or this pattern should fail
to match when it doesn't hold, since a violation is silent today:

  xegpu.create_nd_tdesc %src : memref<4x2x8x16xf32> -> !xegpu.tensor_desc<1x8x16xf32>
  xegpu.load_nd %t[%b, %c0, %c0] : !xegpu.tensor_desc<1x8x16xf32> -> vector<8xf32>

lowers to a wrong base address with no diagnostic.


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