MXNet: wait_to_read 方法

wait_to_read

在mxnet中,類ndarray可以調(diào)用 wait_to_read,官方給出的該函數(shù)解釋是:

Waits until all previous write operations on the current array are finished.
This method guarantees that all previous write operations that pushed
into the backend engine for execution are actually finished.
Examples

    >>> import time
    >>> tic = time.time()
    >>> a = mx.nd.ones((1000,1000))
    >>> b = mx.nd.dot(a, a)
    >>> print(time.time() - tic) # doctest: +SKIP
    0.003854036331176758
    >>> b.wait_to_read()
    >>> print(time.time() - tic) # doctest: +SKIP
    0.0893700122833252

翻譯過來就是,調(diào)用該方法可以保證,之前進行的該ndarray的所有寫操作都完成了。由于MXNet是一個異步框架,我們使用python接口調(diào)用方法的時候,只是將該方法push給了執(zhí)行者,由執(zhí)行者來安排什么時候執(zhí)行。所以,我們就無法知道,某一個操作是不是完成了。

那么,在知道了該方法的功能后,我們的疑問就變成了,mxnet是怎么做到的呢?

在mxnet的類ndarray中定義了WaitToRead方法,如下:

inline void WaitToRead() const {
    if (is_none()) return;
    Engine::Get()->WaitForVar(ptr_->var);
}

其中ptr_->var指向了該ndarray實例對應的唯一的var,engine會根據(jù)該var來進行判斷,當前操作對應的是哪個ndarray。

那么,我們只需要知道Engine對應的WariForVar做了什么就好了。在mxnet中實現(xiàn)了很多種engine。但是我們用的engine都繼承了類ThreadedEngineWaitForVar方法。沒有重載。
在這個方法里面,主要做了一件事情,就是將一個操作push給了engine。
該操作需要完成的函數(shù)是:

this, &done](RunContext, CallbackOnComplete on_complete) {
      if (engine_info_) {
        LOG(INFO) << "Sync is executed";
      }
      {
        std::unique_lock<std::mutex> lock{finished_m_};
        done.store(true);
      }
      finished_cv_.notify_all();
      if (engine_info_) {
        LOG(INFO) << "Sync is notified";
      }
      on_complete();
    }

輸出是var。
函數(shù)將done賦值為true,它被初始化為false。之后會調(diào)用finished_cv_.notify_all();來通知該方法繼續(xù)運行。
因為,該方法會在將操作push之后,調(diào)用

{
    std::unique_lock<std::mutex> lock{finished_m_};
    finished_cv_.wait(lock, [this, &done]() {
        return done.load() || kill_.load();
    });
  }

導致一直等待,直到done為true。

到這里,我們大概明白了ndarray的wait_to_read的方法如何實現(xiàn)等待。
總體思路就是,調(diào)用了engine的WaitForVar方法。這個方法會將一個操作異步push給engine,這個操作將一個原子變量done從false改變?yōu)閠rue。由于push是異步的,它會立刻返回,返回后,WaitForVar一直等待,知道done為true。
而engine會根據(jù)變量來安排執(zhí)行,由于這次push的操作的輸出是ndarray對應的var,因此,engine會保證所有之前的寫入操作完成。

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點,簡書系信息發(fā)布平臺,僅提供信息存儲服務。

相關閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容