Tensorflow之class Allocator




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();
  }
};

                
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值