僅此一文讓你掌握OneFlow框架的系統(tǒng)設(shè)計(jì)

OneFlow開(kāi)源近半年，近期發(fā)布了v0.3.2（https://github.com/Oneflow-Inc/oneflow/tree/v0.3.2）版本，相較于上個(gè)大版本，我們又新增了眾多算子和功能（如亞線性內(nèi)存優(yōu)化、Partial FC、足夠靈活易用的新版Checkpoint...），同時(shí)完備性（如Serving）、易用性（全新的API）也在快速推進(jìn)中，敬請(qǐng)期待。

兩個(gè)月前我們通過(guò)DLPerf（https://github.com/Oneflow-Inc/DLPerf）項(xiàng)目證明了OneFlow是世界上最快的深度學(xué)習(xí)框架：

成誠(chéng)：OneFlow是如何做到世界最快深度學(xué)習(xí)框架的

https://zhuanlan.zhihu.com/p/271740706

其中PK了經(jīng)NVIDIA深度優(yōu)化后的各個(gè)主流框架實(shí)現(xiàn)以及官方實(shí)現(xiàn)，對(duì)比了ResNet50和BERT這兩個(gè)CV和NLP領(lǐng)域應(yīng)用最廣的模型在數(shù)據(jù)并行下的吞吐率和加速比。這兩個(gè)月我們又做了大量的實(shí)驗(yàn)，對(duì)比了其他各個(gè)應(yīng)用場(chǎng)景下的第三方框架，如：

• 超大規(guī)模人臉識(shí)別案例 vs InsightFace
• Wide&Deep vs HugeCTR
• GPT-2 vs Megatron-LM
• SSP vs PipeDream
• Optimizer-Placement-Optimization vs DeepSpeed ZeRO
• Auto Parallelism vs FlexFlow

實(shí)驗(yàn)證明了OneFlow這一套簡(jiǎn)潔的抽象（SBP + Actor）在支持各種模型并行、混合并行、流水并行、自動(dòng)并行、ZeRO等方面是如此的簡(jiǎn)單高效。上述的每個(gè)特性，都有一個(gè)專門(mén)的第三方的框架項(xiàng)目對(duì)主流框架進(jìn)行魔改，至少涉及了數(shù)千行的項(xiàng)目代碼，而這些特性在OneFlow中要么是原生支持的特性，要么是新增一個(gè)Actor類(lèi)型或者是一個(gè)圖優(yōu)化的Pass（幾十行的代碼）就可以支持的非常好。我們認(rèn)為OneFlow這套設(shè)計(jì)不僅是性能最快的框架設(shè)計(jì)，同時(shí)也是分布式深度學(xué)習(xí)訓(xùn)練框架中最簡(jiǎn)單、最易擴(kuò)展的框架設(shè)計(jì)。上述的每個(gè)實(shí)驗(yàn)，近期都會(huì)有專門(mén)的技術(shù)文章分享，感興趣的小伙伴可以期待一下~

如果你對(duì)OneFlow這套致簡(jiǎn)致快的框架設(shè)計(jì)感興趣，或者對(duì)深度學(xué)習(xí)框架、分布式系統(tǒng)感興趣的話，本文就會(huì)讓你完全掌握OneFlow的系統(tǒng)設(shè)計(jì)。相信讀完這篇文章，你就會(huì)理解我們是如何看待分布式深度學(xué)習(xí)訓(xùn)練的，我們?yōu)槭裁匆@樣設(shè)計(jì)，這樣設(shè)計(jì)的好處是什么，以及我們?yōu)槭裁聪嘈臤neFlow這套設(shè)計(jì)是分布式深度學(xué)習(xí)訓(xùn)練框架的最優(yōu)設(shè)計(jì)。本文的主要內(nèi)容如下：

• 深度學(xué)習(xí)框架原理

• OneFlow系統(tǒng)架構(gòu)設(shè)計(jì)（簡(jiǎn)略版）

• OneFlow完整運(yùn)行流程 與 各模塊的交互方式

