refactor: move dynamo/TorchXLA bridge to pytorch/xla repo (#92601)
This is a follow up from the previous PR: https://github.com/pytorch/pytorch/pull/88449 , to move the dynamo/TorchXLA bridge from pytorch repo to xla repo.
Overall the dynamo/TorchXLA integration has the following four layers of code
- pybind layer: This is the bottom layer containing various pybind APIs as the foundation. This part resident in xla repo
- bridge layer: build upon the pybind layer to implement the trace once functionality. This layer and it's corresponding unit test are in pytorch repro previously. This PR (and the corresponding xla pr https://github.com/pytorch/xla/pull/4476 ) moves them to the xla repo.
- dynamo backend registration: this a thin layer registers 4 dynamo backends (training/inference/trace_once/trace_everytime). It remains in pytorch repo.
- benchmark script: the torchbench.py script in dynamo is adapted so it can be used in dynamo/TorchXLA integration. This one remains in pytorch repo.
We think the new code organization is cleaner.
I'll wait for the xla PR in first before trying to merge this one.
Tests
1. run the unit tests moved to the xla repo
2. Test for inference: `GPU_NUM_DEVICES=1 python benchmarks/dynamo/torchbench.py --randomize-input --performance --trace-on-xla --backend=torchxla_trace_once --only resnet18`
3. Test for training: `GPU_NUM_DEVICES=1 python benchmarks/dynamo/torchbench.py --randomize-input --performance --trace-on-xla --training --backend=aot_torchxla_trace_once --only resnet18 --collect-outputs`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/92601
Approved by: https://github.com/wconstab
