xref: /aosp_15_r20/external/pytorch/aten/src/ATen/xpu/CachingHostAllocator.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #pragma once
2 
3 #include <ATen/core/CachingHostAllocator.h>
4 #include <ATen/xpu/XPUEvent.h>
5 #include <c10/core/Allocator.h>
6 #include <c10/xpu/XPUStream.h>
7 
8 namespace at::xpu {
9 
10 TORCH_XPU_API c10::Allocator* getCachingHostAllocator();
11 
12 TORCH_XPU_API bool CachingHostAllocator_recordEvent(
13     void* ptr,
14     void* ctx,
15     c10::xpu::XPUStream stream);
16 
17 TORCH_XPU_API void CachingHostAllocator_emptyCache();
18 
HostAlloc(size_t size)19 inline TORCH_XPU_API at::DataPtr HostAlloc(size_t size) {
20   return getCachingHostAllocator()->allocate(size);
21 }
22 
23 } // namespace at::xpu
24