class RefCounted {
 public:
  RefCounted();
  void Ref() const;
  bool Unref() const;
  bool RefCountIsOne() const;

 protected:
  virtual ~RefCounted();

 private:
  mutable std::atomic_int_fast32_t ref_;

  RefCounted(const RefCounted&) = delete;
  void operator=(const RefCounted&) = delete;
};

inline void RefCounted::Ref() const {
  DCHECK_GE(ref_.load(), 1);
  ref_.fetch_add(1, std::memory_order_relaxed);
}

inline bool RefCounted::Unref() const {
  DCHECK_GT(ref_.load(), 0);
  if (RefCountIsOne() || ref_.fetch_sub(1) == 1) {
    DCHECK((ref_.store(0), true));
    delete this;
    return true;
  } else {
    return false;
  }
}

class TensorBuffer : public core::RefCounted {
 	public:
    explicit TensorBuffer(void* data_ptr) : data_(data_ptr) {}
    ~TensorBuffer() override {}

    void* data() const { return data_; }
    template <typename T>   
    T* base() const {
    	return reinterpret_cast<T*>(data());
    }
 	private:
  	vodi* const data_;
}

class BufferBase : public TensorBuffer {
 	public:
    explicit BufferBase(Allocator* alloc, void* data_ptr)
        : TensorBuffer(data_ptr), alloc_(alloc) {}
  protected:
  	Allocator* const alloc_;
}

template <typename T>
class Buffer : public BufferBase {
public:
    Buffer(Allocator* a, int64 n) : BufferBase(a, TypedAllocator::Allocate<T>(a, n, AllocationAttributes())), elem_(n) {
    }
    Buffer(Allocator* a, int64 n, const AllocationAttributes& allocation_attr)
            : BufferBase(a, TypedAllocator::Allocate<T>(a, n, allocation_attr)), elem_(n) {
    }

    size_t size() const override {
        return sizeof(T) * elem_;
    }

private:
    ~Buffer() {
        if (data()) {
            if (LogMemory::IsEnabled()) {
                RecordDeallocation();
            }
            TypedAllocator::Deallocate<T>(alloc_, static_cast<T*>(data()), elem_);
        }
    }

private:
    T* data_;
    int64 elem_;
};

class Tensor {
public:
    Tensor(DataType type, const TensorShape& shape)  // 1. 朴素构造函数
            : Tensor(get_default_cpu_allocator(), type, shape) {
    }
    Tensor(Allocator* a, DataType type,
           const TensorShape& shape)  // 2. 带分配器的构造函数
            : shape_(shape), buf_(nullptr) {
        set_dtype(type);
        CHECK_NOTNULL(a);
        if (shape_.num_elements() > 0 || a->AllocatesOpaqueHandle()) {
            CASES(type, buf_ = new Buffer<T>(a, shape.num_elements()));
        }
        if (buf_ != nullptr && buf_->data() != nullptr && LogMemory::IsEnabled()) {
            LogMemory::RecordTensorAllocation("Unknown", LogMemory::UNKNOWN_STEP_ID, *this);
        }
    }

private:
    TensorShape shape_;
    TensorBuffer* buf_;
}

tensorflow::Tensor sparse_bias_tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({sparse_bias.size(), 1}));

static Allocator* get_default_cpu_allocator() {
  static Allocator* default_cpu_allocator =
      cpu_allocator(port::kNUMANoAffinity);
  return default_cpu_allocator;
}

Allocator* cpu_allocator(int numa_node) {
  static ProcessStateInterface* ps =
      AllocatorFactoryRegistry::singleton()->process_state();
  if (ps) {
    return ps->GetCPUAllocator(numa_node);
  } else {
    return cpu_allocator_base();
  }
}