• 1. 分布式集群環(huán)境初始化
• 2. Python端搭建計(jì)算圖
• 3. 編譯期：OneFlow(JobSet) -> MergedPlan
• 4. 編譯期：Compiler(Job)->Plan
• 5. 運(yùn)行時(shí)：Runtime(Plan)

全文分上、中、下三篇。本文是上篇。

深度學(xué)習(xí)框架原理

深度學(xué)習(xí)框架是人工智能領(lǐng)域的“操作系統(tǒng)”，為深度學(xué)習(xí)相關(guān)的算法工程師提供一套簡(jiǎn)潔易用的用戶接口，使之能方便的搭建深度學(xué)習(xí)模型，進(jìn)行深度學(xué)習(xí)模型的訓(xùn)練、驗(yàn)證、測(cè)試、調(diào)參、遷移、部署、迭代開(kāi)發(fā)等工作。同時(shí)深度學(xué)習(xí)框架作為底層硬件跟算法工程師之間的中間件，要做到設(shè)備無(wú)關(guān)，使得算法工程師可以不用關(guān)心具體的計(jì)算設(shè)備、存儲(chǔ)設(shè)備的細(xì)節(jié)就能方便的開(kāi)發(fā)模型。

深度學(xué)習(xí)框架本質(zhì)上是一個(gè)基于張量（Tensor）之間的計(jì)算（Operator）表達(dá)式所組成的計(jì)算圖（Graph）編譯執(zhí)行引擎，提供了一系列張量的定義、一元操作、二元操作等數(shù)學(xué)原語(yǔ)，并根據(jù)反向傳播算法（Back Propagation）進(jìn)行梯度自動(dòng)求導(dǎo)以及模型更新。在大量數(shù)據(jù)分批次流入計(jì)算圖進(jìn)行模型訓(xùn)練之后，使得模型學(xué)習(xí)到數(shù)據(jù)中的內(nèi)在關(guān)聯(lián)關(guān)系，從而獲得對(duì)應(yīng)場(chǎng)景中的“智能”感知與判斷能力。

OneFlow系統(tǒng)架構(gòu)設(shè)計(jì)

OneFlow總體分為3個(gè)層次：Python前端、編譯期（Compiler）、運(yùn)行時(shí)（Runtime）。

• Python端是用戶接口，是OneFlow啟動(dòng)、編譯、運(yùn)行的入口，負(fù)責(zé)構(gòu)建邏輯圖（Job），且負(fù)責(zé)運(yùn)行時(shí)跟底層計(jì)算圖執(zhí)行引擎交互，包括發(fā)送控制指令（運(yùn)行一個(gè)global_function / job）、喂數(shù)據(jù)（input）、處理輸出（output，callback）。
• 編譯期（Compiler）負(fù)責(zé)將前端用戶的定義的邏輯上的計(jì)算圖進(jìn)行編譯，產(chǎn)出實(shí)際上的物理計(jì)算圖 （Plan）
• 運(yùn)行時(shí)（Runtime）負(fù)責(zé)根據(jù)Plan創(chuàng)建真正的執(zhí)行圖——即一個(gè)由Actor組成的去中心化流式計(jì)算圖，每個(gè)Actor各司其職，有的Actor負(fù)責(zé)接收Python端的控制信號(hào)，有的Actor負(fù)責(zé)加載數(shù)據(jù)，有的Actor負(fù)責(zé)初始化模型、計(jì)算、更新、存儲(chǔ)、傳輸...，有的Actor負(fù)責(zé)返還給Python端數(shù)據(jù)，數(shù)據(jù)在計(jì)算圖中流動(dòng)，實(shí)現(xiàn)深度學(xué)習(xí)的模型訓(xùn)練功能。

總體架構(gòu)圖如下圖所示：

OneFlow 系統(tǒng)架構(gòu)圖（簡(jiǎn)略版）

QA by @poohRui（https://github.com/poohRui）

想問(wèn)一下今天分享的時(shí)候提到，OneFlow以op為基本單元，相比于PyTorch以Tensor為基本單元，這兩個(gè)視角從整個(gè)框架的設(shè)計(jì)上來(lái)看有什么優(yōu)缺點(diǎn)。

