- do not create MPI operations if no halo exchange is needed
- allow returning sharding information through `!mesh.sharding`
(gets converted into a tuple of tensors)
- lowering `mesh.shard_shape` including fixes to the operation itself
- global symbol `static_mpi_rank` replaced by an DLTI attribute
(now aligned with MPIToLLVM)
- smaller fixes and some minor cleanup
---------
Co-authored-by: Christian Ulmann <christianulmann@gmail.com>
A collection of fixes to the mesh dialect
- allow constants in sharding propagation/spmdization
- fixes to tensor replication (e.g. 0d tensors)
- improved canonicalization
- sharding propagation incorrectly generated too many ShardOps
New operation `mesh.GetShardOp` enables exchanging sharding information
(like on function boundaries)
Follow-up to #102598 : as discussed, move tensor sharding implementation
into separate tensor extension lib.
@sogartar @yaochengji, could you take a look at this PR?
