深度解讀：讓你掌握OneFlow框架的系統(tǒng)設計(中篇)

本文是OneFlow系統(tǒng)設計分享系列文章的中篇，主要介紹OneFlow的編譯期Compiler如何將Job編譯為Plan的。其中最精華的部分是OneFlow的Boxing模塊，負責構(gòu)建兩個邏輯上的Op對應的兩組物理上的Op在任意情形下的物理子圖，完成了分布式訓練中各個機器各個設備之間的數(shù)據(jù)拷貝、切分、傳輸、通信的子圖搭建。值得一提的是，Boxing模塊的代碼實現(xiàn)是非常直觀且易擴展的，使用了設計模式中的責任鏈模式（Chain of Responsibility），未來我們會結(jié)合OneFlow的代碼實現(xiàn)分享一些C++編程技巧的文章，以及為什么OneFlow要使用這些編程技巧，解決了哪些問題，敬請期待~

如果你對OneFlow這套致簡致快的框架設計感興趣，或者對深度學習框架、分布式系統(tǒng)感興趣的話，本文就會讓你全面掌握OneFlow的系統(tǒng)設計。相信讀完這篇文章，你就會理解我們是如何看待分布式深度學習訓練的，我們?yōu)槭裁匆@樣設計，這樣設計的好處是什么，以及我們?yōu)槭裁聪嘈臤neFlow這套設計是分布式深度學習訓練框架的最優(yōu)設計。

本篇為系列內(nèi)容中的中篇，閱讀上篇，請點擊本次推送第一條。

目 錄

1. 深度學習框架原理

2. OneFlow系統(tǒng)架構(gòu)設計（簡略版）

3. OneFlow完整運行流程與各模塊的交互方式

3.1 分布式集群環(huán)境初始化

3.2 Python端搭建計算圖

3.3 編譯期：OneFlow(JobSet) -> MergedPlan

3.4 編譯期：Compiler(Job)->Plan

3.5 運行時：Runtime(Plan)

3.4 編譯期：Compiler(Job)->Plan

為了方便理解，我們再簡要描述一些重要的概念和抽象：

• Job：用戶定義的邏輯上的計算圖，由邏輯上的Op組成。

• Plan：編譯生成的物理上的計算圖，由物理上的Task組成。

• Task：運行時Actor的配置描述，一個Actor與一個Task一一對應，Task內(nèi)部有Op的運行時描述Kernel的配置。Task并不一定關(guān)聯(lián)某個用戶計算圖Job中的邏輯上的Op，因為并編譯期會增加很多物理上的Op用于數(shù)據(jù)搬運、網(wǎng)絡傳輸、切分/拼接等操作。Task中標記了自己是在哪臺機器哪個設備上，并使用哪個線程工作等。Task還需要指定自己產(chǎn)出Regst的regst num、內(nèi)存類型、所屬內(nèi)存塊的偏移量等等信息。

• Regst：運行時的數(shù)據(jù)存儲、Actor之間通信的基本單元。存儲某個具體的Tensor。

編譯期（Compiler）的設計體現(xiàn)了OneFlow作為一個分布式深度學習訓練框架的很多重要的設計原則：

1）一致性視角（Consistent View）

OneFlow把整個分布式集群抽象成為一個超級設備，用戶使用OneFlow做分布式訓練跟做單機單卡的訓練沒有任何區(qū)別。體現(xiàn)在：

• 編譯期Compiler僅在Master機器上編譯整個Plan，其他Worker機器等待獲取Plan啟動運行時即可。Master上編譯的Plan就包含了所有機器所有設備上的基礎單元——Task（Actor）。

• 所有的分布式訓練過程中各個機器各個設備之間的數(shù)據(jù)通信、同步操作均被Compiler自動生成，無需用戶關(guān)心和編寫分布式訓練中的數(shù)據(jù)同步。

2）數(shù)據(jù)搬運是一等公民