我覺(jué)得這正是OneFlow跟PyTorch的重要區(qū)別。

• PyTorch以Tensor為基本單元，更符合算法工程師寫(xiě)Python腳本的直覺(jué)，是一種以面向?qū)ο蟮姆绞竭M(jìn)行模型搭建和訓(xùn)練，對(duì)Tensor進(jìn)行賦值、切片，就像numpy一樣易用。所以這種易用性更高。但在分布式情況下想把一個(gè)Tensor切分到不同機(jī)器上，需要手動(dòng)構(gòu)建傳輸過(guò)程，相當(dāng)于直接對(duì)物理編程，所以對(duì)分布式使用的門(mén)檻更高。Tensor為基本單元，有一個(gè)好處就是mutable修改Tensor的值，是非常直觀和簡(jiǎn)單的。
• OneFlow以O(shè)p為基本單元，模型搭建是用構(gòu)圖的方式實(shí)現(xiàn)的，單機(jī)單卡易用性不如以Tensor為基本單元，但分布式情況下對(duì)用戶的門(mén)檻更低。因?yàn)橛蠸BP可以讓用戶不用關(guān)心多機(jī)多卡情況下的Tensor切分和傳輸。SBP Signature是Op看待自己的輸入輸出Tensor的SBP，同一個(gè)Tensor，在其唯一的生產(chǎn)者Op和眾多消費(fèi)者Op眼中的SBP Parallel可能不一樣，中間的boxing可以讓框架來(lái)做。而單獨(dú)拿一個(gè)Tensor來(lái)說(shuō)，是沒(méi)有SBP這個(gè)概念的，SBP是基于Op的視角來(lái)說(shuō)的。以O(shè)p為基本單元，Tensor一般是不可修改的，僅該Tensor的生產(chǎn)者可以修改它。Tensor只是整個(gè)Op組成的計(jì)算圖中流動(dòng)的數(shù)據(jù)。如果在Op為基本單元的計(jì)算圖中搞mutable消費(fèi)，就會(huì)比較麻煩。

OneFlow完整運(yùn)行流程 & 各個(gè)模塊之間交互方式

我們通過(guò)介紹一次OneFlow完整運(yùn)行的流程來(lái)了解系統(tǒng)中的各個(gè)主要模塊是如何協(xié)同工作的。

1. 初始化環(huán)境（Env）

OneFlow是一個(gè)分布式計(jì)算系統(tǒng)，在Python前端啟動(dòng)時(shí)，第一件要做的就是初始化整個(gè)集群環(huán)境（Env）。環(huán)境由一個(gè)配置文件（EnvProto）（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/env.proto）所描述，里面包含了有多少臺(tái)機(jī)器，每臺(tái)機(jī)器的id、ip地址、控制端口號(hào)、數(shù)據(jù)傳輸端口號(hào)等信息。（Resource、MachineCtx等應(yīng)該被合并進(jìn)Env里）

• 如果是類(lèi)MPI方式啟動(dòng)，各個(gè)機(jī)器會(huì)執(zhí)行相同的Python腳本，每個(gè)機(jī)器在執(zhí)行腳本時(shí)會(huì)判斷得知自己的machine_id，從而知道自己是不是master：

• 如果不是master，則在Python腳本的入口就卡住，進(jìn)入cluster（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/cluster.cpp#L53）的WorkerLoop()中循環(huán)、等待、執(zhí)行集群中master發(fā)來(lái)的指令（Eager::Instruction）、邏輯圖集合（Lazy::JobSet）。
• 如果是master，則真正執(zhí)行python腳本，啟動(dòng)session、進(jìn)入global function、構(gòu)圖...

• 如果是以ssh & worker的方式啟動(dòng)（目前主要使用這種方式，未來(lái)會(huì)替換成類(lèi)MPI方式），則僅在master機(jī)器上啟動(dòng)了python進(jìn)程，master會(huì)把oneflow_worker可執(zhí)行程序通過(guò)ssh的scp命令拷貝到各個(gè)worker機(jī)器上，并根據(jù)配置執(zhí)行oneflow_worker程序，進(jìn)入WorkerLoop()的循環(huán)。

