[Mlir-commits] [mlir] [mlir][shard, mpi] Allowing 2d-grids and simplifying lowering shard.all_gather (PR #180243)

[Rolf Morel]

https://github.com/rolfmorel approved this pull request.

The changes to the tests make sense to me!

The transform changes I skimmed as I am not so familiar with this lowering path/involved dialects.

The one thing that took my notice is what Copilot (also) points out: the going to buffers and then back to tensors and then back to buffers and then back to tensors. Further lowering of that is likely to not look so pretty. Might it be possible to just do the transpose and collapse on the memrefs "in the middle"?

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