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

[Petr Kurapov]

kurapov-peter wrote:

> I feel that option 2 is quite unintuitive due to this mix of abstractions.

Yup, I think so too.

> Ideally, all sg_maps should be consumed at the time of distribution

Yes, that could be a good indication for the pass itself (e.g., as a return condition) to do the distribution.

> Currently I'm mostly leaning toward the 4th idea. Subview of a descriptor could work or maybe the load operation itself could contain some optional offsets, e.g. mapping it to thread ID.

Me too, that looks cleanest to me, so I'm suggesting we go there. For an offset I'm not sure what would it represent in a non-distributed case. You could do a vector of offsets that describes all the offsets for each individual lane but I think the definition of ops and sg_map wanted to avoid such an explicit representation. Anyway, there are clearly multiple ways of doing that.

> Ideally, it would be in a format that allows to easily identify whether IR is pre- or post-distribution.

I think that's not necessary, the transformation in a sense just slices through the dataflow in a way. You theoretically could do that multiple times (here's when this starts resembling tiling, which you could argue it kinda is).

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