二者的區(qū)別：如果是類(lèi)MPI的方式啟動(dòng)，各個(gè)機(jī)器上都需要安裝oneflow的python包，每個(gè)機(jī)器上僅需要一份python腳本即可；而以ssh & worker的方式啟動(dòng)，需要把oneflow_worker的二進(jìn)制文件臨時(shí)拷貝到各個(gè)機(jī)器上，不需要python腳本。

環(huán)境啟動(dòng)時(shí)做了什么事呢？

1. 各個(gè)機(jī)器上啟動(dòng)了oneflow的進(jìn)程
2. 創(chuàng)建CtrlServer（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/control/ctrl_server.h）和CtrlClient（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/control/ctrl_client.h），并互相監(jiān)聽(tīng)直到每臺(tái)機(jī)器跟其他所有機(jī)器（包括自己）都建立了連接

Ctrl就是oneflow的控制平面（control plane），負(fù)責(zé)發(fā)送控制消息和數(shù)據(jù)，如master向worker發(fā)送JobSet、Plan等。

在OneFlow的Runtime階段，每個(gè)機(jī)器都會(huì)創(chuàng)建CommNet（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/comm_network/comm_network.h#L36）全局對(duì)象，這是oneflow的數(shù)據(jù)平面（data plane），運(yùn)行時(shí)各個(gè)機(jī)器上的Actor之間的消息通信、數(shù)據(jù)傳輸均通過(guò)數(shù)據(jù)平面發(fā)送。

控制平面使用rpc方式通信和傳輸數(shù)據(jù)，簡(jiǎn)單直接；數(shù)據(jù)平面通過(guò)高性能的網(wǎng)絡(luò)（epoll，或者ibverbs）通信和傳輸數(shù)據(jù)，效率更高。這里補(bǔ)充一句，在使用ibverbs（RDMA）構(gòu)建數(shù)據(jù)平面的過(guò)程中，RDMA的數(shù)據(jù)傳輸需要使用注冊(cè)內(nèi)存（pinned memory，又稱鎖頁(yè)內(nèi)存， page-locked memory）。而各個(gè)機(jī)器之間需要通信知曉各自的注冊(cè)內(nèi)存地址，這是通過(guò)控制平面rpc的方式傳輸注冊(cè)內(nèi)存的元信息的。見(jiàn)：IBVerbsCommNet::RegisterMemoryDone（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/comm_network/ibverbs/ibverbs_comm_network.cpp#L59）

2. Python端Job構(gòu)圖

在初始化環(huán)境之后，master上的python進(jìn)程會(huì)執(zhí)行用戶在global function中的構(gòu)圖代碼，生成Job（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/job.proto）

Job是對(duì)整個(gè)邏輯圖的基本描述，有兩個(gè)主要部分：net和placement。

• net是一個(gè)op list，表述了整個(gè)網(wǎng)絡(luò)是由哪些op以哪種方式連接起來(lái)的。net可以轉(zhuǎn)化成一個(gè)DAG，圖上的點(diǎn)表示一個(gè)op，圖上的邊表示op之間的產(chǎn)出和消費(fèi)的tensor。
• placement表示了每個(gè)op放置在哪些設(shè)備哪些卡上。對(duì)于env里的所有機(jī)器以及所有設(shè)備，任何一個(gè)op都可以以任何方式分布在這些機(jī)器上。placement表示了邏輯上的op跟物理上的op的映射關(guān)系。

Python端通過(guò)C++（oneflow_internal_helper.h -> c_api_util.py）暴露出來(lái)的接口，實(shí)際上使用JobBuildAndInferCtx（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/job_build_and_infer_ctx.h）的AddAndInferOp接口進(jìn)行推導(dǎo)。JobBuildAndInferCtx會(huì)保存已經(jīng)加入的Op及其相關(guān)狀態(tài)（SBP、shape等），并根據(jù)新加入的OpConf推導(dǎo)生成新的Op及其相關(guān)狀態(tài)。同時(shí)JobBuildAndInferCtx會(huì)給Python端提供一系列查詢接口，這樣在Python的global function中的構(gòu)圖邏輯，后一個(gè)op的python代碼在執(zhí)行前，之前所有的op和tensor（的描述，TensorDesc）都已經(jīng)構(gòu)建好了，這樣就實(shí)現(xiàn)了在global function中“類(lèi)似eager的方式構(gòu)圖。