static T* TypedAllocator::Allocate(Allocator* raw_allocator, size_t num_elements,
                                   const AllocationAttributes& allocation_attr) {
    if (num_elements > (std::numeric_limits<size_t>::max() / sizeof(T))) {
        return nullptr;
    }
    void* p = raw_allocator->AllocateRaw(Allocator::kAllocatorAlignment, sizeof(T) * num_elements, allocation_attr);
    T* typed_p = reinterpret_cast<T*>(p);
    if (typed_p)
        RunCtor<T>(raw_allocator, typed_p, num_elements);
    return typed_p;
}

class AllocatorFactory {
 public:
  virtual bool 					NumaEnabled() { return false; }
  virtual Allocator* 		CreateAllocator() = 0;
  virtual SubAllocator* CreateSubAllocator(int numa_node) = 0;
};

static Allocator* cpu_alloc = AllocatorFactoryRegistry::singleton()->GetAllocator();

class AllocatorFactoryRegistry {
 public:
  void Register(const char* source_file, int source_line, const string& name,
                int priority, AllocatorFactory* factory);
  Allocator* GetAllocator();
  SubAllocator* GetSubAllocator(int numa_node);
  
  static AllocatorFactoryRegistry* singleton();
  ProcessStateInterface* process_state() const { return process_state_; }

 protected:
  friend class ProcessState;
  ProcessStateInterface* process_state_ = nullptr;

 private:
  mutex mu_;
  bool first_alloc_made_ = false;
  struct FactoryEntry {
    const char* source_file;
    int source_line;
    string name;
    int priority;
    std::unique_ptr<AllocatorFactory> factory;
    std::unique_ptr<Allocator> allocator;
    std::vector<std::unique_ptr<SubAllocator>> sub_allocators;
  };
  std::vector<FactoryEntry> factories_ GUARDED_BY(mu_);
  const FactoryEntry* FindEntry(const string& name, int priority) const
      EXCLUSIVE_LOCKS_REQUIRED(mu_);

  TF_DISALLOW_COPY_AND_ASSIGN(AllocatorFactoryRegistry);
};

class AllocatorFactoryRegistration {
 public:
  AllocatorFactoryRegistration(const char* file, int line, const string& name,
                               int priority, AllocatorFactory* factory) {
    AllocatorFactoryRegistry::singleton()->Register(file, line, name, priority, factory);
  }
};

#define REGISTER_MEM_ALLOCATOR(name, priority, factory)                     \
  REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(__COUNTER__, __FILE__, __LINE__, name, \
                                     priority, factory)

#define REGISTER_MEM_ALLOCATOR_UNIQ_HELPER(ctr, file, line, name, priority, \
                                           factory)                         \
  REGISTER_MEM_ALLOCATOR_UNIQ(ctr, file, line, name, priority, factory)

#define REGISTER_MEM_ALLOCATOR_UNIQ(ctr, file, line, name, priority, factory) \
  static AllocatorFactoryRegistration allocator_factory_reg_##ctr(            \
      file, line, name, priority, new factory)

class ProcessStateInterface {
 public:
  virtual ~ProcessStateInterface() {}
  virtual Allocator* GetCPUAllocator(int numa_node) = 0;
};

