xla
c20bcdd3 - Introduce XLAGraphExecutor (#4270)

Commit
3 years ago
Introduce XLAGraphExecutor (#4270) Summary: This pull request moves all graph executor related parts out from XLATensor into XLAGraphExecutor such that it matches the format in upstream and therefore makes it easier to inherit the upstream LazyTensor and LazyGraphExecutor later. A few changes to notice: 1. DeviceContextArena::RegisterTensor/UnregisterTensor are now proxied via XLAGraphExecutor. 2. torch::lazy::IsSpecialScalar is used to replace our own helper. 3. GetDeviceData are moved into XlaDataCacheArena and proxied via XLAGraphExecutor. 4. DeviceBarrier are moved into DeviceLockerArena and proxied via XLAGraphExecutor. 5. Few quirks are added to ease this patch but will be removed later: 5.1. XLATensor(std::shared_ptr data) are made public such that DeviceContextArena::GetLiveTensors can access it. 5.2. XLAGraphExecutor is made to be a friend class of XLATensor in order to access some of the later's private methods/members. Test Plan: CI.
Author
Parents
Loading