OneFlow將所有的數(shù)據(jù)加載、預處理、拷貝、網(wǎng)絡傳輸、Tensor的自動切分/拼接/廣播/求和操作都抽象成了跟計算Op一樣的運行時執(zhí)行體——Actor，即在分布式的物理計算圖上顯式表示了數(shù)據(jù)搬運的操作。這樣做的好處是OneFlow可以感知到所有的數(shù)據(jù)搬運、同步操作，因此編譯期Compiler可以更好的在整個物理計算圖上做全局調(diào)度，使得這些數(shù)據(jù)搬運操作盡可能被計算操作所掩蓋，對數(shù)據(jù)搬運操作的性能優(yōu)化轉(zhuǎn)而變成了圖分析與圖優(yōu)化。

3）編譯期全局調(diào)度，運行時去中心化調(diào)度

OneFlow的運行時是一個極其簡單的抽象——Actor，每個Actor僅需要關(guān)心和自己相關(guān)的上下游Actor的消息就可以知道自己能否工作，這樣做的好處是運行時系統(tǒng)不會因為有中心調(diào)度節(jié)點導致性能瓶頸（在計算圖非常大的情況下）。為了做到這一點，OneFlow的調(diào)度工作大多都是在編譯期完成的，Compiler會做好全局的內(nèi)存調(diào)度、Op執(zhí)行調(diào)度、通信調(diào)度等工作，使得運行時的調(diào)度開銷盡可能的低，從而達到更快的訓練速度。

4）天然支持流水線，解決流控問題

編譯期Compiler通過推導和設置Task產(chǎn)出Regst的regst num，可以使得運行時相鄰Actor之間可以流水并行起來。同時還可以通過背壓機制（Back Pressure）解決流控問題（Control Flow）。具體的Actor機制如何解決流水線和流控問題的討論我放在下篇中介紹。

在原生的OneFlow設計中，Compiler輸入是一個Job（用戶定義的op list），經(jīng)過編譯生成OneFlow的中間表示IR（Intermediate Representation）——Plan，Plan是一個被Runtime直接讀取就能生成運行時執(zhí)行圖的描述。而上面介紹的OneFlow(JobSet)->MergedPlan是為了支持Python前端交互 + 多Job（Train/Eval同時做）而后設計出來的。我們下面介紹OneFlow的Compiler做了哪些事。

3.4.1 JobCompleter().Complete(job)

第一步，經(jīng)過JobCompleter將Job不斷重寫。經(jīng)過多個Pass以生成最終的Job。中間借助OpGraph抽象不斷推導新的Job對應的邏輯圖。這些Pass包括一些優(yōu)化如插入KeepHeaderOnly節(jié)點；增加Source/Sink的Tick節(jié)點使得圖成為一個單源節(jié)點和單匯節(jié)點；增加控制邊；計算臨界區(qū)；以及使用XRT框架重新構(gòu)建Job。

XRT框架會將Job中的OpGraph進行有選擇的合并，并選取使用XLA或者TensorRT來進行編譯生成優(yōu)化后的Kernel。對于OneFlow而言，這些都是XrtLaunchOpConf，其Kernel都是XrtLaunchKernel。OneFlow系統(tǒng)并不關(guān)心其實現(xiàn)細節(jié)，實際上，經(jīng)過XRT優(yōu)化后的Kernel實現(xiàn)都是在其框架內(nèi)定義的頂層抽象：Executable 中存儲的，在XrtLaunchKernel的計算過程中調(diào)用executable->Run()去執(zhí)行。

3.4.2 生成OpGraph

Graph是OneFlow中的一個重要基礎抽象，各個重要的圖相關(guān)的概念（OpGraph、LogicalGraph、TaskGraph、ChainGraph、ExecGraph...）都繼承自Graph。Graph表示一個圖，里面保存著這個圖中的所有的節(jié)點Node和節(jié)點之間的連邊Edge。Graph上面提供一系列共用的遍歷方法（普通遍歷、拓撲遍歷、BFS、DFS...），以及圖改寫（插入、刪除 節(jié)點/邊）圖查詢方法。

其實在第一階段JobCompleter在修改Job的過程中就需要多次Build OpGraph，在最終版本的Job生成以后，我們還需要在全局創(chuàng)建一個OpGraph，用于后續(xù)編譯過程中對各個邏輯Op和邏輯Tensor的查詢。

生成OpGraph分為幾步：（核心邏輯：OpGraph::InferLogicalBlobDesc https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/graph/op_graph.cpp#L563）