Python端構(gòu)圖跟C++的交互流程

在整個(gè)global function中的代碼都執(zhí)行完之后，JobBuildAndInferCtx會(huì)被調(diào)用Complete（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/job_build_and_infer_ctx.cpp#L934），生成最終的用戶定義的Job。

在Complete過(guò)程中，會(huì)調(diào)用執(zhí)行多個(gè)JobPass，每個(gè)pass是對(duì)Job進(jìn)行一次圖修改、重寫(xiě)。其中最重要的pass就是生成后向op以及Optimizer的Pass（其中調(diào)用了AutoGrad）：“GenerateBackwardAndOptimizerOpConfs”。每個(gè)pass輸入是一個(gè)job，輸出是重寫(xiě)后的job。很多性能優(yōu)化的pass也是這個(gè)時(shí)期做，比如“FuseAddToOutputPass”、自動(dòng)混合精度"AutoMixedPrecision"等。

用戶可能會(huì)定義多個(gè)global function（如cnn的train job和eval job），所有用戶定義的Job構(gòu)成一個(gè)集合（JobSet）。而OneFlow的C++主體對(duì)象Oneflow就只接收一個(gè)JobSet對(duì)象啟動(dòng)Complier和Runtime。

3. 編譯期：OneFlow(JobSet) -> MergedPlan

由于歷史原因，Oneflow的Complier僅編譯單個(gè)Job，多Job的編譯、Job間內(nèi)存復(fù)用、MainPlan等均在Oneflow的CompileAndMergePlanOnMaster（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/oneflow.cpp#L897）接口中執(zhí)行。我們先假定Complier已經(jīng)將Job編譯成對(duì)應(yīng)的Plan了（Compiler的編譯過(guò)程我們放在后面講）。

下面我們介紹Oneflow生成最終的MergedPlan的流程：

輸入是用戶定義的多個(gè)job（已經(jīng)過(guò)前后向展開(kāi)以及各種圖優(yōu)化），我們稱之為user job。

3.1 構(gòu)建Model IO Job

Model IO Job中的每個(gè)Variable由全部user job中的Variable op name唯一確定。如果多個(gè)user job中有完全相同的Variable，則這兩個(gè)Variable是內(nèi)存共享的。即，Variable op的name是全局唯一的，是一個(gè)全局變量。舉例：train job和eval job中的同名Variable共享同一份內(nèi)存。

另外，Model IO Job（包含了三類(lèi) Model Init Job、Model Save Job、Model Load Job）中的Variable op，跟多個(gè)user job中的同名Variable也是內(nèi)存共享的。

目前，OneFlow里有兩種構(gòu)建Model IO Job的方式（MakeModelIoJobs（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/model_io_job.cpp#L218）， MakeModelIoV2Jobs（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/model_io_v2_job.cpp#L207）），分別表示這些Variable是共用一個(gè)Init/Load/Save op去處理，還是每個(gè)Variable單獨(dú)一個(gè)Op去處理。下圖展示了Model IO Job的幾種形式：

Model IO Job

請(qǐng)注意，Model IO中真正存放各個(gè)模型的Op類(lèi)型是Output，是InterfaceOp的一種。而不是Variable。Output1產(chǎn)出的Tensor::Var1跟其他某幾個(gè)user job中的VariableOp::Var1內(nèi)存共享。以此類(lèi)推。

OneFlow中有幾種類(lèi)型的InterfaceOp:

• Input（Python端的global function輸入Tensor）
• Output
• Return（Python端的global function的return Tensor）

Job之間的數(shù)據(jù)傳遞和綁定均通過(guò)InterfaceOp來(lái)實(shí)現(xiàn)。

