Add Per Device Threaded Engine Policy, Explicit use copy

.travis.yml

@@ -11,6 +11,7 @@ env:
  - TASK=python CXX=g++
  - TASK=python3 CXX=g++
  - TASK=python_naive CXX=g++
+ - TASK=python_perdev CXX=g++
  - TASK=cpp_unittest CXX=g++
 
 # dependent apt packages

include/mxnet/engine.h

@@ -33,8 +33,10 @@ typedef Opr* OprHandle;
 enum class FnProperty {
  /*! \brief Normal operation */
  kNormal,
- /*! \brief Copy operation between CPU and GPU */
- kCopy,
+ /*! \brief Copy operation from GPU to other devices */
+ kCopyFromGPU,
+ /*! \brief Copy operation from CPU to other devices */
+ kCopyToGPU,
  /*! \brief Asynchronous function call */
  kAsync
 }; // enum class FnProperty

scripts/travis_script.sh

@@ -40,15 +40,23 @@ if [ ${TASK} == "python3" ]; then
  make all || exit -1
  export MXNET_ENGINE_TYPE=ThreadedEngine
  nosetests tests/python/unittest || exit -1
- nosetests tests/python/train || exit -1 
+ nosetests tests/python/train || exit -1
 fi
 
 if [ ${TASK} == "python_naive" ]; then
  echo "USE_CUDA=0" >> config.mk
  make all || exit -1
  export MXNET_ENGINE_TYPE=NaiveEngine
  nosetests tests/python/unittest || exit -1
- nosetests tests/python/train || exit -1 
+ nosetests tests/python/train || exit -1
+fi
+
+if [ ${TASK} == "python_perdev" ]; then
+ echo "USE_CUDA=0" >> config.mk
+ make all || exit -1
+ export MXNET_ENGINE_TYPE=ThreadedEnginePerDevice
+ nosetests tests/python/unittest || exit -1
+ nosetests tests/python/train || exit -1
 fi
 
 if [ ${TASK} == "cpp_unittest" ]; then

src/engine/thread_pool.h

@@ -6,7 +6,7 @@
 
 #include <dmlc/base.h>
 #include <cstddef>
-#include <array>
+#include <vector>
 #include <thread>
 #include <utility>
 #include "mxnet/base.h"
@@ -24,11 +24,17 @@ class ThreadPool {
  * \param size size of the thread pool.
  * \param func the function to run on the thread pool.
  */
- explicit ThreadPool(size_t size, std::function<void()> func);
- /*!
- * \brief Destructor.
- */
- ~ThreadPool() noexcept(false);
+ explicit ThreadPool(size_t size, std::function<void()> func)
+ : worker_threads_(size) {
+ for (auto& i : worker_threads_) {
+ i = std::thread(func);
+ }
+ }
+ ~ThreadPool() noexcept(false) {
+ for (auto&& i : worker_threads_) {
+ i.join();
+ }
+ }
 
  private:
  /*!
@@ -44,20 +50,6 @@ class ThreadPool {
  */
  DISALLOW_COPY_AND_ASSIGN(ThreadPool);
 };
-
-ThreadPool::ThreadPool(size_t size, std::function<void()> func)
- : worker_threads_(size) {
- for (auto& i : worker_threads_) {
- i = std::thread(func);
- }
-}
-
-ThreadPool::~ThreadPool() noexcept(false) {
- for (auto&& i : worker_threads_) {
- i.join();
- }
-}
-
 } // namespace engine
 } // namespace mxnet
 #endif // MXNET_ENGINE_THREAD_POOL_H_

src/engine/threaded_engine_perdevice.cc

@@ -7,6 +7,7 @@
 #include <dmlc/logging.h>
 #include <dmlc/parameter.h>
 #include <dmlc/concurrency.h>
+#include <array>
 #include "./threaded_engine.h"
 #include "./thread_pool.h"
 #include "./stream_manager.h"
