Do not touch GPU 0 during ReleaseAll

src/storage/pooled_storage_manager.h

@@ -53,11 +53,14 @@ class GPUPooledStorageManager final : public StorageManager {
  public:
   /*!
    * \brief Default constructor.
+   *
+   * \param initial_context context used by this Storage Manager
    */
-  GPUPooledStorageManager() {
+  explicit GPUPooledStorageManager(Context initial_context) {
     reserve_ = dmlc::GetEnv("MXNET_GPU_MEM_POOL_RESERVE", 5);
     page_size_ = dmlc::GetEnv("MXNET_GPU_MEM_POOL_PAGE_SIZE", 4096);
     large_alloc_round_size_ = dmlc::GetEnv("MXNET_GPU_MEM_LARGE_ALLOC_ROUND_SIZE", 2 * 1024 * 1024);
+    initial_context_ = initial_context;
     if (large_alloc_round_size_ <= 0) {
       LOG(FATAL) << "MXNET_GPU_MEM_LARGE_ALLOC_ROUND_SIZE cannot be set to a value <= 0, found: "
                  << large_alloc_round_size_;
@@ -123,6 +126,8 @@ class GPUPooledStorageManager final : public StorageManager {
   int reserve_;
   // number of devices
   const size_t NDEV = 32;
+  // context used by this Storage Manager
+  Context initial_context_;
   // memory pool
   std::unordered_map<size_t, std::vector<void*>> memory_pool_;
   DISALLOW_COPY_AND_ASSIGN(GPUPooledStorageManager);
@@ -177,6 +182,7 @@ void GPUPooledStorageManager::ReleaseAll() {
       Storage::Handle handle;
       handle.dptr = j;
       handle.size = i.first;
+      handle.ctx = initial_context_;
       DirectFreeNoLock(handle);
     }
   }
@@ -201,11 +207,14 @@ class GPUPooledRoundedStorageManager final : public StorageManager {
  public:
   /*!
    * \brief Default constructor.
+   *
+   * \param initial_context context used by this Storage Manager
    */
-  GPUPooledRoundedStorageManager() {
+  explicit GPUPooledRoundedStorageManager(Context initial_context) {
     reserve_ = dmlc::GetEnv("MXNET_GPU_MEM_POOL_RESERVE", 5);
     page_size_ = dmlc::GetEnv("MXNET_GPU_MEM_POOL_PAGE_SIZE", 4096);
     cut_off_ = dmlc::GetEnv("MXNET_GPU_MEM_POOL_ROUND_LINEAR_CUTOFF", 24);
+    initial_context_ = initial_context;
     if (page_size_ < 32) {
       LOG(FATAL) << "MXNET_GPU_MEM_POOL_PAGE_SIZE cannot be set to a value smaller than 32. " \
                  << "Got: " << page_size_ << ".";
@@ -290,6 +299,8 @@ class GPUPooledRoundedStorageManager final : public StorageManager {
   size_t cut_off_;
   // percentage of reserved memory
   int reserve_;
+  // context used by this Storage Manager
+  Context initial_context_;
   // memory pool
   std::vector<std::vector<void*>> memory_pool_;
   DISALLOW_COPY_AND_ASSIGN(GPUPooledRoundedStorageManager);
@@ -345,6 +356,7 @@ void GPUPooledRoundedStorageManager::ReleaseAll() {
       Storage::Handle handle;
       handle.size = size;
       handle.dptr = j;
+      handle.ctx = initial_context_;
       DirectFreeNoLock(handle);
     }
     memory_pool_[i].clear();

src/storage/storage.cc

@@ -104,13 +104,13 @@ void StorageImpl::Alloc(Storage::Handle* handle) {
             std::string strategy = type;
 
             if (strategy == "Round") {
-              ptr = new storage::GPUPooledRoundedStorageManager();
+              ptr = new storage::GPUPooledRoundedStorageManager(handle->ctx);
               LOG(INFO) << "Using GPUPooledRoundedStorageManager.";
             } else {
               if (strategy != "Naive") {
                 LOG(FATAL) << "Unknown memory pool strategy specified: " << strategy << ".";
               }
-              ptr = new storage::GPUPooledStorageManager();
+              ptr = new storage::GPUPooledStorageManager(handle->ctx);
             }
 #else
             LOG(FATAL) << "Compile with USE_CUDA=1 to enable GPU usage";