Add partition info message to NetDef (#33616)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/33616
Att. We start by assign `node_name` of DeviceOption in each of the op in the net. The for each unique node_name, we will have a PartitionInfo describing the partition, including logic devices that it can be assigned and we establish the link by partition names.
Test Plan:
unittests
Canaries:
AF: https://our.intern.facebook.com/intern/ads/canary/424817103900710410
AI: https://our.intern.facebook.com/intern/ads/canary/424737510862189908
Reviewed By: ipiszy, bangshengtang, jfix71
Differential Revision: D20015493
fbshipit-source-id: 0bb0f30cfc3892f7b8709d87b8bc1fbab2f2c46d