背景
[作者：DeepLearningStack，阿里巴巴算法工程師，開源TensorFlow Contributor]
歡迎大家關注我的公眾號，“互聯網西門二少”，我將繼續輸出我的技術干貨~
作為一款優秀的異構深度學習算法框架，TensorFlow可以在多種設備上運行算法程序，包括CPU，GPU，Google開發的TPU等。因為TensorFlow的架構特性非常好，可擴展性很強，所以也支持用戶自定義補充其他計算設備，比如可以接入FPGA甚至是自定義芯片等。雖然在Google發布的TensorFlow white paper中並沒有過多的描述設備管理相關的內容，只是從較高層面上闡述了Device以及Job的命名規則，但是其設備管理模塊確實是處於架構中比較核心的地位。本文將從架構層面出發，詳細闡述當前TensorFlow源碼中關於設備管理的設計思想和相關細節，理解這部分內容不但可以加深對TensorFlow源碼的理解，還可以有能力接入一些自定義的設備。本文是TensorFlow設備管理的第一篇文章，為了能讓讀者更好的切入到TensorFlow源碼閱讀過程中，先從較為簡單地Device的創建和注冊機制開始。讀者也可以一邊對照本文一邊對照源碼進行閱讀和梳理，並歡迎大家提出各種相關的意見和建議。
計算設備（Device）定義
Google在2015年發布的第一版TensorFlow white paper中，從功能角度上闡述了Device的相關內容，我們可以總結出關鍵的幾點如下：
1. 在TensorFlow中的Device有着特殊的命名規則，無論在單機還是分布式任務中，都能依靠命名確定唯一的Device，它是Device的唯一標識符；
2. TensorFlow使用注冊機制將實現多種Device的添加管理；
3. 每個Device自己管理Memory的分配和釋放。
Device的命名一般使用/job:{job_name}/task:{job_id}/device:{type}:{device_id}的格式，這是為了更好的支持分布式任務。例如/job:worker/task:17/device:gpu:3就表示該Device是ID為17的worker上的ID為GPU設備。至於分布式中的相關概念會在其他文章中詳細闡述，在這里我們只需要知道Device的命名可以幫助我們在任務中定位到具體的某個唯一Device即可。
Device的注冊機制
TensorFlow有兩處涉及到了設備管理，一處存在於TensorFlow的core中，另一部分存在於XLA中。本文只會闡述TensorFlow core中的內容，關於XLA部分的講解可以參見其他blog。TensorFlow使用工廠來創建各種各樣的Device，並且幾乎為每一種Device都實現了對應的DeviceFactory。初讀代碼時可能會被各種Device類名搞混，下面先從TensorFlow中已經有的Device類出發，給出各種Device的類說明。
Device相關類圖
TensorFlow對不同種類的Device做了多層級的抽象，下面的類圖是從當前TensorFlow源碼中梳理出的比較重要的部分。
上圖中的每個類（class）都可以在TensorFlow的源碼中找到，因為當前TensorFlow的進化過程比較快，代碼結構並不處於一個十分穩定的狀態，所以上述類圖中的類結構關系可能在未來發生一些變化，這一點從注釋中也可以看出一些端倪，但是大的架構不會發生變化，所以梳理類結構也是十分有意義的。下面將對每個類的作用進行簡單地闡述，讀者在理解這些類的作用以及關系后再去閱讀源碼就會非常清晰了。
1. DeviceAttributes：在TensorFlow源碼中並不能直接找到這個類的c++定義，其實它是由protobuf編譯出來的。其含義也很好理解，是對特定Device屬性的封裝，比如Device的type，存儲的限制等等；
2. DeviceBase：定義了Device用到的基本方法，比較重要的是GetAllocator和MakeTensorFromProto，前者返回存儲器的分配器，后者是從Proto中生成Tensor，該方法必須被重寫；
3. Device：這個類比DeviceBase更加具體，新包含了一些用於計算調用的方法，比如Compute函數就會調用某Op的Compute計算；
4. RemoteDevice：這個類會在分布式時使用，在此暫時不進行闡述；
5. SingleThreadedCpuDevice：這是一個僅有單個線程的CPU Device抽象，它和ThreadPoolDevice不同，只被用於in expensive的Op計算，這樣做的好處是避免了一些thread初始化工作；
6. ThreadPoolDevice：這就是CPU Device的實現；
7. RenamedDevice：Device的封裝類，封裝時會再取一個新的Device name；
8. GPUCompatibleCPUDevice：這也是CPU Device的實現，和ThreadPoolDevice不同的是，它更多的是為了和GPU進行交互而存在，從其使用的CudaHostAllocator就可以看出這一點；
9. BaseGPUDevice，GPUDevice：這兩個類都和GPU Device的實現有關，其中GPU Device類只是在繼承BaseGPUDevice的基礎上重寫了Allocator，但沒有理解這樣設計的深層次原因。
DeviceFactory相關類圖
上文提到過，TensorFlow中的Device是通過注冊機制添加到運行的進程中的。注冊機制在開源代碼中是十分常見的設計技巧，它涉及到了一種非常經典的設計模式——工廠模式。在定義每個Device時，通過利用C++事先定義好的宏（Macro）將類對象主動注冊到工廠中，這樣就可以達到在程序啟動完畢時，工廠里已經儲備有各種各樣所需要的內容。在TensorFlow中存在多處使用工廠模式的例子，比如本文闡述的Device管理，以及Session管理等。在其他開源框架中我們也能夠看到這一模式，比如Caffe中的Layer也使用的是工廠模式。
從源碼中可以看到，TensorFlow在啟動時會調用一系列static的函數，這些函數是通過宏（Macro）展開得到的。對於設備管理模塊來說，每種Device都由對應的Factory負責管理，而每種DeviceFactory會在程序啟動時注冊到全局唯一的static device factory表中。下面的類圖展示了各種DeviceFactory之間的繼承關系。
DeviceFactory的注冊機制
在理清了上一小節中Device相關類的繼承關系和說明之后，對於上圖中各種DeviceFactory之間的繼承關系就很好理解了。這里面比較陌生的類是Registrar，這可以看做是一個帶有模板的控制類，它只負責一件事——各種DeviceFactory向全局表的注冊。在代碼層面，注冊函數的調用是通過宏實現的，該宏通過傳入DeviceFactory的類名稱即可觸發Regsitrar的調用邏輯，在每個DeviceFactory的C++實現文件后面都會引用此宏。下圖形象的展示了注冊的過程。
有了上述的DeviceFactory的注冊后，就可以在使用時根據使用的Device類型，從對應的DeviceFactory中“獲取”想要的Device了。
Device的創建
有了DeviceFactory之后，我們就可以從Factory中拿到各種各樣的Device了。真正從Factory中取出Device的過程是在Session創建時進行的，調用的函數是DeviceFactory中的static函數AddDevices。它會遍歷全局device factories表中全部的DeviceFactory並取出，然后逐個調用每個具體XXDeviceFactory的CreateDevices函數，將創建的Device放進vector數組中。下面給出一個簡化版的時序圖。
上述的時序圖描述的較為簡單，實際上DeviceFactory的static函數調用AddDevices時會先將CPU Device創建出來，如果沒有可用的CPU Device，那么程序就會直接報錯退出（一般情況下不會發生此類情況）。這是因為TensorFlow需要保證當沒有其他Device存在時，至少還有CPU可以完成整體程序的計算和調度運行。創建CPU Device之后，就會去遍歷所有DeviceFactory，把所有能夠創建的Device全部創建出來放入vector數組中。
總結
本文主要闡述了TensorFlow設備管理模塊中的設備創建於注冊機制，它是TensorFlow進行設備管理的第一步，也是最簡單的部分。想要深入TensorFlow源碼的新人可以先從此模塊開始閱讀，進而熟悉TensorFlow的Coding style。Device的創建和注冊過程觸發於程序運行的初始階段，因為創建Device時使用了工廠模式，所以此處涉及到了各種DeviceFactory的定義和注冊。在TensorFlow的C++代碼中，各種DeviceFactory在實現文件中通過宏主動將自己注冊到全局表中，這樣做的目的不但減少了大量重復的注冊代碼，還與Device的創建解耦合，是一個非常經典常見的編碼技巧。至於Device的創建是在創建Session時才會觸發，該過程十分簡單。