Add support for lowering info during serialize_module, and add padding/partial to...
authorJordan Fix <jfix@fb.com>
Thu, 9 Sep 2021 19:59:54 +0000 (12:59 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Thu, 9 Sep 2021 20:01:28 +0000 (13:01 -0700)
commit6cc8cc6e56431be75745fd6cff602ef587cf3e53
tree057abd40ea65cfba985b7be631b7afb91ff579e1
parentd43fb75a2196742c806a2e938d10cd2fba1cc168
Add support for lowering info during serialize_module, and add padding/partial to it (#5810)

Summary:
Pull Request resolved: https://github.com/pytorch/glow/pull/5810

Pull Request resolved: https://github.com/pytorch/pytorch/pull/64725

- Any info added to the dict in node.meta["lowering_info"] will be added to the node_rep during serialization.
- Use this to add annotations on placeholders that allow partial inputs and require padding.
- Check for these annotations and set them in the NNPICompiledFunction as expected

Test Plan: Validated working on inline_cvr in stack. Additionally existing fx_glow end to end tests should still pass.

Reviewed By: 842974287

Differential Revision: D30824192

fbshipit-source-id: def64ef097aa35c337abb494415f7d437c6c7fa9
torch/fx/experimental/graph_manipulation.py