目錄
- 什么是allocator
- 內存分配器的管理
- 內存分配追蹤
- 其它結構
- 關系圖
- 涉及的文件
- 迭代記錄
1. 什么是allocator
Allocator是所有內存分配器的基類，它定義了內存分配器需要實現的接口。
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);
};
這些API可以被分為兩類，一是內存分配與返還，二是內存分配信息追蹤。前者是內存分配器的本職工作，后者提供了對內存分配器所分配的內存進行追蹤和管理的功能。若想實現后者的功能，需要提供一定的數據結構支持，下文中的“內存分配追蹤”會詳細講到。
另外還遇到了一個新的結構AllocatorStats，這個類是對內存分配器當前分配內存量的一個宏觀統計，它的定義如下：
struct AllocatorStats {
int64 num_allocs;//內存分配次數
int64 bytes_in_use;//分配的內存中，當前正在使用的大小
int64 max_bytes_in_use;//使用中的內存大小的峰值
int64 max_alloc_size;//最大的單次內存分配大小
int64 bytes_limit;//當前內存分配器能分配的最大內存量，如果申請內存大小超過這個閾值，返回0
//...
}
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);
};
除了抽象的內存分配器接口之外，TF還為最常用的CPU內存分配器，提供了一個默認實現：
class CPUAllocator : public Allocator {
public:
//...
void GetStats(AllocatorStats* stats) override {
mutex_lock l(mu_);
*stats = stats_;
}
size_t AllocatedSizeSlow(void *ptr) override {
return port::MallocExtension_GetAllocatedSize(ptr);
}
private:
mutex mu_;
AllocatorStats stats_ GUARDED_BY(mu_);
};
需要注意兩點，第一，它為內存分配宏觀統計加入了一個數據成員stats_，在需要時直接將其返回，第二，為追蹤內存分配大小提供了一個“慢”版本的實現，這個實現的意義在於，如果需要知道某個指針對應的分配內存的大小，而我們又沒有專門為其准備這樣的數據記錄時，可以直接調用操作系統層面的函數，來獲取指針對應的分配內存的大小。這個操作會相對耗時，但聊勝於無。
2. 內存分配器的管理
不同類型的設備，可能需要不同的內存分配器，即便對於相同類型的設備，考慮到效率問題，也可能會提供不同的內存分配器版本。因此，需要一個對內存分配器進行集中管理的地方，TF為我們提供了類AllocatorRegistry：
class AllocatorRegistry {
public:
//內存分配器注冊
void Register(const string& name, int priority, Allocator* allocator);
//返回最高優先級的內存分配器
Allocator* GetAllocator();
//返回一個全局的內存分配器的注冊器
static AllocatorRegistry* Global();
private:
//內存分配器存儲位置
std::vector<AllocatorRegistryEntry> allocators_;
//...
}
所有的內存分配器都被保存在allocators_這個向量里，但向量中保存的並不是內存分配器本身，而是對它的一個封裝，我們看下這個封裝的結構：
typedef struct {
string name;
int priority;
Allocator* allocator;
} AllocatorRegistryEntry;
除了內存分配器之外，這個entry里還存放了內存分配器的名稱和優先級。當向AllocatorRegistry請求一個內存分配器時，它返回的是具有最高優先級的分配器，如果多個分配器有相同的優先級，就返回其中的一個。
AllocatorRegistry實際上是一個單例對象，它的Global接口返回一個全局靜態的注冊器對象。為了方便進行注冊，TF還設計了一個統一的注冊入口類：
class AllocatorRegistration {
public:
AllocatorRegistration(const string& name, int priority, Allocator* allocator){
AllocatorRegistry::Global()->Register(name,priority,allocator);
}
};
另外，TF還設計了一個宏來簡化注冊過程，感興趣的讀者可以參考源代碼。
3. 內存分配追蹤
剛才提到，在內存分配器的公共API中，有一類專門用於追蹤內存分配，需要有一些專用的數據結構，來保存每一次內存分配的信息。這些工作被TrackingAllocator類實現：
class TrackingAllocator : public Allocator {
public:
std::tuple<size_t, size_t, size_t> GetSizeAndUnRef();
//...
private:
bool UnRef() EXCLUSIZE_LOCKS_REQUIRED(mu_);
Allocator* allocator_;
mutex mu_;
int ref_ GUARDED_BY(mu_);
//當前仍在使用的分配內存大小，如果allocator_不支持內存追蹤，則為0
size_t allocated_ GUARDED_BY(mu_);
//allocated_的峰值
size_t high_watermark_ GUARDED_BY(mu_);
//當前內存分配器總共分配的內存大小
size_t total_bytes_ GUARDED_BY(mu_);
const bool track_sizes_locally_;
struct Chunk {
size_t requested_size;
size_t allocated_size;
int64 allocation_id;
};
std::unordered_map<void*, Chunk> in_use_ GUARDED_BY(mu_);
};
除了提供allocated_，high_watermark_，total_bytes_三個數據成員記錄內存分配的統計信息之外，更重要的是加入了in_use_這個數據成員。它從一個指針映射到一個Chunk，而這個Chunk中保存了每次內存分配需要的內存大小、實際分配的內存大小、本次內存分配的唯一標識。用這樣一個結構保存了每次內存分配的詳細信息。
再細說一下ref_成員的作用。這個數據成員存在的意義在於，保存當前內存分配器分配內存的次數，當所有分配的內存全部回收之后，就刪除掉當前的內存分配器。但你會發現，這個ref_成員在TrackingAllocator對象初始化的時候，本身已經賦值為1了，那即便我們每次返回內存時都調用UnRef函數將它減一，最終不還是不能為0嗎？原因在於，我們希望對於GetSizeAndUnRef這個函數在對象的生命周期內只調用一次，因此這個函數調用后會將ref_減一。而不論是GetSizeAndUnRef函數還是調用UnRef，只要ref_值減到0，就刪除這個對象。這要求我們必須僅調用GetSizeAndUnRef函數一次，否則就會出現內存泄漏。
4. 其它結構
在TF中，計算是發生在節點上的，而節點被分配在具體的設備上。對於一個GPU設備，在它上面運行的節點，是不是就不需要CPU內存了呢？顯然不是，比如，為了使用DMA給某些設備傳送數據，運行在GPU上的節點仍然需要申請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位被保留為設備相關的設置。各設備的實現可以根據需要自行解析，為這些設備實現的操作也需要正確的解析它
}
AllocatorAttributes很容易與另外一個類混淆，那就是AllocationAttributes。后者記錄的是為內存分配器的某一次具體的內存分配的屬性信息，使用時機完全不一樣。
class AllocationAttributes {
bool no_retry_on_failure = false; //如果首次內存分配失敗了，不再嘗試。
bool allocation_will_be_logged = false;//本次內存分配是否會被記錄
}
另外，有時候我們想對某個內存分配器進行封裝，以便在某個API上實現定制化。TF為此准備了類AllocatorWrapper類，它本質上就是對Allocator類的直接封裝，感興趣的讀者可以去看下源碼。
5. 關系圖
6. 涉及的文件
- allocator
- allocator_registry
- tracking_allocator
7. 迭代記錄
- v1.0 2018-08-25 文檔創建
- v2.0 2018-09-08 文檔重構