[Mlir-commits] [mlir] [mlir][xegpu] Fix high-D nd base_height for non-packed plane strides (PR #225556)

[Sang Ik Lee]

================
@@ -214,27 +214,30 @@ translateStoreXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
 
 //
 // High-D (>2D) nd descriptors are lowered by viewing the source as a single
-// flattened 2D plane: `base_height` is the product of all dims but the
-// innermost, so the 2D-block surface covers every leading (batch) plane at
-// once, and a batch position becomes a row offset into it. Leaving `base_ptr`
-// at the true base means an out-of-range batch index lands past `base_height`,
-// where the HW boundary check handles it, instead of aiming the surface at
-// unmapped memory. Encoding the leading strides as row counts (`stride[d] /
-// stride[R-2]`) also makes them dimensionless, so they survive element-type
-// repacking (e.g. the f16 -> i32 transpose repack) without a unit conversion.
+// flattened 2D plane, so the 2D-block surface covers every leading (batch)
+// plane at once and a batch position becomes a row offset into it. The surface
+// height is the row extent the leading dims reach:
+//   base_height = size[R-2] + sum_d (size[d] - 1) * (stride[d] / stride[R-2])
+//
+// Leaving `base_ptr` at the true base means an out-of-range batch index lands
+// past `base_height`, where the HW boundary check handles it, instead of aiming
+// the surface at unmapped memory. Encoding the leading strides as row counts
+// (`stride[d] / stride[R-2]`) also makes them dimensionless, so they survive
+// element-type repacking (e.g. the f16 -> i32 transpose repack) without a unit
+// conversion.
 //
 // Limitations of the flattened-plane view:
 //  1. Each leading stride must be a whole number of rows. A source with gaps
 //     between planes (`stride[d] % stride[R-2] != 0`) is not lowered. This can
 //     only be checked when the strides are static; for dynamic strides the
 //     divisibility is assumed.
-//  2. `base_height` grows to the product of the leading dims, so a source with
-//     a large batch x head x sequence extent can exceed the HW 2D-block
-//     surface height.
+//  2. `base_height` grows with the leading extent, so a source with a large
+//     batch x head x sequence extent can exceed the HW 2D-block surface height.
 //  3. Plane boundaries are invisible to the boundary check: a tile whose rows
-//     run past `size[R-2]` reads the next plane's rows instead of the zero
-//     padding a per-plane surface would return. This only matters when
-//     `size[R-2]` is not a multiple of the tile height.
+//     run past `size[R-2]` reads the following rows of the surface (the next
----------------
silee2 wrote:

Limitation is written in terms of reads. What about writes?

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