目錄

1. core/framework
   1. resource
   2. allocator

核心概念

給出的只是內存分配器的接口，沒有給出具體實現。

Allocator

Allocator是一個內存分配器的接口類，它規定了一個內存分配器需要具有哪些API。具體看代碼：

class Allocator {
public:
    virtual void* AllocateRaw(size_t alignment, size_t num_bytes) = 0;
    virtual void DeallocateRaw(void* ptr) = 0;
    T* Allocate(size_t num_elements);
    T* Allocate(size_t num_elements, const AllocationAttributes& allocation_attr);
    void Deallocate(T* ptr, size_t num_elements);
    virtual bool TracksAllocationSizes();
    virtual bool ShouldAllocateEmptyTensors();
    virtual size_t RequestedSize(void* ptr);
    virtual size_t AllocatedSize(void* ptr);
    virtual int64 AllocationId(void* ptr);//本次內存分配的編號
    virtual size_t AllocatedSizeSlow(void* ptr);
    virtual void GetStats(AllocatorStats* stats);
}
 

另外，Allocator除了提供申請內存的接口之外，還提供了為申請好的內存調用默認構造和析構函數的接口。如果在申請的時候指定了對象的類型，就可以選擇調用對象所屬類的構造和析構方法。Allocator提供了針對三種常用類的構造方法，分別是String，ResourceHandle,Variant。

class Allocator {
public:
    //...
private:
    void RunCtor(T* p, size_t n);
    virtual void RunStringCtor(string* p, size_t n);
    virtual void RunStringDtor(string* p, size_t n);
    virtual void RunResourceCtor(ResourceHandle* p, size_t n);
    virtual void RunResourceDtor(ResourceHandle* p, size_t n);
    virtual void RunVariantCtor(Variant* p, size_t n);
    virtual void RunVariantDtor(Variant* p, size_t n);
}
 

AllocatorAttributes

不同的設備分配內存的方法並不相同，那是不是各設備只需要實現自身的內存分配器就可以了呢？如果在計算中每個設備只需要用到自己的內存，當然是沒有問題的，但在TF中，有些情況下為了效率，GPU也需要用到CPU內存，比如，為了使用DMA給某些設備傳送數據，我們仍然需要申請CPU內存。因此，當我們向一個設備索要內存分配器時，需要給它提供一些信息，告訴設備我們想要申請哪種類型的內存，這些信息就存儲在AllocatorAttributes類中。

struct AllocatorAttributes {
    void set_on_host(bool v);
    bool on_host() const;
    void set_nic_compatible(bool v);
    bool nic_compatible() const;
    void set_gpu_compatible(bool v);
    bool gpu_compatible() const;
    void set_track_sizes(bool v);
    bool track_sizes() const;
    void Merge(AllocatorAttributes other);
    bool IsEqualOrLessRestrictiveThan(const AllocatorAttributes& other);
    uint32 value = 0;//這個數值的高8位被保留為設備相關的設置。各設備的實現可以根據需要自行解析，為這些設備實現的操作也需要正確的解析它
}
 

AllocationAttributes

AllocatorAttributes很容易與另外一個類混淆，那就是AllocationAttributes。後者是為內存分配器的某一次具體的內存分配準備信息的，而前者是為向設備索要合適的內存分配器提供給設備的，使用時機完全不一樣。

class AllocationAttributes {
    bool no_retry_on_failure = false; //如果首次內存分配失敗了，不再嘗試。
    bool allocation_will_be_logged = false;//本次內存分配是否會被記錄
}

AllocatorWrapper

有時候我們想對某個內存分配器進行封裝，以便在某個API上實現定制化。這時就需要用到AllocatorWrapper類，它本質上就是對Allocator類的直接封裝。

AllocatorStats

為了對某個內存分配器已分配的內存進行統計，TF還設計了一個結構，AllocatorStats。

struct AllocatorStats {
    int64 num_allocs;//內存分配次數
    int64 bytes_in_use;//分配的內存中，當前正在使用的大小
    int64 max_bytes_in_use;//使用中的內存大小的峰值
    int64 max_alloc_size;//最大的單次內存分配大小
    int64 bytes_limit;//當前內存分配器能分配的最大內存量，如果申請內存大小超過這個閾值，返回0
    //...
}

tensorflow源碼解析之framework-allocator