[example] Simplifying ingress sharding and using sharding propagation

[fschlimb]

Using the sharding propagation pass in the mlp-mpi example. This allows simpler ingress IR, because not all operands and results require annotation.

Also letting the payload return a result that is sharded like the activations, since that's what a real MLP would need (the result is the input for next layer).

This requires llvm/llvm-project#183028

[Copilot]

Pull request overview

Updates the mlp-mpi example to reduce manual sharding annotations by relying on sharding propagation during lowering, and simplifies host-side gather logic used for reference checking.

Changes:

• Simplifies the MLIR payload by removing several explicit intermediate shard annotations and reusing a shared sharding for matmul ingress/egress.
• Adds a sharding-propagation pass to the transform schedule (with optional IR printing when verbose).
• Refactors repeated gather calls in Python into a _gather helper and uses it for reference computation and correctness checks.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
examples/mlp-mpi/mlp_weight_stationary.mlir	Removes several explicit shard annotations and adjusts result sharding to rely more on propagation.
examples/mlp-mpi/mlp-mpi.py	Introduces _gather, updates reference/correctness paths to use it, and adds sharding-propagation to the schedule.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.