class Allocator { public: #ifdef EIGEN_VECTORIZE_AVX512 static constexpr size_t kAllocatorAlignment = 64; #else static constexpr size_t kAllocatorAlignment = 32; #endif virtual ~Allocator(); virtual string Name() = 0; virtual void* AllocateRaw(size_t alignment, size_t num_bytes) = 0; virtual void* AllocateRaw(size_t alignment, size_t num_bytes, const AllocationAttributes& allocation_attr) { return AllocateRaw(alignment, num_bytes); } virtual void DeallocateRaw(void* ptr) = 0; template <typename T> T* Allocate(size_t num_elements) { return Allocate<T>(num_elements, AllocationAttributes()); } template <typename T> T* Allocate(size_t num_elements, const AllocationAttributes& allocation_attr) { if (num_elements > (std::numeric_limits<size_t>::max() / sizeof(T))) { return NULL; } void* p = AllocateRaw(kAllocatorAlignment, sizeof(T) * num_elements, allocation_attr); T* typed_p = reinterpret_cast<T*>(p); if (typed_p) RunCtor<T>(typed_p, num_elements); return typed_p; } template <typename T> void Deallocate(T* ptr, size_t num_elements) { if (ptr) { RunDtor<T>(ptr, num_elements); DeallocateRaw(ptr); } } virtual bool TracksAllocationSizes() { return false; } virtual bool ShouldAllocateEmptyTensors() { return false; } virtual size_t RequestedSize(void* ptr) { CHECK(false) << "allocator doesn't track sizes"; return size_t(0); } virtual size_t AllocatedSize(void* ptr) { return RequestedSize(ptr); } virtual int64 AllocationId(void* ptr) { return 0; } virtual size_t AllocatedSizeSlow(void* ptr) { if (TracksAllocationSizes()) { return AllocatedSize(ptr); } return 0; } virtual void GetStats(AllocatorStats* stats) { stats->Clear(); } private: template <typename T> void RunCtor(T* p, size_t n) { static_assert(is_simple_type<T>::value, "T is not a simple type."); } template <typename T> void RunDtor(T* p, size_t n) {} virtual void RunStringCtor(string* p, size_t n) { for (size_t i = 0; i < n; ++p, ++i) new (p) string(); } virtual void RunStringDtor(string* p, size_t n) { for (size_t i = 0; i < n; ++p, ++i) p->~string(); } virtual void RunResourceCtor(ResourceHandle* p, size_t n) { for (size_t i = 0; i < n; ++p, ++i) new (p) ResourceHandle(); } virtual void RunResourceDtor(ResourceHandle* p, size_t n) { for (size_t i = 0; i < n; ++p, ++i) p->~ResourceHandle(); } virtual void RunVariantCtor(Variant* p, size_t n) { for (size_t i = 0; i < n; ++p, ++i) new (p) Variant(); } virtual void RunVariantDtor(Variant* p, size_t n) { for (size_t i = 0; i < n; ++p, ++i) p->~Variant(); } };
Tensorflow之class Allocator
最新推荐文章于 2022-05-01 10:58:36 发布