Allocator* ProcessState::GetCPUAllocator(int numa_node) {
  if (!numa_enabled_ || numa_node == port::kNUMANoAffinity) numa_node = 0;
  mutex_lock lock(mu_);
  while (cpu_allocators_.size() <= static_cast<size_t>(numa_node)) {
    const bool alloc_visitors_defined =
        (!cpu_alloc_visitors_.empty() || !cpu_free_visitors_.empty());
    bool use_bfc_allocator = false;
    Status status = ReadBoolFromEnvVar("TF_CPU_ALLOCATOR_USE_BFC", alloc_visitors_defined, &use_bfc_allocator);
    Allocator* allocator = nullptr;
    SubAllocator* sub_allocator =
        (numa_enabled_ || alloc_visitors_defined || use_bfc_allocator)
            ? new BasicCPUAllocator(
                  numa_enabled_ ? numa_node : port::kNUMANoAffinity,
                  cpu_alloc_visitors_, cpu_free_visitors_)
            : nullptr;
    if (use_bfc_allocator) {
      int64 cpu_mem_limit_in_mb = -1;
      Status status = ReadInt64FromEnvVar("TF_CPU_BFC_MEM_LIMIT_IN_MB", 1LL << 16 /*64GB max by default*/, &cpu_mem_limit_in_mb);
      int64 cpu_mem_limit = cpu_mem_limit_in_mb * (1LL << 20);
      DCHECK(sub_allocator);
      allocator = new BFCAllocator(sub_allocator, cpu_mem_limit, true , "bfc_cpu_allocator_for_gpu");
    } else if (sub_allocator) {
      DCHECK(sub_allocator);
      allocator = new PoolAllocator(100, true, sub_allocator, new NoopRounder, "cpu_pool");
    } else {
      DCHECK(!sub_allocator);
      allocator = cpu_allocator_base();
    }
    if (LogMemory::IsEnabled() && !allocator->TracksAllocationSizes()) {
      allocator = new TrackingAllocator(allocator, true);
    }
    cpu_allocators_.push_back(allocator);
    if (!sub_allocator) {
      DCHECK(cpu_alloc_visitors_.empty() && cpu_free_visitors_.empty());
    }
  }
  return cpu_allocators_[numa_node];
}

class CPUAllocatorFactory : public AllocatorFactory {
 public:
  Allocator* CreateAllocator() override { return new CPUAllocator; }
  SubAllocator* CreateSubAllocator(int numa_node) override {
    return new CPUSubAllocator(new CPUAllocator);
  }

 private:
  class CPUSubAllocator : public SubAllocator {
   public:
    explicit CPUSubAllocator(CPUAllocator* cpu_allocator)
        : SubAllocator({}, {}), cpu_allocator_(cpu_allocator) {}

    void* Alloc(size_t alignment, size_t num_bytes) override {
      return cpu_allocator_->AllocateRaw(alignment, num_bytes);
    }

    void Free(void* ptr, size_t num_bytes) override {
      cpu_allocator_->DeallocateRaw(ptr);
    }

   private:
    CPUAllocator* cpu_allocator_;
  };
};

void* CPUAllocator::AllocateRaw(size_t alignment, size_t num_bytes) override {
    if (num_bytes > LargeAllocationWarningBytes() && single_allocation_warning_count_ < kMaxSingleAllocationWarnings) {
      	++single_allocation_warning_count_;
    }
    void* p = port::AlignedMalloc(num_bytes, alignment);
    if (cpu_allocator_collect_stats) {
        const std::size_t alloc_size = port::MallocExtension_GetAllocatedSize(p);
        mutex_lock l(mu_);
        ++stats_.num_allocs;
        stats_.bytes_in_use += alloc_size;
        stats_.peak_bytes_in_use = std::max<int64>(stats_.peak_bytes_in_use, stats_.bytes_in_use);
        stats_.largest_alloc_size = std::max<int64>(stats_.largest_alloc_size, alloc_size);

        if (stats_.bytes_in_use > TotalAllocationWarningBytes() && 
            total_allocation_warning_count_ < kMaxTotalAllocationWarnings) {
          	++total_allocation_warning_count_;
      	}
    }
    return p;
}

void CPUAllocator::DeallocateRaw(void* ptr) override {
    if (cpu_allocator_collect_stats) {
        const std::size_t alloc_size = port::MallocExtension_GetAllocatedSize(ptr);
        mutex_lock l(mu_);
        stats_.bytes_in_use -= alloc_size;
    }
    port::AlignedFree(ptr);
}

void* AlignedMalloc(size_t size, int minimum_alignment) {
#if defined(__ANDROID__)
    return memalign(minimum_alignment, size);
#else  // !defined(__ANDROID__)
    void* ptr = nullptr;
    const int required_alignment = sizeof(void*);
    if (minimum_alignment < required_alignment) return Malloc(size);
    int err = posix_memalign(&ptr, minimum_alignment, size);
    if (err != 0) {
      return nullptr;
    } else {
      return ptr;
    }
#endif
}