按照拓撲序遍歷每個Op（OpNode）

1） 推導ParallelSignature （Eager所需）

2） 推導BatchAxis（將要被廢棄，描述了哪一個維度是batch維，或者沒有batch維，如Variable那一支路上的op）

3） 推導MirroredSignature （推導每個Tensor是否是Mirrored，我認為這個應該跟Sbp成為同一級的東西：SBPM）

4） 推導SbpSignature

SBP是oneflow非常重要的概念，我在知乎文章——《都2020年了，為什么我們相信OneFlow會成功》 中有初步解釋了SbpParallel的語義：一種邏輯上的Tensor跟物理上的多個Tensor的映射關(guān)系。SbpSignature是一個SbpParallel的集合，在OneFlow的設計里是Op的屬性，它描繪了一個邏輯上的Op被映射成各個設備上的多個物理上的Op以后，這些物理上的Op是如何看待他們輸入輸出Tensor在邏輯上和物理上的映射關(guān)系的。

這里的推導SbpSignature，就是在每個Op多個合法的SbpSignature中搜索到一個最優(yōu)的（傳輸代價最低的）作為本次訓練實際采用的SbpSIgnature。

在自動并行（by @Yipeng1994 ）完成以后，推導SbpSignature就不再是按照拓撲序貪心算法推導，而是在全局搜索一個近似次優(yōu)解。

5） 推導Logical BlobDesc

此處是推導每個邏輯Op的邏輯Tensor的Shape、DType、is_dynamic等信息。

Op最重要的概念就是推導SBP，并根據(jù)SBP來推導Tensor的Shape。編譯期僅需要靜態(tài)推導出每個Tensor的形狀，以及特殊Op需要推導其Op/Kernel的特殊屬性：Inplace、TempBufferSize...

OpGraph是邏輯上的概念，當OpGraph構(gòu)建完成后，每個（邏輯上的）Op、每個（邏輯上的）Tensor的描述信息都被推導、創(chuàng)建完成了。

3.4.3 生成LogicalGraph 【即將過時】

LogicalGraph是OneFlow的歷史遺留產(chǎn)物，在遠古時期負責邏輯圖展開、后向生成、Model IO等工作。后面隨著OneFlow系統(tǒng)設計的演化，其功能逐步被OpGraph + JobCompleter + Pass所替代。之所以目前還保留，是因為Op與TaskType的映射關(guān)系還保留在LogicalNode的不同子類中。在未來一段時間內(nèi)會移除掉LogicalGraph抽象，完全由OpGraph所取代。

3.4.4 生成TaskGraph

TaskGraph的生成過程是OneFlow編譯期最重要也是最精華的一部分。Task是Actor的編譯期抽象，Actor是Task的運行時抽象。所以TaskGraph就描繪了整個運行時計算圖的全貌。TaskGraph的生成過程分為兩部分：

• 構(gòu)圖部分 （https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/graph/task_graph.cpp#L156）

• Build/Infer部分 （https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/compiler.cpp#L77）

3.4.4.1 構(gòu)圖部分

如何根據(jù)邏輯圖生成物理計算圖？

1） 遍歷LogicalGraph的每個LogicalNode，根據(jù)每個LogicalNode的placement：生成有序的ComputeTaskNode（專用于計算的TaskNode）

2） 遍歷LogicalGraph的每個LogicalEdge，根據(jù)前后LogicalNode的類型，找到對應的生成這部分SubTaskGraph的方法：BuildSubTaskGraphMethod，執(zhí)行該方法給這兩個LogicalNode對應的ComputeTaskNode連邊、新增節(jié)點構(gòu)圖。

在遠古時期的OneFlow設計中，生成SubTaskGraph的方法是跟前后LogicalNode的類型相關(guān)。而在后面的boxing重構(gòu)中（by: @liujuncheng, 見：Oneflow-Inc/oneflow#2248 和 Oneflow-Inc/oneflow#2846 等），生成SubTaskGraph的方法被SubTskGphBuilder所推導，根據(jù)情況構(gòu)建紛繁多樣的SubTaskGraph。下面會粗略介紹一下其中的設計。

