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

[Frank Schlimbach]

fschlimb wrote:

@rolfmorel  Thanks for your review.

> 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"?

The only thing that might not pretty is the pipeline itself. The resulting code should be comparable to a direct memref-based formulation.

There is some unclarity about when the ShardToMPI should be applied. In theory, it could be called after bufferization and so it would not deal with any tensor input/output. The drawback would be that cases like this would not offer tensor-level optimizations. Maybe this is another motivation for separating the communication primitives from the shard dialect.

MPI has buffer-semantics, so there is no way around going to buffer-land before calling MPI. 

Right now, the shard-dialect does not support oneshot-bufferization, so it can only be applied before bufferization. Until that's solved, this seems to be the best tradeoff, since we only go to buffers to call MPI. Before and after we are in tensor-land. 

The worst case you are referring to can only happen, if the compiler pipeline bypasses tensor land. For shard this is not possible if it is used normally (e.g. more than plain communication operations).

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