#ifdef JEMALLOC_MANGLE
#  ifndef JEMALLOC_NO_DEMANGLE
#    define JEMALLOC_NO_DEMANGLE
#  endif
#  define malloc_conf je_malloc_conf
#  define malloc_message je_malloc_message
#  define malloc je_malloc
#  define calloc je_calloc
#  define posix_memalign je_posix_memalign
#  define aligned_alloc je_aligned_alloc
#  define realloc je_realloc
#  define free je_free
#  define mallocx je_mallocx
#  define rallocx je_rallocx
#  define xallocx je_xallocx
#  define sallocx je_sallocx
#  define dallocx je_dallocx
#  define sdallocx je_sdallocx
#  define nallocx je_nallocx
#  define mallctl je_mallctl
#  define mallctlnametomib je_mallctlnametomib
#  define mallctlbymib je_mallctlbymib
#  define malloc_stats_print je_malloc_stats_print
#  define malloc_usable_size je_malloc_usable_size
#  define valloc je_valloc
#endif

allocator = new BFCAllocator(sub_allocator, cpu_mem_limit, true /*allow_growth*/, "bfc_cpu_allocator_for_gpu" /*name*/);

struct Chunk {
    size_t size = 0;  // Full size of buffer.
    size_t requested_size = 0;
    int64 allocation_id = -1;
    void* ptr = nullptr;  // pointer to granted subbuffer.
    ChunkHandle prev = kInvalidChunkHandle;
    ChunkHandle next = kInvalidChunkHandle;
    // What bin are we in?
    BinNum bin_num = kInvalidBinNum;
    // Optional count when this chunk was most recently made free.
    uint64 freed_at_count = 0;
    bool in_use() const { return allocation_id != -1; }
};

struct Bin {
    size_t bin_size = 0;
    class ChunkComparator {
     public:
        explicit ChunkComparator(BFCAllocator* allocator)
            : allocator_(allocator) {}
        // Sort first by size and then use pointer address as a tie breaker.
        bool operator()(const ChunkHandle ha,
                        const ChunkHandle hb) const NO_THREAD_SAFETY_ANALYSIS {
          const Chunk* a = allocator_->ChunkFromHandle(ha);
          const Chunk* b = allocator_->ChunkFromHandle(hb);
          if (a->size != b->size) {
            return a->size < b->size;
          }
          return a->ptr < b->ptr;
        }

     private:
      	BFCAllocator* allocator_;  // The parent allocator
    };

    typedef std::set<ChunkHandle, ChunkComparator> FreeChunkSet;
    // List of free chunks within the bin, sorted by chunk size.
    FreeChunkSet free_chunks;
    Bin(BFCAllocator* allocator, size_t bs) : bin_size(bs), free_chunks(ChunkComparator(allocator)) {}
};

// Map from bin size to Bin
Bin* BinFromIndex(BinNum index) {
  return reinterpret_cast<Bin*>(&(bins_space_[index * sizeof(Bin)]));
}
size_t BinNumToSize(BinNum index) {
  return static_cast<size_t>(256) << index;
}
BinNum BinNumForSize(size_t bytes) {
  uint64 v = std::max<size_t>(bytes, 256) >> kMinAllocationBits;
  int b = std::min(kNumBins - 1, Log2FloorNonZero(v));
  return b;
}
Bin* BinForSize(size_t bytes) { return BinFromIndex(BinNumForSize(bytes)); }

char bins_space_[sizeof(Bin) * kNumBins];

for (BinNum b = 0; b < kNumBins; b++) {
    size_t bin_size = BinNumToSize(b);
    new (BinFromIndex(b)) Bin(this, bin_size);
}

