1 #pragma once 2 3 #include <c10/core/CachingDeviceAllocator.h> 4 #include <c10/xpu/XPUStream.h> 5 6 namespace c10::xpu::XPUCachingAllocator { 7 8 C10_XPU_API Allocator* get(); 9 10 C10_XPU_API void init(DeviceIndex device_count); 11 12 C10_XPU_API void emptyCache(); 13 14 C10_XPU_API void resetPeakStats(DeviceIndex device); 15 16 C10_XPU_API void resetAccumulatedStats(DeviceIndex device); 17 18 C10_XPU_API c10::CachingDeviceAllocator::DeviceStats getDeviceStats( 19 DeviceIndex device); 20 21 C10_XPU_API void* raw_alloc(size_t size); 22 23 C10_XPU_API void raw_delete(void* ptr); 24 25 C10_XPU_API void recordStream(const DataPtr& dataPtr, XPUStream stream); 26 27 } // namespace c10::xpu::XPUCachingAllocator 28