下圖展示了一種可能的SubTaskGraph構(gòu)建方式：在LogicalGraph中，邏輯上的Op_a產(chǎn)出一個Tensor X 供Op_b消費，其中Op_a和Op_b的Placement分別是4，3，而Op_a和Op_b對X的SBP parallel根據(jù)各自Op的屬性、用戶指定/推導/自動SBP的結(jié)果確定。Tensor X就是一條LogicalEdge。第一步：分別生成所有LogicalNode對應的有序的ComputeTaskNode，Op_a的LogicalNode展開成4個CompTaskNode，Op_b的LogicalNode展開成3個ComputeTaskNode。第二步：這些ComputeTaskNode在Boxing架構(gòu)中會根據(jù)實際情況新增節(jié)點，并連邊，使得下面Op_b的3個TaskNode可以拿到其想要的那部分X的數(shù)據(jù)。

每個LogicalEdge就是一個邏輯上的Tensor，前后兩個邏輯上的Op對同一個Tensor的SBP、Placement看待可能一致也可能不一致。如何構(gòu)建這部分SubTaskGraph對應的子圖呢？OneFlow提供了一系列SubTskGphBuilder，根據(jù)各種情況生成不同的子圖。

SubTskGphBuilder

構(gòu)建該子圖需要的全部信息是：源節(jié)點的CompTaskNode列表，匯節(jié)點對應的CompTaskNode列表，源節(jié)點與匯節(jié)點的并行屬性（ParallelDesc，SBP），傳輸?shù)倪壿婽ensor的信息（Shape、Dtype、LogicalBlobId...）

目前OneFlow內(nèi)部有7種SubTskGphBuilder，每種Builder下面都可以根據(jù)SBP、Placement等信息自定義多種實際的構(gòu)圖方案，如SliceBoxingSubTskGphBuilder下面就有5種不同的構(gòu)圖情況，CollectiveBoxingSubTskGphBuilder下面又有7種集合通信的Builder。我們這里簡單介紹幾個常見的子圖構(gòu)建方式：

1） one to one

這是最常見的連接方式，即LogicalEdge的兩端節(jié)點在ParallelNum、SBP上對中間邏輯Tensor的看待方式完全一致，可以一對一的直連。在常見的數(shù)據(jù)并行情況下（如Mirror的方式），前后向Op都是一對一直連的。下圖展示了兩種一對一直連情況。

左邊是GPU內(nèi)部的一對一直連，右邊是當兩個ComputeTaskNode不在同一個設備上時，我們會插入傳輸節(jié)點：CopyD2H（Device to Host），CopyCommNet（網(wǎng)絡傳輸），CopyH2D（Host to Device），使得一對一直連的匯節(jié)點CompTaskNode可以拿到對應的Tensor。

2） collective boxing

集合通信（Collective Communication）大多采用NCCL的實現(xiàn)，包含了：AllReduce、ReduceScatter、AllGather、Reduce、Broadcast等操作。需要注意的是：

• 由于NCCL多個設備上的通信是在NCCL內(nèi)部實現(xiàn)的，在OneFlow的TaskGraph上，這些NcclTaskNode之間沒有顯式的連邊，但其實中間有隱含的同步操作。這樣如果在NCCL結(jié)點前后連控制邊不當，可能會造成死鎖，所以系統(tǒng)中對NCCL附近的順序化連邊需要非常小心。

• 使用NCCL進行集合通信操作，在構(gòu)圖上是one to one連接的。

下圖展示了OneFlow中使用NCCL進行集合通信的collective boxing操作構(gòu)圖。

3） slice boxing

這種boxing涵蓋了oneflow中遇到的大多數(shù)跟SBP相關(guān)的Boxing。在SliceBoxingSubTskGphBuilder中提供了支持S2B、S2S、P2S、P2B、B2S等5種不同的SBP情形。

slice boxing會根據(jù)上下游兩組CompTaskNode的ParallelDesc、SBP的不同，把上面一組物理上的Tensor按照下游期望的SBP的方式分配給下游的一組CompTaskNode，同時考慮Machine id、CPU/GPU的不同，同時希望傳輸開銷、構(gòu)圖開銷盡可能少。