void* BFCAllocator::FindChunkPtr(BinNum bin_num, size_t rounded_bytes, size_t num_bytes, uint64 freed_before) {
  for (; bin_num < kNumBins; bin_num++) {
    Bin* b = BinFromIndex(bin_num);
    for (auto citer = b->free_chunks.begin(); citer != b->free_chunks.end(); ++citer) {
      const BFCAllocator::ChunkHandle h = (*citer);
      BFCAllocator::Chunk* chunk = ChunkFromHandle(h);
      DCHECK(!chunk->in_use());
      if (freed_before > 0 && freed_before < chunk->freed_at_count) { continue; }
      if (chunk->size >= rounded_bytes) {
        RemoveFreeChunkIterFromBin(&b->free_chunks, citer);
        const int64 kMaxInternalFragmentation = 128 << 20;  // 128mb
        if (chunk->size >= rounded_bytes * 2 || static_cast<int64>(chunk->size) - rounded_bytes >= kMaxInternalFragmentation) {
          SplitChunk(h, rounded_bytes);
          chunk = ChunkFromHandle(h);  // Update chunk pointer in case it moved
        }
        chunk->requested_size = num_bytes;
        chunk->allocation_id = next_allocation_id_++;

        ++stats_.num_allocs;
        stats_.bytes_in_use += chunk->size;
        stats_.peak_bytes_in_use = std::max(stats_.peak_bytes_in_use, stats_.bytes_in_use);
        stats_.largest_alloc_size = std::max<std::size_t>(stats_.largest_alloc_size, chunk->size);

        return chunk->ptr;
      }
    }
  }
  return nullptr;
}

void BFCAllocator::SplitChunk(BFCAllocator::ChunkHandle h, size_t num_bytes) {
  // Allocate the new chunk before we do any ChunkFromHandle
  ChunkHandle h_new_chunk = AllocateChunk();

  Chunk* c = ChunkFromHandle(h);

  // Create a new chunk starting num_bytes after c
  BFCAllocator::Chunk* new_chunk = ChunkFromHandle(h_new_chunk);
  new_chunk->ptr = static_cast<void*>(static_cast<char*>(c->ptr) + num_bytes);
  region_manager_.set_handle(new_chunk->ptr, h_new_chunk);

  // Set the new sizes of the chunks.
  new_chunk->size = c->size - num_bytes;
  c->size = num_bytes;

  // The new chunk is not in use.
  new_chunk->allocation_id = -1;

  // It inherits the freed time.
  new_chunk->freed_at_count = c->freed_at_count;

  // Maintain the pointers.
  // c <-> c_neighbor becomes
  // c <-> new_chunk <-> c_neighbor
  BFCAllocator::ChunkHandle h_neighbor = c->next;
  new_chunk->prev = h;
  new_chunk->next = h_neighbor;
  c->next = h_new_chunk;
  if (h_neighbor != kInvalidChunkHandle) {
    Chunk* c_neighbor = ChunkFromHandle(h_neighbor);
    c_neighbor->prev = h_new_chunk;
  }

  // Add the newly free chunk to the free bin.
  InsertFreeChunkIntoBin(h_new_chunk);
}

void* BFCAllocator::AllocateRawInternal(size_t unused_alignment,
                                        size_t num_bytes,
                                        bool dump_log_on_failure,
                                        uint64 freed_before) {
  size_t rounded_bytes = RoundedBytes(num_bytes);
  BinNum bin_num = BinNumForSize(rounded_bytes);

  mutex_lock l(lock_);
  if (!timestamped_chunks_.empty()) {
    MergeTimestampedChunks(0);
  }
  void* ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, freed_before);
  if (ptr != nullptr) {
    return ptr;
  }
  if (Extend(unused_alignment, rounded_bytes)) {
    ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, freed_before);
    if (ptr != nullptr) {
      return ptr;
    }
  }

  if ((freed_before == 0) && (!timestamped_chunks_.empty())) {
    if (MergeTimestampedChunks(rounded_bytes)) {
      ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, freed_before);
      if (ptr != nullptr) {
        return ptr;
      }
    }
  }
  if (DeallocateFreeRegions(rounded_bytes) &&
      Extend(unused_alignment, rounded_bytes)) {
    ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, freed_before);
    if (ptr != nullptr) {
      return ptr;
    }
  }
  return nullptr;
}