InterfaceOp產(chǎn)出的Tensor的“RegstNum”恒為一，即僅有一份內(nèi)存塊，不支持流水；同時(shí)這塊內(nèi)存是被這個(gè)Tensor所獨(dú)占的，不會(huì)跟系統(tǒng)中的其他Op產(chǎn)出的Tensor內(nèi)存進(jìn)行內(nèi)存復(fù)用。

注：目前的Model IO是非常靜態(tài)的，非常不利于用戶對(duì)Checkpoint靈活使用的需求。@daquexian的新Model IO會(huì)徹底解決這個(gè)靈活性的問(wèn)題。PR見(jiàn)：Oneflow-Inc/oneflow#3540（https://github.com/Oneflow-Inc/oneflow/pull/3540）

3.2 構(gòu)建Push/Pull Job

遍歷所有user job中的Input Op和Return Op，針對(duì)每個(gè)Input Op，分別構(gòu)建一個(gè)對(duì)應(yīng)的Push Job，針對(duì)每個(gè)Return Op，分別構(gòu)建一個(gè)對(duì)應(yīng)的Pull Job。Push/Pull的原理見(jiàn)下圖：

Push Job和Pull Job

其中ForeignInput Op內(nèi)部維護(hù)一個(gè)buffer，該buffer等待Python端喂數(shù)據(jù)，當(dāng)數(shù)據(jù)喂完時(shí)該Op/Kernel執(zhí)行完畢。ForeignOutput Op內(nèi)部也有一個(gè)buffer，當(dāng)往該buffer內(nèi)填完數(shù)據(jù)以后，python端對(duì)應(yīng)的of blob對(duì)象中的numpy就拷貝了對(duì)應(yīng)的數(shù)據(jù)。

參見(jiàn)：

• ForeignInputKernel（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/kernel/foreign_input_kernel.cpp#L27） 與 ForeignOutputKernel（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/kernel/foreign_output_kernel.cpp#L27）
• Python端：OfBlob._CopyBodyFromNdarray()（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/python/framework/ofblob.py#L113）
  OfBlob._CopyToNdarrayListAndIsNewSliceStartMask()（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/python/framework/ofblob.py#L138）
• C++與Python端連接處：Dtype_GetOfBlobCurTensorCopyToBufferFuncName（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/python/oneflow_internal.e.h#L40）

為什么OneFlow與Python端的數(shù)據(jù)交換需要通過(guò)兩種獨(dú)立的Job子圖實(shí)現(xiàn)？

有兩個(gè)目的：

1. 新增Push/Pull Job，并使用內(nèi)存共享的方式，對(duì)原有的Job沒(méi)有構(gòu)圖上的破壞。
2. （主要目的）為了盡可能重疊Python與C++數(shù)據(jù)交換的過(guò)程。如何重疊？需要依賴OneFlow構(gòu)圖上的重要設(shè)計(jì)：MainJob和CriticalSection。我們放在下一節(jié)講。

3.3 編譯所有的job

順序編譯所有的user job和Model IO Job、Push/Pull Job。每個(gè)Job編譯時(shí)，都是用Compiler完整編譯至plan。且各個(gè)job之間不知道彼此的存在（歷史原因）。

Compiler將一個(gè)Job編譯成Plan的過(guò)程放在下一章節(jié)（中篇）講。

3.4 生成MainJob并得到最終的MergedPlan

這個(gè)過(guò)程分為幾步。

1. 將每個(gè)Job生成的Plan（SubPlan）合并到一個(gè)大的MergedPlan中
2. Job之間的內(nèi)存復(fù)用和內(nèi)存共享 （OneFlow中的內(nèi)存共享和內(nèi)存復(fù)用是一個(gè)很大的話題，我們后面會(huì)專門(mén)單獨(dú)寫(xiě)一篇文章分享其中的設(shè)計(jì)）
3. 計(jì)算CriticalSection（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/oneflow.cpp#L737）
4. 生成MainJob（https://github.com/Oneflow-Inc/oneflow/blob/v0.2.0/oneflow/core/job/oneflow.cpp#L573）
5. 編譯MainJob得到MainPlan
6. 將MainPlan和MergedPlan中每個(gè)Job生成的SubPlan進(jìn)行l(wèi)ink，得到最終的MergedPlan