下圖展示了一種可能的S2S的slice boxing 情形。邏輯圖上SrcOp產(chǎn)出Tensor X供DstOp消費，其中SrcOp在Machine0的GPU0、1以及Machine1的GPU2、3上產(chǎn)出SBP Parallel = Split(0) 的Tensor X，而DstOp在Machine3的GPU4、5以及Machine 4的GPu6上消費SBP Parallel = Split(0)的Tensor X。故需要把已經(jīng)分成4份的Tensor X 先concat起來，然后再split成3份分發(fā)給各個DstOp。

我們再舉一個可能的P2B的例子。Src Op產(chǎn)出的兩個物理Tensor X分別是邏輯Tensor X的一部分值，經(jīng)過Add和Clone操作發(fā)給后面以Broadcast消費X的兩個DstOp。

?

OneFlow中的Boxing設計是其分布式易用性以及分布式性能上最精華的一部分，這里僅介紹了其概況，后續(xù)會單獨出一篇文章分享其中的設計。

我們通過SubTaskGraphBuilder給每個LogicalEdge對應的物理子圖構(gòu)圖，這樣就搭建起了整個TaskGraph，完成了邏輯圖到物理圖的映射。構(gòu)圖過程中，根據(jù)節(jié)點類型等信息可以給每個TaskNode分配Thread id、Area id等屬性。

Thread id 標記了每個TaskNode（即Actor）工作在哪個線程上。由于分布式環(huán)境下每個機器上是一個進程，所以每個TaskNode都會設置Machine id和Thread id。線程id分配的方式：CPU上是平均分配各個thread id；GPU上，同一個GPU的所有計算Task在同一個計算線程中；所有集合通信的Task在同一個NCCL線程中。這樣分配線程id的方式是因為經(jīng)過實驗驗證，計算Task在相同線程中速度最快（最小切換開銷）。

Area id 【即將過時】標記了不同類型的Op、TaskNode分別從屬于整個TaskGraph上的哪一個區(qū)域。有一些特殊的Area如kMdUpdtArea 標記了這些Task是在Optimizer子圖部分的。然而Area id是一個過時設計。應該被完備的Scope概念所取代，同時一個Op從屬于哪個Area也不是Op的類型決定的，而是Op在編譯期圖重寫的哪一個階段被插入所決定的。后續(xù)會把Area id移除。

ChainGraph?【即將過時】

在目前的設計中，TaskGraph還會生成ChainGraph，進行Chain的合并，給每個Task上新增Chain id的屬性，用于將一組Task子圖標記出來（方便做內(nèi)存復用）。

被合并到一個Chain中的這組Task有一個共性：在相同的Thread/Stream中執(zhí)行，當Chain子圖中的源節(jié)點可以執(zhí)行以后，Chain子圖的所有后繼節(jié)點可以一股腦的執(zhí)行完，不需要依賴或者再等其他的節(jié)點。

Chain的合并算法在遠古時期以Layer為單位進行遍歷合并時是可以較好工作的，但是目前以Op為單位就顯得有些過時，尤其是在一些特殊的網(wǎng)絡（如包含where op）中會因為圖的拓撲遍歷順序的不同而有較大的合并效果差異，甚至是成環(huán)的BUG。

目前僅在內(nèi)存復用算法中依賴了Chain的合并結(jié)果。后續(xù)會重構(gòu)掉這塊，將Chain的概念從TaskNode中去掉。

3.4.4.2 TaskGraph的Build/Infer階段

在TaskGraph的構(gòu)圖完畢之后，Compiler會按照TaskGraph中TaskNode的拓撲序遍歷，依次構(gòu)建每個TaskNode對應的各種信息：

1） 生成每個TaskNode的所有Regst，并把Regst綁定到TaskNode的出邊TaskEdge上。TaskNode::ProduceAllRegstsAndBindEdges （https://github.com/Oneflow-Inc/oneflow/blob/64c20462f245b5cbef4230a62fa06edff85411b3/oneflow/core/job/compiler.cpp#L77）

2） 將每個TaskNode的入邊TaskEdge中的Regst關(guān)聯(lián)到TaskNode中

3） 執(zhí)行每個TaskNode的Build過程

TaskNode

