tensorflow原始碼解析之framework-allocator

1. 核心概念

allocator給出的只是記憶體分配器的介面，沒有給出具體實現。allocator_registry用單例模式實現了一個全域性的記憶體分配器註冊類，負責管理和註冊所有的記憶體分配器。

2. allocator

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

2.2 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位被保留為裝置相關的設定。各裝置的實現可以根據需要自行解析，為這些裝置實現的操作也需要正確的解析它}

2.3 AllocationAttributes

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

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

2.4 AllocatorWrapper

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

2.5 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
    //...}

3. allocator_registry

3.1 AllocatorRegistry

先看一下，註冊器內部是怎樣儲存記憶體分配器的：

class AllocatorRegistry {    //...private:    std::vector allocators_;    //...}

其實就是把記憶體分配器儲存在一個向量裡，但並不是直接儲存記憶體分配器本身，而是對它的一個封裝，我們看下這個封裝的結構：

typedef struct {    string name;    int priority;
    Allocator* allocator;
} AllocatorRegistryEntry;

除了記憶體分配器之外，這個entry裡還存放了記憶體分配器的名稱和優先順序。當我們向AllocatorRegistry請求一個記憶體分配器時，它返回的是具有最高優先順序的分配器，如果多個分配器有相同的優先順序，就返回其中的一個。

原文出處：https://www.cnblogs.com/jicanghai/p/9535808.html