CriticalSection

CriticalSection是OneFlow構(gòu)圖中一個(gè)非常重要的概念——臨界區(qū)。多個(gè)Job編譯的多個(gè)Plan分布在各個(gè)臨界區(qū)中。每個(gè)Job都關(guān)聯(lián)多個(gè)臨界區(qū)，臨界區(qū)有兩種類(lèi)型：InputOutput 和 Total。其中InputOutput是根據(jù)這個(gè)Job的Input、Output、Return等特殊類(lèi)型的Op專門(mén)設(shè)立的臨界區(qū)，Total是每個(gè)Job必有的臨界區(qū)，Job內(nèi)的所有Op都被包含在Total臨界區(qū)里。

整個(gè)JobSet會(huì)劃分成眾多臨界區(qū)，臨界區(qū)之間最重要的關(guān)系就是——互斥。如果兩個(gè)臨界區(qū)互斥，則其中一個(gè)臨界區(qū)在執(zhí)行的時(shí)候，另一個(gè)臨界區(qū)必須等待。如果兩個(gè)臨界區(qū)不互斥，則可以并行同時(shí)執(zhí)行。如何判斷兩個(gè)臨界區(qū)是否互斥？借助全局概念的Op——InterfaceOp和VariableOp，如果兩個(gè)臨界區(qū)中的全局Op有同名，則這兩個(gè)臨界區(qū)在執(zhí)行的時(shí)候會(huì)訪問(wèn)同一個(gè)全局的Op，則這兩個(gè)臨界區(qū)必然互斥，無(wú)法同時(shí)訪問(wèn)同一個(gè)全局Op。

臨界區(qū)是比Job更細(xì)粒度的概念（但跟Op相比，仍然是粗粒度的）。

• 為什么要把Job分成多個(gè)臨界區(qū)？
• 為什么要有InputOutput和Total兩種類(lèi)型的臨界區(qū)？

原因是想讓不同的Job之間盡可能流水并行起來(lái)。如何使得相鄰的兩個(gè)有消費(fèi)關(guān)系的Job（Job A -> Job B, A的output被B的input消費(fèi)/共享）同時(shí)執(zhí)行？借助CriticalSection以及MainJob里的幾個(gè)重要組件，我們就能實(shí)現(xiàn)多Job之間盡可能的并行執(zhí)行。

Idea by @Ldpe2G
臨界區(qū)的互斥可以區(qū)分讀寫(xiě)互斥，這樣多個(gè)只讀的臨界區(qū)可以并行執(zhí)行，在某些場(chǎng)景下可以更好的流水并行。

MainJob

MainJob的結(jié)構(gòu)圖如下：

Main Job結(jié)構(gòu)圖

Main Job 的結(jié)構(gòu)大體上反映了運(yùn)行時(shí)Python端跟OneFlow系統(tǒng)的交互情況：

1) Python端每調(diào)用一次global_function，都會(huì)向WaitAndSendIds op發(fā)送一個(gè)job id，WaitAndSendIds會(huì)把收到的job id對(duì)應(yīng)的多個(gè)CriticalSection id發(fā)送給ReentrantLock op。

2) ReentrantLock——可重入鎖，里面維護(hù)了所有臨界區(qū)之間的互斥情況，并且會(huì)維護(hù)一個(gè)等待隊(duì)列。其輸入有兩個(gè)：

• 一個(gè)是python端發(fā)來(lái)的控制指令說(shuō)要執(zhí)行哪個(gè)Job對(duì)應(yīng)的多個(gè)CriticalSection id，稱之為“start”輸入
• 另一個(gè)是esac返還回來(lái)的CriticalSection id，稱之為“end”輸入

start表示需要執(zhí)行哪個(gè)CriticalSection，end表示哪個(gè)CriticalSection已經(jīng)執(zhí)行完了。每個(gè)輸入都會(huì)更新可重入鎖中的等待隊(duì)列。由可重入鎖來(lái)判斷哪個(gè)CriticalSection可以被執(zhí)行。

