[simple-engine] implement missing functions

src/dag_engine/simple_engine.cc

@@ -6,6 +6,9 @@
 #include <cassert>
 #include <algorithm>
 #include <utility>
+#include <condition_variable>
+#include <mutex>
+#include "../common/cuda_utils.h"
 
 namespace mxnet {
 
@@ -15,7 +18,8 @@ SimpleVar* SimpleVar::CastFromBase(Var* v) { return v->Cast<SimpleVar>(); }
 
 SimpleOpr* SimpleOpr::CastFromBase(Opr* o) { return o->Cast<SimpleOpr>(); }
 
-SimpleEngine::SimpleEngine() : thread_pool_{[this]() { ThreadWorker(); }} {}
+SimpleEngine::SimpleEngine()
+    : pending_{0}, thread_pool_{[this]() { ThreadWorker(); }} {}
 
 SimpleEngine::~SimpleEngine() noexcept(false) { task_queue_.SignalForKill(); }
 
@@ -42,9 +46,9 @@ SimpleEngine::Operator SimpleEngine::NewOperator(
 void SimpleEngine::DeleteOperator(Operator op) { delete op; }
 
 void SimpleEngine::Push(Operator op, Context exec_ctx) {
-  static_cast<void>(exec_ctx);
   auto opr = SimpleOpr::CastFromBase(op);
   auto opr_block = new OprBlock{};
+  ++pending_;
   opr_block->wait.store(opr->use_vars.size() + opr->mutate_vars.size() + 1);
   // Add reading dependencies.
   auto add_dependency = [&opr_block](SimpleVar* i) {
@@ -97,15 +101,60 @@ void SimpleEngine::Push(Operator op, Context exec_ctx) {
     previous->lock.unlock();
   }
   auto callback = [this, first]() { OnComplete(first); };
-  // TODO(hotpxl) do something useful
-  RunContext ctx{};
-  ctx.stream = nullptr;
-  opr_block->fn = [opr, ctx, callback]() { opr->fn(ctx, callback); };
+  RunContext rctx{};
+  rctx.stream = nullptr;
+  opr_block->fn = [exec_ctx, opr, rctx, callback]() {
+    if (exec_ctx.dev_mask == gpu::kDevMask) {
+#if MXNET_USE_CUDA
+      CUDA_CALL(cudaSetDevice(exec_ctx.dev_id));
+#else  // MXNET_USE_CUDA
+      LOG(FATAL) << "Please compile with CUDA enabled";
+#endif  // MXNET_USE_CUDA
+    }
+    opr->fn(rctx, callback);
+  };
   if (--opr_block->wait == 0) {
     task_queue_.Push(opr_block);
   }
 }
 
+void SimpleEngine::PushAsync(AsyncFn fn, Context exec_ctx,
+                             std::vector<Variable> const& use_vars,
+                             std::vector<Variable> const& mutate_vars) {
+  auto&& opr = NewOperator(fn, use_vars, mutate_vars);
+  Push(opr, exec_ctx);
+  DeleteOperator(opr);
+}
+
+void SimpleEngine::PushDelete(Fn delete_fn, Context exec_ctx, Variable var) {
+  auto&& callback = [delete_fn, var](RunContext ctx) {
+    delete var;
+    // If you used `var` after `PushDelete`, then the following will be
+    // undefined
+    delete SimpleVar::CastFromBase(var)->var;
+    delete_fn(ctx);
+  };
+  Push(callback, exec_ctx, {}, {var});
+}
+
+void SimpleEngine::WaitForVar(Variable var) {
+  std::condition_variable cv;
+  std::mutex m;
+  std::unique_lock<std::mutex> lock{m};
+  std::atomic<bool> done{false};
+  auto&& callback = [&cv, &done](RunContext) {
+    done.store(true);
+    cv.notify_all();
+  };
+  cv.wait(lock, [&done]() { return done.load(); });
+  Push(callback, Context{}, {var}, {});
+}
+
+void SimpleEngine::WaitForAll() {
+  std::unique_lock<std::mutex> lock{finished_m_};
+  finished_cv_.wait(lock, [this]() { return pending_.load() == 0; });
+}
+
 void SimpleEngine::OnComplete(VersionedVarBlock* trigger) {
   auto head = trigger;
   while (head != nullptr) {
@@ -145,6 +194,9 @@ void SimpleEngine::ThreadWorker() {
     assert(opr_block->wait.load() == 0);
     opr_block->fn();
     delete opr_block;
+    if (--pending_ == 0) {
+      finished_cv_.notify_all();
+    }
   }
 }
 

src/dag_engine/simple_engine.h

@@ -9,6 +9,8 @@
 #include <vector>
 #include <functional>
 #include <atomic>
+#include <condition_variable>
+#include <mutex>
 #include "mxnet/dag_engine.h"
 #include "dag_engine_impl.h"
 #include "thread_pool.h"
@@ -76,15 +78,17 @@ class SimpleEngine final : public DAGEngine {
    * \brief Overriding methods.
    */
   Variable NewVar() override;
-  Operator NewOperator(AsyncFn, std::vector<Variable> const&,
-                       std::vector<Variable> const&) override;
+  Operator NewOperator(AsyncFn fn, std::vector<Variable> const& use_vars,
+                       std::vector<Variable> const& mutate_vars) override;
   void DeleteOperator(Operator op) override;
-  void Push(Operator op, Context) override;
-  void PushAsync(AsyncFn, Context, std::vector<Variable> const&,
-                 std::vector<Variable> const&) override{};
-  void PushDelete(Fn, Context, Variable) override{};
-  void WaitForVar(Variable) override{};
-  void WaitForAll() override{};
+  void Push(Operator op, Context exec_ctx) override;
+  using DAGEngine::Push;
+  void PushAsync(AsyncFn exec_fun, Context exec_ctx,
+                 std::vector<Variable> const& use_vars,
+                 std::vector<Variable> const& mutate_vars) override;
+  void PushDelete(Fn delete_fn, Context exec_ctx, Variable var) override;
+  void WaitForVar(Variable var) override;
+  void WaitForAll() override;
   /*!
    * \brief Callback on operation completion.
    *
@@ -103,6 +107,12 @@ class SimpleEngine final : public DAGEngine {
    * \brief Concurrency for thread pool.
    */
   static constexpr std::size_t kNumWorkingThreads = 16;
+  /*!
+   * \brief Number of pending operations.
+   */
+  std::atomic<std::size_t> pending_;
+  std::condition_variable finished_cv_;
+  std::mutex finished_m_;
   /*!
    * \brief Task queue.
    */

test/test_simple_engine.cc

@@ -30,7 +30,8 @@ int main() {
         {var}, {}));
     engine->Push(oprs.at(i), mxnet::Context{});
   }
-  std::this_thread::sleep_for(std::chrono::seconds{1});
+  engine->WaitForAll();
+  // std::this_thread::sleep_for(std::chrono::seconds{1});
 
   printf("============= Test #2 ==============\n");
   var = engine->NewVar();
@@ -44,19 +45,8 @@ int main() {
         {}, {var}));
     engine->Push(oprs.at(i), mxnet::Context{});
   }
-  std::this_thread::sleep_for(std::chrono::seconds{1});
+  // std::this_thread::sleep_for(std::chrono::seconds{1});
+  engine->WaitForAll();
 
-  // usleep(1000000);
-
-  // // Test #2
-  // cout << "============= Test #2 ==============" << endl;
-  // for (int i = 0; i < 10; ++i) {
-  //   engine->Push([i] (RunContext rctx) { Foo(rctx, i); },
-  //       exec_ctx, {}, vars);
-  // }
-
-  // usleep(1000000);
-
-  // // Test #3
   return 0;
 }