File size: 496 Bytes
9dd3461 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | #pragma once
#include <cstddef>
#include <c10/cuda/CUDACachingAllocator.h>
namespace at {
namespace 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);
}
};
}
}
|