[PyTorch Mobile] Enable partial loading of GPU models on linux CPU machines (#51236)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51236
The problem we currently have with tracing is that GPU models can't load on devvm CPU machines. Here's why:
1. The Metal GPU ops don't exist so the validation that checks for missing ops kicks in and prevents loading
2. Even if the check for missing ops is commented out, the actual model contents can't be succssfully loaded (see t83364623 for details)
Hence, to work around these problems and allow tracing to detect GPU models, and skip actual tracing for these (as discussed in the meeting 2 weeks ago and based on recommendations from raziel, iseeyuan, and xta0), we're adding code to detect these GPU models based on the set of operators that show up in the file `extra/mobile_info.json`.
The code then skips tracing, and picks up the root operators from the model itself.
The diff below this one will be removed before landing since we don't want to check in the model - I've kept it here in case anyone wants to patch this diff in and run the command on their devvm locally.
ghstack-source-id: 120638092
Test Plan:
See {P168657729} for a successful run of tracing on a GPU model (person segmentation tier-0, v1001) provided by xta0
Also ran `buck test //xplat/pytorch_models/build/...` successfully.
Reviewed By: ljk53
Differential Revision: D26109526
fbshipit-source-id: 6119b0b59af8aae8b1feca0b8bc29f47a57a1a67
