[Mlir-commits] [mlir] [MLIR][XeGPU] Allow some nd ops to have argument shapes mismatch for … (PR #120566)

[Adam Siemieniuk]

adam-smnk wrote:

> The tensor descriptor is a uniform value which all SIMD lanes share. There is only one copy of tensor descriptor being created for all SIMD lanes, and the creation doesn't involve lane id.

Good to know, that adds more context to the proposed solution as I initially thought this distribution split could be more arbitrary.
It would be helpful to add these details to the dialect documentation as XeGPU operations start covering both cooperative and distributed abstractions.

After closer looks and going through the detailed explanations (thanks @Jianhui-Li), the main competing options 2 and 4 revolve around making the distribution implicit(-ish as it is still captured by `sg_map`) vs explicit new nop subview. The mismatch is fine as option 2 distributed `load_nd` has built-in subview that follows strict semanctics captured by descriptor's map.

> The instruction does not live in a vacuum though. Since we are composing this with vector distribution, the resulting vector of a load would be consumed by some vector IR that describes what a logical thread does. It will have to consume some portion of the data that "belongs" to that thread.

That's been my concern, however, so far I couldn't come up with a case that would break without explicit mapping. Primarily, as I imagine most complex transformations would occur still at SIMD/cooperative level before SIMT.
If consumers expect different mapping compared to the one described by the load, then distribution might fail or extra step that changes layout through shared memory might be necessary anyway.

> I suggest we first go with option 2. When it become clear that we really need a subview op, we can revisit it.

+1 on taking option 2 as the first step as it is simpler - no extra ops, changes that don't break any existing code.


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