MXNet: wait_to_read 方法

weixin_34019929發表於2019-02-07

wait_to_read

在mxnet中,類ndarray可以呼叫 wait_to_read,官方給出的該函式解釋是:

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

翻譯過來就是,呼叫該方法可以保證,之前進行的該ndarray的所有寫操作都完成了。由於MXNet是一個非同步框架,我們使用python介面呼叫方法的時候,只是將該方法push給了執行者,由執行者來安排什麼時候執行。所以,我們就無法知道,某一個操作是不是完成了。

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

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

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

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

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

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
函式將done賦值為true,它被初始化為false。之後會呼叫finished_cv_.notify_all();來通知該方法繼續執行。
因為,該方法會在將操作push之後,呼叫

{
    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的方法如何實現等待。
總體思路就是,呼叫了engine的WaitForVar方法。這個方法會將一個操作非同步push給engine,這個操作將一個原子變數done從false改變為true。由於push是非同步的,它會立刻返回,返回後,WaitForVar一直等待,知道done為true。
而engine會根據變數來安排執行,由於這次push的操作的輸出是ndarray對應的var,因此,engine會保證所有之前的寫入操作完成。

相關文章