TaskNode根據(jù)其不同的TaskType 有對應的TaskNode子類特化。每種類型的TaskNode其構(gòu)建過程都不同。最常見的是NormalForwardCompTaskNode，對應了所有用戶定義的計算Op的Actor。每種TaskNode對應一種Actor，其Actor內(nèi)部執(zhí)行的狀態(tài)機也不同。oneflow/core/graph/路徑下列出了目前所有種類的TaskNode子類及實現(xiàn)。

TaskNode的構(gòu)建過程中，內(nèi)部需要構(gòu)建Regst。

Regst

Regst是OneFlow中數(shù)據(jù)存儲、傳遞的基本單元。運行時Actor之間的消息通信，數(shù)據(jù)傳遞都使用Regst。在目前的Regst設計中，一個Regst會包含多個Blob（廣義上的Tensor概念），但越來越多的需求是需要一個Regst僅包含一個Blob，后續(xù)的重構(gòu)中，會把Blob概念整合進Tensor中，精簡這里的概念。

Tensor是用戶級別的概念，是獨立的的一塊數(shù)據(jù)，而Regst是Actor級別的概念，記錄了這個Regst是由哪個Actor生產(chǎn)的，并被哪些Actor所消費的。

TaskNode的Build過程

TaskNode內(nèi)部會有一個ExecGraph（執(zhí)行子圖），執(zhí)行子圖上的節(jié)點稱之為ExecNode，邊稱之為ExecEdge。在遠古的OneFlow設計中，每個TaskNode里是由多個Op組成的執(zhí)行子圖構(gòu)成的，每個Op對應一個ExecNode，后面隨著性能優(yōu)化變成了一個TaskNode對應一個Op。我們?nèi)匀槐Ａ袅薊xecGraph的設計，雖然在目前的絕大多數(shù)場景中ExecGraph里只有一個ExecNode，沒有ExecNode。

ExecNode 和 Op 的區(qū)別：（雖然我不止一次希望把ExecNode和Op合并）

在OneFlow最初的設計中，Op是一個描述概念，并不關(guān)心具體的某個Blob/Tensor，僅提供一系列方法用于推導，Op是無狀態(tài)的。而ExecNode是在某個具體的TaskNode內(nèi)部，同時要關(guān)聯(lián)具體的Regst，是有狀態(tài)的。

1） ExecGraph：絕大多數(shù)TaskNode的Build過程都是根據(jù)LogicalNode中的Op（CompTaskNode）/ 新建Op （CopyTaskNode），先構(gòu)建ExecGraph。

2） ExecNode：bind regst。在TaskNode中，入邊消費的Regst和出邊生產(chǎn)的Regst內(nèi)部都維護了一個或多個lbi（logical blob id），用于標識一個Blob（Tensor）。TaskNode的構(gòu)建過程中需要把這些Regst里的lbi跟Op內(nèi)部的BnInOp綁定起來。

3） ExecNode：InferBlobDesc。推導每個Blob/Tensor的Shape等信息（存儲在BlobDesc中）

3.4.4.3 TaskGraph & Plan 優(yōu)化

在TaskGraph Build結(jié)束以后，原本的Compiler還會對整個TaskGraph進行一些優(yōu)化。在前后向分離、python前端的重構(gòu)中，Compiler這里的優(yōu)化被精簡成了幾步：

1） 移除空的Regst

2） 增加Chain內(nèi)的控制邊保證執(zhí)行順序

3） 推導Inplace的內(nèi)存共享

Inplace的推導使用了 InplaceLbiGraph 進行推導。需要注意的是，我們在Op（UserKernel）里定義的SetInplaceProposalFn 僅是一種“建議”，而實際上這個Op的輸出和輸入能否Inplace，還需要經(jīng)過InplaceLbiGraph進行推導以后才能決定。一些顯而易見的約束是，一個Tensor不能同時被兩個消費它的Op進行Inplace，因為Inplace會改寫輸入的Tensor數(shù)據(jù)，是一種Mutable消費。Inplace在一些情況下可以加速計算。

4） 推導時間形狀（time shape）

