tensorflow源碼解析之common_runtime-executor-下

目錄

1. 核心概念
2. executor.h
   1. Executor
   2. NewLocalExecutor
   3. ExecutorBarrier
3. executor.cc
   1. structs
   2. GraphView
   3. ExecutorImpl
   4. ExecutorState
   5. details

3.4 ExecutorState

在執行器的執行圖計算的時候，需要一個結構來儲存當前計算的即時資訊，TF為此設計了類ExecutorState，它被用來儲存每一個對ExecutorImpl::Run調用的狀態資訊。它會在一個節點已經準備好之後調度這個節點，並且儲存每個節點尚未完成的輸入資訊。
下面讓我們先來看一下這個類的結構：

class ExecutorState {  public:    ExecutorState(const Executor::Args& args, ExecutorImpl* impl);    void RunAsync(Executor::DoneCallback done);  private:    DeviceContextMap device_context_map_;        typedef gtl::InlinedVector<TaggedNode, 8> TaggedNodeSeq;    typedef gtl::InlinedVector<Entry, 4> EntryVector;        const bool vlog_;    const bool log_memory_;    int64 step_id_;        //未擁有    Rendezvous* rendezvous;    SessionState* session_state_;    TensorStore* tensor_store_;    //每個執行步層級的容器    ScopedStepContainer* step_container_;    StepStatesCollector* stats_collector_;        checkpoint::TensorSliceReaderCacheWrapper* slice_reader_cache_;    FunctionCallFrame* call_frame;    const ExecutorImpl* impl_;    CancellationManager* cancellation_manager_;    Executor::Args::Runner runner_;    bool sync_on_finish_;        //擁有    bool dumped_on_error_ = false;    //當前執行步驟開始的幀    FrameState* root_frame_;    //執行器結束時需要調用的回呼函數    Executor::DoneCallback done_cb_;    std::atomic_int_fast32_t num_outstanding_ops_;    mutex mu_;    Status status_ GUARDED_BY(mu_);        //從幀名稱到實際幀的映射。在當前幀的某個迭代周期內，可能會產生一個新的幀。新的子幀的唯一索引值必須由父幀的名稱、迭代編號、以及由nodedef推斷出來的新幀的名稱組合而成    gtl::FlatMap<string, FrameState*> outstanding_frames_ GUARDED_BY(mu_);        //一個幀的名稱    inline string MakeFrameName(FrameState* frame, int64 iter_id, const string& name);        //找到一個現存的幀，或者建立一個新幀，在幀frame的iter迭代周期    void FindOrCreateChildFrame(FrameState* frame, int64 iter, const Node* node, FrameState** child);        //刪除一個幀，當幀調用結束時使用    void DeleteFrame(FrameState* frame, TaggedNodeSeq* ready);        //清除那些起源於幀frame和迭代iter的幀，當一個子幀結束時調用    void CleanupFramesIterations(FrameState* frame, int64 iter, TaggedNodeSeq* ready);        //在當前的線程中處理一個已準備好的節點    void Process(TaggedNode node, int64 scheduled_usec);        //在調用item->kernel之前，先填入其輸入    Status PrepareInputs(const NodeItem& item, Entry* first_input, TensorValueVec* inputs, DeviceContextVec* input_device_contexts, AllocatorAttributeVec* input_alloc_attrs, bool* is_input_dead);        //在item->kernel計算結束之後，處理輸出    Status ProcessOutputs(const NodeItem& item, OpKernelContext* ctx, EntryVector* outputs, NodeExecStats* stats);        //在處理完輸出之後，將輸入傳遞給下一個輸入    void PropagateOutputs(const TaggedNode& tagged_node, const NodeItem* item, EntryVector* outputs, TaggedNodeSeq* ready);        //節點計算結束後，接管stats，如果執行完成則返回true    bool NodeDone(const Status& s, const Node* node, const TaggedNodeSeq& ready, NodeExecStats* stats, TaggedNodeReadyQueue* inline_ready);        //調度ready中的所有複雜節點，然後將ready中的非複雜節點放入inline_ready    void ScheduleReady(const TaggedNodeSeq& ready, TaggedNodeReadyQueue* inline_ready);        //僅用作調試或記錄    inline void MaybeMarkCompleted(FrameState* frame, int64 iter, int64 id);        //輸出一個未完成或者活躍節點的資訊    void DumpPendingNodeState(const int node_id, const Entry* input_vector, bool show_nodes_with_no_ready_inputs);    void DumpActiveNodeState(const FrameState* frame, IterationState* iteration);        //提供執行器的狀態資訊    void DumpState();    const Tensor* GetTensorValueForDump(const Entry& input);        //當執行器結束執行時，清理    void Finish();};

從API上來看，ExecutorState幾乎擔當了執行器的職責，從後面的介紹也可以看出，實際上確實如此。執行器內部實際調用的就是ExecutorState內部的API。從類的結構中，我們還是看到了許多未曾相識的結構，下面我們先一一分析這些類的意義和結構。

首先來看Entry，Entry要麼是一個張量指標，要麼是一個張量值，為計算圖中的節點的輸入或輸出提供了一種統一的類型。

struct Entry {    Entry(const Entry& other);    Entry& operator=(const Entry& other);        void ClearVal();//清除val欄位    ManualConstructor<Tensor> val;//一個張量的值，如果val_filed_is_set是true的話    Tensor* ref = nullptr;//一個張量引用    mutext* ref_mu = nullptr;//為上述張量引用的互斥量    bool has_value = false;//值是否存在，不論是val或者ref    bool val_filed_is_set = false;//val欄位是否被設定        AllocatorAttributes alloc_attr;//為當前的張量分配記憶體的記憶體 Clerk的屬性        DeviceContext* device_context = nullptr;//包含了關於這個張量如何建立的裝置相關的資訊};

接下來看看IterationState，它代表了一輪迭代的狀態。

struct IterationState {  public:    //一輪迭代的狀態，每個迭代輪次都由一個單獨的拷貝。對於第k輪迭代，第i個節點的第j個輸入在input_tensors[k][impl_->nodes[i].input_start+j]。注意，沒有必要對input_tensors做互斥鎖，其中的內容只會被邊的前一個節點寫入，被邊的後一個節點擦除，而每條邊的前後兩個節點是不可能同時啟動並執行    Entry* input_tensors;        //每一輪迭代中未完成的op數量    size_t outstanding_ops;        //每一輪迭代中未完成的幀數量    int outstanding_frame_count;    int pending(PendingCounts::Handle h);    int decrement_pending(PendingCounts::Handle int v);        //標記一個merge節點為live    void mark_live(PendingCounts::Handle h);    //標記一個節點為處理開始    void mark_started(PendingCounts::Handle h);    //標記一個節點為處理結束    void mark_completed(PendingCounts::Handle h);    //擷取節點狀態    PendingCounts::NodeState node_state(PendingCounts::Handle h);    int dead_count(PendingCounts::Handle h);    void increment_dead_count(PendingCounts::Handle h);    void adjust_for_activation(PendingCounts::Handle h, bool increment_dead, int* pending_result, int* dead_result);    private:    PendingCounts counts_;};

接下來是FrameState，代表了一個幀的狀態。對於幀和迭代輪次，有以下幾點需要說明：

• 對於計算圖中的迴圈來說，每個迴圈都需要建立一個新的幀。執行從第0個迭代開始。當第0個迭代的某個數值通過了一個NextIteration節點時，第1輪迭代就被建立並開始運行了。注意這時第0輪迭代可能仍在進行，所以多輪迭代可能會同時在運行。幀保持了多種資料結構來儲存每輪迭代的狀態。當第0輪迭代結束後，我們對其對應的狀態進行記憶體回收。
• 一個幀，當它的所有輸入都已經被傳入，所有的迭代都被計算完成時，這個幀就被認為是完成了，可以被進行記憶體回收了。
• 一個幀儲存了其中每一輪迭代的狀態。如果以下三個條件都被滿足，那麼第i輪迭代就會被認為是已經完成了，第一，第i輪迭代已經沒有未完成的節點了，第二，所有該輪的接收操作都已經完成了，第三，第i-1輪已完成。對於第0輪迭代，當幀的所有輸入都已完成，我們就認為它已經結束了。
• 幀和迭代輪次在結束後，都會進行記憶體回收。我們需要儲存的狀態量，跟調度器允許的並行度高度相關。我們希望調度器能夠動態控制未完成的並行幀和迭代的數量。為了減少記憶體消耗，調度器可能需要優先調度內層的幀和較低的迭代輪次。
• 幀的狀態一般總是在需要的時候才會被初始化，因此我們沒有引入額外的損耗。

下面我們來具體看下FrameState的結構：

struct FrameState {    const ExecutorImpl* executor = nullptr;//幀所在的執行器    string frame_name;//當前幀的名稱，是父幀，迭代輪次，和frame_name欄位拼合起來得到的    uint64 frame_id;//當前幀的唯一標識    int64 parent_iter = -1;//父幀的迭代輪次，frame_name和parent_iter共同標識了當前的FrameState    FrameState* parent_frame = nullptr;//父幀的FrameState    const int         max_parallel_iterations;//最大允許的並行迭代數量    int num_pending_inputs = 0;//當前幀仍然在等待的輸入數量    int64 iteration_count GUARDED_BY(mu) = 0;//當前幀中到達過的最大的迭代數量    int num_outstanding_iterations GUARDED_BY(mu) = 1;//未完成的迭代數量        gtl::InlinedVecotr<IterationState*,12> iterations;//當前幀活躍的迭代狀態    std::vector<std::pair<const Node*, Entry>> next_iter_roots GUARDED_BY(mu);    std::vector<std::pair<const Node*, Entry>> inv_values GUARDED_BY(mu);    std::vector<const Node*> dead_exits GUARDED_BY(mu);        //屬於當前幀的靜態資訊    PendingCounts* pending_counts = nullptr;    int total_input_tensors = 0;    std::vector<const Node*>* nodes = nullptr;        void InitializeFrameInfo(const string& enter_name);    inline IterationState* GetInteration(int64 iter);    inline void SetIteration(int64 iter, IterationState* state);        //減少未完成的運算元量，清理幀中的迭代資訊。如果幀執行結束則返回true    inline bool DecrementOutputstandingOps(const GraphView* gview, int64 iter, TaggedNodeSeq* ready);    inline bool DecrementOutstandingOpsLocked(const GraphView* gview, int64 iter, TaggedNodeSeq* ready);        //如果幀中的計算都已經完成則返回true    inline bool IsFrameDone();    //如果迭代的計算已經結束則返回true    bool IsIterationDone(int64 iter);    //增加迭代的編號，如果是一個新迭代，就初始化它    void IncrementIteration(const GraphView* gview, TaggedNodeSeq* ready);    //啟用一個新的迭代輪次中所有的NextIteration節點    void ActivateNexts(const GraphView* gview, int64 iter, TaggedNodeSeq* ready);    void ActivateLoopInvs(const GraphView* gview, int64 iter, TaggedNodeSeq* ready);    void AddLoopInv(const NodeItem* item, const Entry& value, TaggedNodeSeq* ready);    void ActivateNodes(const NodeItem* item, const bool is_dead, int64 iter, EntryVector* outputs, TaggedNodeSeq* ready);    bool CleanupIterations(const GraphView* gview, int64 iter, TaggedNodeSeq* ready);};

最後讓我們來看下最後的兩個結構體，TaggedNode和TaggedNodeReadyQueue。其中TaggedNode非常簡單，就是一個<frame, iter, node>的結構體，而後者就是前者的一個Queue，用來表示已經準備好的節點的隊列。

struct TaggedNode {    const Node* node = nullptr;    FrameState* input_frame = nullptr;    int64 input_iter = -1;    bool is_dead = false;        TaggedNode(const Node* t_node, FrameState* in_frame, int64 in_iter, bool dead);};class TaggedNodeReadyQueue {  public:    void push_back(TaggedNode node);    void pop_front();    bool empty();    const TaggedNode* begin();    const TaggedNode* end();  private:    gtl::InlinedVector<TaggedNode, 16> ready_;    int front_index_;};

關於TaggedNodeReadyQueue，我們要說明一下，本來這裡很自然的可以使用std::deque這個標準的雙端列表來實現的，但因為在待運行序列中我們通常並沒有太多的節點，所以為了效率我們只使用了一個vector來實現，並且省去了節點消耗後釋放空間的麻煩。

3.5 details

終於快要接近終點了。在前文中我們講了那麼多結構，最終計算圖的執行過程究竟是怎樣的，我們仍然不得而知。因為具體的實現細節都隱藏在函數的實現中，而我們上文中全部都在探討介面。現在我們就來看下，具體的實現方法。
首先，執行器的入口是Run函數，先來看下ExecutorImpl中的Run函數是如何?的吧。

void ExecutorImpl::RunAsync(const Args& args, DoneCallback done){    (new ExecutorState(args,this))->RunAsync(std::move(done));}

這驗證了我們上文中提到的，ExecutorImpl仍然只是一個介面，真正的執行是被推到ExecutorState類中完成的。在上述函數中，我們首先定義了一個ExecutorState對象，然後調用了它的RunAsync函數。在建構函式中，首先初始化了root_frame和iteration 0，我們具體看看RunAsync是如何?的：

void ExecutorState::RunAsync(Executor::DoneCallback done){    const Graph* graph = impl_->graph_;//擷取計算圖指標    TaggedNodeSeq ready;//構建ready節點序列        //讓裝置填充裝置上下文映射    Device* device = impl_->params_.device;    Status fill_status = device->FillContextMap(graph, &device_context_map_);    if(!fill_status.ok()){        done(fill_status);        return;    }        //初始化ready隊列    for(const Node* n : impl_->root_nodes){        DCHECK_EQ(n->in_edges().size(),0);        ready.push_back(TaggedNode{n,root_frame_,0,false});    }    if(ready.empty()){        done(Status::OK());    } else {        num_outstanding_ops = ready.size();        root_frame_->iterations[0]->outstanding_ops = ready.size();        done_cb_ = std::move(done);        ScheduleReady(ready,nullptr);    }}

可見，主要做了兩件事，第一是初始化了ready queue，第二是啟動了ScheduleReady函數。
下面我們再來看一下SheduleReady函數的運行機制：

void ExecutorState::ScheduleReady(const TaggedNodeSeq& ready, TaggedNodeReadyQueue* inline_ready){    if(ready.empty()) return ;        int64 scheduled_usec = 0;    if(stats_collector_){        scheduled_usec = nodestats::NowInUsec();    }    if(inline_ready == nullptr){        //線上程池中調度所有已經準備好的op        for(auto& tagged_node : ready){            runner_([=]() {Process(tagged_node, scheduled_usec);});        }        return;    }    const GraphView& gview = impl_->gview_;    const TaggedNode* curr_expensive_node = nullptr;    for(auto& tagged_node : ready){        const NodeItem& item = *gview.node(tagged_node.node->id());        if(tagged_node.is_dead || ! item.kernel_is_expensive){            //內聯化這個非複雜節點            inline_ready->push_back(tagged_node);        } else {            if(curr_expensive_node){                //將複雜節點丟給其它線程去做，因為當前線程還有很多事情要做                runner_(std::bind(&ExecutorState::Process, this, *curr_expensive_node, scheduled_usec));            }            curr_expensive_node = &tagged_node;        }    }    if(curr_expensive_node){        if(inline_ready->empty()){            //尾遞迴最佳化            inline_ready->push_back(*curr_expensive_node);        } else {            //我們仍然有內聯節點需要運行，因此把這個複雜節點丟給其它線程去運行            runner_(std::bind(&ExecutorState::Process, this, *curr_expensive_node, scheduled_usec));        }    }}

這個函數包含了兩個輸入，一個是待執行的節點隊列，一個是待執行的內聯節點序列。一共分兩種情況處理：

• 第一種情況，inline_ready隊列為空白，這種情況下，我們會為ready隊列中的每一個節點，單獨新增一個執行線程，這也是ExecutorState::RunAsync函數調用調度函數時的執行方式。也就是說，執行的起點是，對於根執行隊列中的節點，分別新增一個線程來執行；
• 第二種情況，inline_ready隊列非空，這種情況下，我們需要明確一點，調度函數不會進行任何實際的執行，只會對執行進行分配。它會遍曆ready中的每個節點，如果這個節點是非複雜節點或者節點已死亡，就放入inline_ready隊列待執行，否則就單獨開啟一個線程執行它，同時，這個遍曆過程進行完之後，會保留最後一個複雜節點（curr_expensive_node），這時候如果inline_ready隊列是空的，就把這個複雜節點放入內聯隊列，否則就開啟一個線程執行。

下面來看下Process函數，它是整個執行的核心，這個函數包含的代碼量比較大，因為我們的核心目標是說明執行過程，所以細枝末節暫時略去，僅保留架構：

void ExecutorState::Process(TaggedNode tagged_node, int64 scheduled_usec){    const GraphView& gview = impl_->gview_;    TaggedNodeSeq ready;    TaggedNodeReadyQueue inline_ready;        //為OpKernel::Compute準備參數        inline_ready.push_back(tagged_node);    while(!inline_ready.empty()){        //資訊提取、記錄處理                //當這個節點是非死亡節點，或者這個節點是send/recv這樣的資料轉送節點時，才執行這個節點，對於傳輸節點，我們需要把dead這個位元組傳輸下去        bool launched_asynchronously = false;        if(tagged_node.is_dead & !IsTransferNode(node)){            outputs.resize(item.num_outputs);        } else {            //調用PreparedInputs準備輸入            //設定計算參數            if(item.kernel_is_async){                AsyncOpKernel* async = item.kernel->AsAsync();                launched_asynchronously = true;                AsyncState* state = new AsyncState(params, tagged_node, &item, first_input, stats);                                auto done = [this, state](){                    //調用ProcessOutputs處理輸出                    //清理輸入                    //調用PropagateOutputs傳遞輸出                    //調用NodeDone清理戰場                };                device->ComputeAsync(async, &state->ctx, done);            } else {                device->Compute(op_kernel,&ctx);                //調用ProcessOutputs處理輸出            }//同步處理結束        }//非死亡、非傳輸節點處理結束        if(!launched_asynchronously){            //清理輸入            //調用PropagateOutputs傳遞輸出            //調用NodeDone打掃戰場        }//後續處理結束    }//while迴圈結束}

這個函數把節點計算分為同步和非同步計算分別處理，且同樣遵循以下的處理方式：

graph LRPrepareInput-->ComputeCompute-->ProcessOutputProcessOutput-->PropagateOutputPropagateOutput-->NodeDone

下面我們分別來看一下這四個函數的實現，首先是準備輸入函數和輸出處理函數

Status ExecutorState::PrepareInputs(const NodeItem& item, Entry* first_input, TensorValueVec* inputs, DeviceContextVec* input_device_contexts, AllocatorAttributeVec* input_alloc_attrs, bool* is_input_dead);Status ExecutorState::ProcessOutputs(const NodeItem& item, OpKernelContext* ctx, EntryVector* outputs, NodeExecStats* stats);

這兩個函數沒有什麼花花腸子，只是對輸入和輸出的屬性進行設定和填充，比較瑣碎，感興趣的讀者可以去看看源碼。
比較要緊的是推廣輸出和節點完成這兩個函數，我們先看一下推廣輸出的函數：

void ExecutorState::PropagateOutputs(const TaggedNode& tagged_node, const NodeItem* item, EntryVector* outputs, TaggedNodeSeq* ready){    //沿著輸出邊傳遞輸出，把新準備好的節點放入ready隊列        //判斷當前節點的類型，選擇合適的處理方法，期間會調用ActivateNodes，DecrementOutstandingOpsLocked，AddLoopInv，DecrementOutstandingOps，IncrementIteration等函數進行處理        //節點處理完成後，判斷當前幀是否執行完畢，以及遞迴的判斷父幀有沒有執行完畢}

其中，在ActivateNodes函數中，有這樣一個結構：

for(size_t out_index=0; out_index<num_output_edges;out_index++){    //其它處理    if(dst_ready){        if(dst_item->is_control_trigger)            dst_dead = false;        ready->push_back(TaggedNode(dst_item->node, this, iter, dst_dead));        iter_state->outstanding_ops++;    }}

也就是說，在啟用節點的時候，會順勢把該節點加入待執行隊列ready。這一點很重要，因為在Process函數中，我們是有一個while迴圈對inline_ready隊列（也就是啟用節點函數中的ready隊列）做處理的，但剛開始這個隊列中只有一個節點，正式因為在PropagateOutputs函數中會調用ActivateNodes函數，不斷向ready中添加節點，才使得這個while迴圈能夠跑起來。
最後再看一下，節點完成這個函數：

bool ExecutorState::NodeDone(const Status& s, const Node* node, const TaggedNodeSeq& ready, NodeExecStats* stats, TaggedNodeReadyQueue* inline_ready){    //其它處理    if(s.ok()){        ScheduleReady(ready, inline_ready);    }    return completed;}

感覺又回到開頭了？還記得剛開始講到過的，ExecutorState::RunAsync函數中就調用了ScheduleReady這個函數嗎？這個函數究竟起到了什麼樣的作用呢？
現在我們有必要把上述的過程總結一下了，這個過程也是本篇最核心的內容了：

• ExecutorImpl::RunAsync作為執行器的入口，其實它是把實際執行的工作交給了ExecutorState::RunAsync，這個函數進一步調用了ScheduleReady函數來調度執行，記住這個函數有兩個輸入，ready和inline_ready，在這裡，inline_ready是空的。也就是說，這裡調用ScheduleReady的作用是，給根節點隊列裡的節點，分別分配一個線程執行，執行的程序呼叫的是Process函數。
• 在Process函數內部，我們還要記住一點，這個函數的輸入只有節點，沒有ready和inline_ready，這兩個變數都是在Process函數內部新建立的。也就是說，一旦把一個節點交給Process函數去處理，這個節點所在的隊列跟Process函數就沒有任何關係了。處理的過程分為輸入準備、實際計算、輸出準備、輸出傳遞、節點完成五個步驟。其中只有輸出傳遞和節點完成會對ready和inline_ready結構產生影響。
• 我們把節點分按照非同步節點和同步節點分開處理。對於非同步節點，NodeDone函數的最後一個參數inline_ready是空，也就是說，在非同步執行時，調用NodeDone中的ScheduleReady時，跟RunAsync中的情形是一樣的，直接調度ready中的節點就好了，不需要處理inline_ready的情況。對於同步節點，NodeDone函數的最後一個參數inline_ready是當前Process函數中新建立的inline_ready，也就是說，傳遞給ScheduleReady的inline_ready是非空的，這也就有可能對inline_ready的結構做修改，注意這裡的inline_ready是從Process函數中建立的，每個Process函數都對應一個全新的線程，也就是說，每個全新的線程裡面只有一個inline_ready結構，其中的函數不斷的修改它的內容，然後不斷的對它進行調度執行。注意Process中的while大迴圈是針對inline_ready隊列執行的。