File size: 528 Bytes
c1af2fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
#pragma once
#include <cstddef>
#include <c10/cuda/CUDACachingAllocator.h>
namespace at::cuda {
/// Allocator for Thrust to re-route its internal device allocations
/// to the THC allocator
class ThrustAllocator {
public:
typedef char value_type;
char* allocate(std::ptrdiff_t size) {
return static_cast<char*>(c10::cuda::CUDACachingAllocator::raw_alloc(size));
}
void deallocate(char* p, size_t size) {
c10::cuda::CUDACachingAllocator::raw_delete(p);
}
};
} // namespace at::cuda
|