class AllocationRegion {
 public:
    AllocationRegion(void* ptr, size_t memory_size)
        : ptr_(ptr),
          memory_size_(memory_size),
          end_ptr_(
                static_cast<void*>(static_cast<char*>(ptr_) + memory_size_)) {
        DCHECK_EQ(0, memory_size % kMinAllocationSize);
        const size_t n_handles =
            (memory_size + kMinAllocationSize - 1) / kMinAllocationSize;
        handles_.reset(new ChunkHandle[n_handles]);
        for (size_t i = 0; i < n_handles; i++) {
          	handles_[i] = kInvalidChunkHandle;
        }
    }

    void* ptr() const { return ptr_; }
    void* end_ptr() const { return end_ptr_; }
    size_t memory_size() const { return memory_size_; }
    ChunkHandle get_handle(const void* p) const { return handles_[IndexFor(p)];}
    void set_handle(const void* p, ChunkHandle h) { handles_[IndexFor(p)] = h; }
    void erase(const void* p) { set_handle(p, kInvalidChunkHandle); }

 private:
    void Swap(AllocationRegion* other) {
        std::swap(ptr_, other->ptr_);
        std::swap(memory_size_, other->memory_size_);
        std::swap(end_ptr_, other->end_ptr_);
        std::swap(handles_, other->handles_);
    }

    size_t IndexFor(const void* p) const {
      	std::uintptr_t p_int = reinterpret_cast<std::uintptr_t>(p);
        std::uintptr_t base_int = reinterpret_cast<std::uintptr_t>(ptr_);
        DCHECK_GE(p_int, base_int);
        DCHECK_LT(p_int, base_int + memory_size_);
        return static_cast<size_t>(((p_int - base_int) >> kMinAllocationBits));
    }


    void* ptr_ = nullptr;
    size_t memory_size_ = 0;
    void* end_ptr_ = nullptr;

    std::unique_ptr<ChunkHandle[]> handles_;

    TF_DISALLOW_COPY_AND_ASSIGN(AllocationRegion);
};

class RegionManager {
 public:
    RegionManager() {}
    ~RegionManager() {}

    void AddAllocationRegion(void* ptr, size_t memory_size) {
      auto entry = std::upper_bound(regions_.begin(), regions_.end(), ptr, &Comparator);
      regions_.insert(entry, AllocationRegion(ptr, memory_size));
    }

    std::vector<AllocationRegion>::iterator RemoveAllocationRegion(std::vector<AllocationRegion>::iterator it) {
      return regions_.erase(it);
    }

    ChunkHandle get_handle(const void* p) const {
      return RegionFor(p)->get_handle(p);
    }

    void set_handle(const void* p, ChunkHandle h) {
      return MutableRegionFor(p)->set_handle(p, h);
    }
    void erase(const void* p) { return MutableRegionFor(p)->erase(p); }

    const std::vector<AllocationRegion>& regions() const { return regions_; }

 private:
    static bool Comparator(const void* ptr, const AllocationRegion& other) {
      return ptr < other.end_ptr();
    }

    AllocationRegion* MutableRegionFor(const void* p) {
      return const_cast<AllocationRegion*>(RegionFor(p));
    }

    const AllocationRegion* RegionFor(const void* p) const {
      auto entry = std::upper_bound(regions_.begin(), regions_.end(), p, &Comparator);
      if (entry != regions_.end()) { return &(*entry);}
      return nullptr;
    }

 private:
  	std::vector<AllocationRegion> regions_;
};