Tensorflow之class Allocator

来源:互联网 发布:火龙果软件 杨秀峰 编辑:程序博客网 时间:2024/06/05 09:01
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();  }};
 
原创粉丝点击