在OneFlow的Regst中，除了其中的數(shù)據(jù)有物理上的形狀，Regst本身也有時間形狀（time shape），表示整個網(wǎng)絡執(zhí)行一個Batch的數(shù)據(jù)，該Regst需要被生產(chǎn)幾次。time shape 有2維，最常見的是（1， 1），表示一個batch執(zhí)行一次。一些特殊的Op/Actor會修改時間形狀：Repeat/Acc、Unpack/Pack。由于這些特殊Op可能會嵌套，所以我們讓時間形狀有兩維，表示最多允許兩層嵌套。當網(wǎng)絡中插入一個Repeat Op，會把該Tensor重復發(fā)送k次，其時間形狀就是（1，k）。當網(wǎng)絡中插入Unpack Op，會把一個Tensor切分成k段，按k次分別發(fā)送給后面的Op（相當于在時間上一種數(shù)據(jù)并行）。

如果網(wǎng)絡中連續(xù)插入多個RepeatOp，比如第一個Repeat將輸出的時間形狀修改為（1， k₁），后續(xù)的Regst均為該時間形狀；再插入第二個Repeat，則輸出的時間形狀會被修改為（ k₂， k₁）。

3.4.4.4 生成Plan

最終TaskGraph中的每個TaskNode會生成Plan中的TaskProto，得到一個naive的Plan。

Plan里最重要的內(nèi)容就是所有的TaskProto，每個TaskProto就描述了運行時的一個Actor所需的所有信息。

3.4.5 Improver(naive_plan) -> complete_plan

在naive的Plan生成之后，Improver會把Plan進行改寫。

Improver的最初設計是為了推導RegstNum。

在我之前的兩篇知乎文章中，都提到了運行時Actor機制的相鄰Actor流水線是通過RegstNum > 1來實現(xiàn)的。naive_plan中沒有推導RegstNum，所以所有的RegstNum均=1。而Improver中設計了一套算法，用于推導每個TaskNode對應的RegstNum，但是算法依賴每個Actor的實際執(zhí)行時間。所以需要有試跑。

在TaskGraph中我們hack了代碼，使得所有的CopyHdTaskNode的MinRegstNum=2，也就是RegstNum=2，目的是為了讓數(shù)據(jù)預處理跟GPU計算可以流水并行起來，未來會刪除掉這個hack。TaskGraph上對于每個Regst都推導了其min、max的regst num，一般的數(shù)據(jù)Regst min = 1， max = inf。也有的regst，我們不希望有任何多余的備份，故讓這些Regst min = 1，max = 1。

由于試跑對于后續(xù)的OneFlow開發(fā)非常不友好，于是Improver這里的試跑一直都沒有被啟用。而且即使所有的RegstNum = 1，在非相鄰的兩個Actor之間也可以流水并行起來。

目前Improver中最重要的目的是為了推導內(nèi)存復用。內(nèi)存復用也經(jīng)歷了多個階段，一開始是使用一種染色算法對Regst進行染色，相同顏色的共用一段內(nèi)存。后續(xù)我設計開發(fā)了內(nèi)存復用2.0（見 Oneflow-Inc/oneflow#2267、Oneflow-Inc/oneflow#2319），采用了Chunk、MemBlock、Regst三級內(nèi)存結(jié)構(gòu)，仍使用Improver作為入口。所以complete_plan中會比na?ve plan新增了Chunk和MemBlock的信息。OneFlow中的內(nèi)存復用設計后續(xù)會單獨出一篇文章進行分享。

后續(xù)會將內(nèi)存復用算法放在Compiler中，使得Compiler的結(jié)果就是最終的Plan。

至此，我們就描述清楚了如何從一個Job編譯成一個Plan的全過程。

本文是OneFlow系統(tǒng)設計分享文章的中篇，主要介紹OneFlow完整運行流程的中間部分：編譯期Compiler將Job編譯成Plan的過程。在下一篇《僅此一文讓您掌握OneFlow框架的系統(tǒng)設計（下篇）》中，我們會介紹OneFlow的運行時（Runtime）以及倉庫源碼下的主要各目錄的模塊簡介，其中會包含Actor運行時如何高效的調(diào)度，以及如何解決流水線和流控問題。閱讀下篇，請點擊第三條推送。

點擊“閱讀原文”，前往OneFlow代碼倉庫。