2 #include <ATen/xpu/XPUContext.h>
3 #include <ATen/xpu/XPUGeneratorImpl.h>
5 #include <c10/xpu/XPUCachingAllocator.h>
6 #include <c10/xpu/XPUFunctions.h>
20 static bool in_bad_fork = false; // True for children forked after xpu init
23 // Called in the forked child if xpu has already been initialized
30 // Should be called before the first xpu call. It is mainly called in lazy_init.
31 // Note: This is distinct from initExtension because a stub xpu implementation
40 // XPU management methods
53   c10::xpu::set_device(device_index);  in THXPModule_setDevice_wrap()
69   auto current_device = c10::xpu::exchange_device(device_index);  in THXPModule_exchangeDevice_wrap()
86   auto current_device = c10::xpu::maybe_exchange_device(device_index);  in THXPModule_maybeExchangeDevice_wrap()
95   auto device_index = c10::xpu::current_device();  in THXPModule_getDevice_wrap()
104   return THPUtils_packUInt64(at::xpu::device_count());  in THXPModule_getDeviceCount_wrap()
115   auto stream = at::xpu::getCurrentXPUStream(c10_device_index);  in THXPModule_getCurrentStream_wrap()
138       &at::xpu::getCurrentXPUStream(c10_device_index).queue());  in THXPModule_getCurrentStream_raw()
165   auto stream = at::xpu::XPUStream::unpack3(  in THXPModule_setStream_wrap()
170   auto device = c10::xpu::current_device();  in THXPModule_setStream_wrap()
172     c10::xpu::set_device(stream.device_index());  in THXPModule_setStream_wrap()
174   at::xpu::setCurrentXPUStream(stream);  in THXPModule_setStream_wrap()
187     c10::xpu::syncStreamsOnDevice(device_index);  in THXPModule_xpuSynchronize()
195   c10::xpu::XPUCachingAllocator::emptyCache();  in THXPModule_emptyCache()
231       c10::xpu::XPUCachingAllocator::getDeviceStats(device_index);  in THXPModule_memoryStats()
248   c10::xpu::XPUCachingAllocator::resetPeakStats(device_index);  in THXPModule_resetPeakMemoryStats()
261   c10::xpu::XPUCachingAllocator::resetAccumulatedStats(device_index);  in THXPModule_resetAccumulatedMemoryStats()
266 // XPU module initialization
270   using namespace c10::xpu;  in registerXpuDeviceProperties()
350   // Add method to torch.xpu  in bindGetDeviceProperties()
354       [](c10::DeviceIndex device) -> c10::xpu::DeviceProp* {  in bindGetDeviceProperties()
355         return at::xpu::getDeviceProperties(device);  in bindGetDeviceProperties()
368   auto m = THPObjectPtr(PyImport_ImportModule("torch.xpu"));  in THXPModule_initExtension()
378   auto num_gpus = c10::xpu::device_count();  in THXPModule_initExtension()
382     const auto& gen = at::xpu::detail::getDefaultXPUGenerator(i);  in THXPModule_initExtension()
437 namespace torch::xpu {  namespace
443 } // namespace torch::xpu