[DTensor][EZ] op schema comparison so that no redistribute is called (#106158)
When looking at traces of TP more carefully, I found that for cases when input reshard is not needed, we also call redistribute within sharding propogation. Upon carefully checking, looks like the way we compare different op_schema is not correct.
One example can be seen in the following trace:
<img width="1146" alt="image" src="https://github.com/pytorch/pytorch/assets/6937752/7322d26f-7029-41f9-8f8c-5f27a6bb98f9">
As you can see, no collectives are called, and this redistribute is not needed.
With this change:
<img width="1491" alt="image" src="https://github.com/pytorch/pytorch/assets/6937752/eb4a971f-44c1-4d83-8671-fce94cfa926c">
Pull Request resolved: https://github.com/pytorch/pytorch/pull/106158
Approved by: https://github.com/Skylion007, https://github.com/wanchaol