舉例：start來(lái)了一個(gè)CriticalSection id 3。我們假設(shè)CriticalSection 3 與 0 互斥，且當(dāng)前CriticalSection 0 正在被執(zhí)行中，所以可重入鎖會(huì)讓3進(jìn)入等待隊(duì)列，直到0的執(zhí)行完畢信號(hào)還回來(lái)時(shí)（end 來(lái)了 0），3可以執(zhí)行了，那么才放3執(zhí)行。

3） ReentrantLock會(huì)根據(jù)內(nèi)部的臨界區(qū)互斥情況和等待隊(duì)列來(lái)判斷要向下發(fā)送真正可以立即執(zhí)行的CriticalSection id，發(fā)給Case op，Case Op執(zhí)行的就是一個(gè)switch的操作，觸發(fā)下面對(duì)應(yīng)id的CriticalSection去執(zhí)行。

4） MainJob的主體部分是所有的CriticalSection，注意在MainJob里每個(gè)CriticalSection用一個(gè)identity的tick op來(lái)標(biāo)識(shí)。當(dāng)整個(gè)MainJob編譯成MainPlan后，會(huì)執(zhí)行Link操作，將每個(gè)SubPlan連接替換MainPlan中的identity tick op。

5） 當(dāng)某個(gè)CriticalSection執(zhí)行完畢后，會(huì)給Esac op發(fā)消息。“Esac” 的命名是“Case”的字母逆序，因?yàn)槠涔δ芫褪歉鶦ase完全對(duì)稱相反的。Esac會(huì)把執(zhí)行結(jié)束的CriticalSection id發(fā)給ReentrantLock op（end輸入）用于更新?tīng)顟B(tài)。

另外圖中還有另外一個(gè)Esac op，僅連接了各個(gè)job對(duì)應(yīng)的Total Critical Section，該op接收某個(gè)Job執(zhí)行完畢的消息，并通過(guò)CallbackNotify Op發(fā)送給Python用于通知Python某個(gè)Job執(zhí)行結(jié)束了，可以執(zhí)行對(duì)應(yīng)的Callback了（如loss收集、acc統(tǒng)計(jì)等）。

ReentrantLock + CriticalSection 實(shí)現(xiàn)Job之間的流水并行

我們假設(shè)一種Job之間的消費(fèi)情況：Job A -> Job B，B消費(fèi)A的輸出，A和B均對(duì)應(yīng)了多個(gè)CriticalSection，A對(duì)應(yīng)0，1，2；B對(duì)應(yīng)3，4，5，其中1，4是TotalCriticalSection類(lèi)型，其余是InputOutputCriticalSection類(lèi)型。由于兩個(gè)Job僅在輸入輸出之間有消費(fèi)關(guān)系，所以僅有2，3互斥，其余均不互斥。互斥關(guān)系如下圖：

Job之間的流水并行

所以在ReentrantLock那里，僅會(huì)把2，3互斥相互block住。而兩個(gè)Job的主體：(A, 1), （B, 4）是不互斥的。故當(dāng)條件允許時(shí)，Job A和Job B可以流水并行執(zhí)行。如果沒(méi)有InputOutputCriticalSection，則A和B是一定會(huì)串行執(zhí)行的。

在OneFlow中，Push/Pull Job跟對(duì)應(yīng)的UserJob就是通過(guò)上述方式進(jìn)行流水并行的。

小結(jié)

本文是OneFlow系統(tǒng)設(shè)計(jì)分享文章的上篇，介紹了深度學(xué)習(xí)框架原理、OneFlow系統(tǒng)架構(gòu)的簡(jiǎn)略版，以及OneFlow完整運(yùn)行流程的前半部分：從環(huán)境初始化到編譯期MergedPlan生成。在下一篇《僅此一文讓你掌握OneFlow框架的系統(tǒng)設(shè)計(jì)（中篇）》中，我們會(huì)介紹編譯期Complier如何將Job編譯成Plan的過(guò)程，其中會(huì)簡(jiǎn)要介紹OneFlow編譯期最精華的Boxing章節(jié)，敬請(qǐng)期待~