@@ -92,7 +93,8 @@ class ThreadedEnginePerDevice : public ThreadedEngine {
  */
  inline ThreadWorkerBlock *GetGPUWorkerBlock(size_t dev_id,
  FnProperty prop) {
- bool is_copy = (prop == FnProperty::kCopy);
+ bool is_copy = (prop == FnProperty::kCopyFromGPU ||
+ prop == FnProperty::kCopyToGPU);
  CHECK_LT(dev_id, kMaxNumGPUs)
  << "GPU Device index " << dev_id
  << " exceed bound " << kMaxNumGPUs;
@@ -130,7 +132,7 @@ class ThreadedEnginePerDevice : public ThreadedEngine {
  dmlc::ConcurrentBlockingQueue<OprBlock*>* task_queue) {
  #if MXNET_USE_CUDA
  // allocate stream
- mshadow::SetDevice(dev_id);
+ mshadow::SetDevice<gpu>(dev_id);
  RunContext run_ctx;
  mshadow::Stream<gpu> *stream;
  if (is_copy_worker) {

src/engine/threaded_engine_pooled.cc

@@ -86,7 +86,9 @@ class ThreadedEnginePooled : public ThreadedEngine {
  LOG(FATAL) << "Please compile with CUDA enabled";
  #endif // MXNET_USE_CUDA
  }
- auto&& rctx = opr_block->opr->prop == FnProperty::kCopy
+ bool is_copy = (opr_block->opr->prop == FnProperty::kCopyFromGPU ||
+ opr_block->opr->prop == FnProperty::kCopyToGPU);
+ auto&& rctx = is_copy
  ? streams_.GetIORunContext(opr_block->ctx)
  : streams_.GetRunContext(opr_block->ctx);
  this->ExecuteOprBlock(rctx, opr_block);
@@ -97,7 +99,8 @@ class ThreadedEnginePooled : public ThreadedEngine {
  */
  void DoPushToQueue(OprBlock* opr_block) {
  switch (opr_block->opr->prop) {
- case FnProperty::kCopy: {
+ case FnProperty::kCopyFromGPU:
+ case FnProperty::kCopyToGPU: {
  io_task_queue_.Push(opr_block);
  break;
  }

src/ndarray/ndarray.cc

@@ -169,7 +169,7 @@ void CopyFromTo(const NDArray &from, NDArray *to) {
  ret.ptr_->CheckAndAlloc();
  TBlob tmp = ret.data();
  ndarray::Copy<cpu, cpu>(from.data(), &tmp,
- from.ctx(), ret.ctx(), ctx);
+  from.ctx(), ret.ctx(), ctx);
  }, from.ctx(), const_vars, {ret.ptr_->var});
  } else {
 #if MXNET_USE_CUDA
@@ -178,28 +178,28 @@ void CopyFromTo(const NDArray &from, NDArray *to) {
  ret.ptr_->CheckAndAlloc();
  TBlob tmp = ret.data();
  ndarray::Copy<cpu, gpu>(from.data(), &tmp,
- from.ctx(), ret.ctx(), ctx);
+  from.ctx(), ret.ctx(), ctx);
  // Wait GPU kernel to complete
  ctx.get_stream<gpu>()->Wait();
- }, ret.ctx(), const_vars, {ret.ptr_->var});
+ }, ret.ctx(), const_vars, {ret.ptr_->var}, FnProperty::kCopyToGPU);
  } else if (a == gpu::kDevMask && b == cpu::kDevMask) {
  Engine::Get()->PushSync([from, ret](RunContext ctx) {
  ret.ptr_->CheckAndAlloc();
  TBlob tmp = ret.data();
  ndarray::Copy<gpu, cpu>(from.data(), &tmp,
- from.ctx(), ret.ctx(), ctx);
+  from.ctx(), ret.ctx(), ctx);
  // Wait GPU kernel to complete
  ctx.get_stream<gpu>()->Wait();
- }, from.ctx(), const_vars, {ret.ptr_->var});
+ }, from.ctx(), const_vars, {ret.ptr_->var}, FnProperty::kCopyFromGPU);
  } else if (a == gpu::kDevMask && b == gpu::kDevMask) {
  Engine::Get()->PushSync([from, ret](RunContext ctx) {
  ret.ptr_->CheckAndAlloc();
  TBlob tmp = ret.data();
  ndarray::Copy<gpu, gpu>(from.data(), &tmp,
- from.ctx(), ret.ctx(), ctx);
+  from.ctx(), ret.ctx(), ctx);
  // Wait GPU kernel to complete
  ctx.get_stream<gpu>()->Wait();
- }, from.ctx(), const_vars, {ret.ptr_->var});
+ }, from.ctx(), const_vars, {ret.ptr_->var}, FnProperty::kCopyFromGPU);
  } else {
  LOG(FATAL) << "unknown device mask";
  }

src/symbol/graph_executor.cc

@@ -205,9 +205,6 @@ GraphExecutor::GetOpExecEntry(uint32_t nid) {
  }
  }
 
- for (const Resource& r : op_node.op_ctx.requested) {
- exec.mutate_vars.push_back(static_cast<DAGEngine::Variable>(r.var));
- }
  // start setup exec function.
  for (const Resource& r : op_node.op_ctx.requested) {
  exec.mutate_vars.push_back(static_cast<Engine::VarHandle>(r.var));