mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
66b75693ae
# Motivation Reuse the shared code. Pull Request resolved: https://github.com/pytorch/pytorch/pull/165508 Approved by: https://github.com/EikanWang
29 lines
698 B
C++
29 lines
698 B
C++
#pragma once
|
|
|
|
#include <c10/core/AllocatorConfig.h>
|
|
#include <c10/core/CachingDeviceAllocator.h>
|
|
#include <c10/xpu/XPUStream.h>
|
|
|
|
namespace c10::xpu::XPUCachingAllocator {
|
|
|
|
C10_XPU_API Allocator* get();
|
|
|
|
C10_XPU_API void init(DeviceIndex device_count);
|
|
|
|
C10_XPU_API void emptyCache();
|
|
|
|
C10_XPU_API void resetPeakStats(DeviceIndex device);
|
|
|
|
C10_XPU_API void resetAccumulatedStats(DeviceIndex device);
|
|
|
|
C10_XPU_API c10::CachingDeviceAllocator::DeviceStats getDeviceStats(
|
|
DeviceIndex device);
|
|
|
|
C10_XPU_API void* raw_alloc(size_t size);
|
|
|
|
C10_XPU_API void raw_delete(void* ptr);
|
|
|
|
C10_XPU_API void recordStream(const DataPtr& dataPtr, XPUStream stream);
|
|
|
|
} // namespace c10::xpu::XPUCachingAllocator
|