src/storage/storage.cc - master - githubmirror/incubator-mxnet

src/storage/storage.cc
274 строки9 KB
Vladimir Cherepanov
Port convolutions to cuDNN v8 API (#20635)
15 ноя 2021, 22:12
15 ноя 2021, 22:1236ed5e0
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include <mxnet/storage.h>
#include "./storage_manager.h"
#include "./naive_storage_manager.h"
#include "./pooled_storage_manager.h"
#include "./cpu_shared_storage_manager.h"
#include "./cpu_device_storage.h"
#include "./gpu_device_storage.h"
#include "./pinned_memory_storage.h"
#include "../common/lazy_alloc_array.h"
#include "../profiler/storage_profiler.h"

namespace mxnet {
namespace storage {

// consider change storage as a pure abstract class
class StorageImpl : public Storage {
 public:
  void Alloc(Handle* handle, bool failsafe) override;
  void Free(Handle handle) override;
  void DirectFree(Handle handle) override;
  void ReleaseAll(Context ctx) override {
    storage_manager(ctx)->ReleaseAll();
  }

  void SharedIncrementRefCount(Handle handle) override;
  StorageImpl()           = default;
  ~StorageImpl() override = default;

 private:
  std::shared_ptr<StorageManager> storage_manager(const Context& ctx) {
    auto&& device                           = storage_managers_.at(ctx.dev_type);
    std::shared_ptr<StorageManager> manager = device.Get(ctx.real_dev_id(), []() {
      LOG(FATAL) << "Cannot Free space to a device you have not allocated";
      return nullptr;
    });
    return manager;
  }

  static constexpr size_t kMaxNumberOfDevices = Context::kMaxDevType + 1;
  // internal storage managers
  std::array<common::LazyAllocArray<StorageManager>, kMaxNumberOfDevices> storage_managers_;
  profiler::DeviceStorageProfiler profiler_;
};  // struct Storage::Impl

StorageManager* CreateStorageManager(const Context& ctx,
                                     const char* context,
                                     int num_gpu_device,
                                     std::string* pStrategy) {
  const auto env_var = env_var_name(context, pool_type);
  const char* type   = getenv(env_var.c_str());
  if (type == nullptr) {
    type = "Naive";  // default pool
  }

  *pStrategy          = type;
  StorageManager* ptr = nullptr;
  if (*pStrategy == "Round") {
    ptr = new PooledStorageManager<RoundPower2, VectorContainer>(ctx, num_gpu_device);
  } else if (*pStrategy == "Naive") {
    ptr = new PooledStorageManager<RoundMultiple, UnorderedMapContainer>(ctx, num_gpu_device);
  } else if (*pStrategy == "Unpooled") {
    if (ctx.dev_type == Context::kCPU || num_gpu_device == 0)
      ptr = new NaiveStorageManager<CPUDeviceStorage>();
#if MXNET_USE_CUDA
    else if (ctx.dev_type == Context::kGPU)
      ptr = new NaiveStorageManager<GPUDeviceStorage>();
    else  // Context::kCPUPinned
      ptr = new NaiveStorageManager<PinnedMemoryStorage>();
#endif
  }
  return ptr;
}

void StorageImpl::Alloc(Storage::Handle* handle, bool failsafe) {
  // Set dptr to nullptr when handle size is 0.
  if (handle->size == 0) {
    handle->dptr = nullptr;
    return;
  }

  // space already recycled, ignore request
  auto&& device                           = storage_managers_.at(handle->ctx.dev_type);
  std::shared_ptr<StorageManager> manager = device.Get(handle->ctx.real_dev_id(), [handle]() {
    const auto dev_type = handle->ctx.dev_type;
    int num_gpu_device  = 0;
#if MXNET_USE_CUDA
    switch (dev_type) {
      case Context::kGPU:
      case Context::kCPUPinned:
        if (cudaGetDeviceCount(&num_gpu_device) != cudaSuccess)
          num_gpu_device = 0;
      default:
        break;
    }
#endif

    const char* context = nullptr;
    switch (dev_type) {
      case Context::kCPU:
        context = "CPU";
        break;
      case Context::kGPU:
#if MXNET_USE_CUDA
        context = "GPU";
        CHECK_GT(num_gpu_device, 0) << "GPU usage requires at least 1 GPU";
#else
        LOG(FATAL) << "Compile with USE_CUDA=1 to enable GPU usage";
#endif
        break;
      case Context::kCPUPinned:
        context = "CPU_PINNED";
        break;
      case Context::kCPUShared:
        // We will not generate the log messages for CPUShared
        // It could be as many of them as the number of "workers".
#if !defined(ANDROID) && !defined(__ANDROID__)
        break;  // For Android shared memory is not implemented
#endif
      default:
        LOG(FATAL) << "Unimplemented device " << dev_type;
    }

    // By default, the Pooled Storage Manager will be used, if it is available
    int naive_storage_manager = dmlc::GetEnv("MXNET_USE_NAIVE_STORAGE_MANAGER", 0);
    if (!naive_storage_manager) {
      // Because, the pooled storage managers are NOT implemented yet for
      // following dev_type's, we will also use the naive storage managers
      switch (dev_type) {
#if MXNET_USE_CUDA
        case Context::kCPUPinned:
          if (num_gpu_device > 0)
            break;
#endif
        case Context::kCPUShared:
          naive_storage_manager = true;
        default:
          break;
      }
    }

    StorageManager* ptr = nullptr;
    std::string strategy, storage_manager_type;
    if (naive_storage_manager) {
      storage_manager_type = "Naive";
      switch (dev_type) {
#if MXNET_USE_CUDA
        case Context::kGPU:
          ptr = new NaiveStorageManager<GPUDeviceStorage>();
          break;
        case Context::kCPUPinned:
          if (num_gpu_device > 0) {
            ptr = new NaiveStorageManager<PinnedMemoryStorage>();
            break;
          }
#else
        case Context::kCPUPinned:
#endif
        case Context::kCPU:
          ptr = new NaiveStorageManager<CPUDeviceStorage>();
          break;
#if !defined(ANDROID) && !defined(__ANDROID__)
        case Context::kCPUShared:
          ptr = new CPUSharedStorageManager();
#endif
        default:
          break;
      }
    } else {
      // Some Pooled Storage Manager will be used
      std::string strategy;
      ptr = CreateStorageManager(handle->ctx, context, num_gpu_device, &strategy);
      if (ptr) {
        if (strategy != "Unpooled")
          storage_manager_type = "Pooled (" + strategy + ")";
        else
          storage_manager_type = "Unpooled";
      } else {
        LOG(FATAL) << "Unknown memory pool strategy specified: " << strategy << ".";
      }
    }

    if (context)
      LOG(INFO) << "Using " << storage_manager_type << " StorageManager for " << context;

    return ptr;
  });

  manager->Alloc(handle, failsafe);
  if (!failsafe || handle->dptr != nullptr)
    profiler_.OnAlloc(*handle);
}

void StorageImpl::Free(Storage::Handle handle) {
  // Do nothing if dtpr is nullptr because the handle may have already
  // been freed or have not been allocated memory yet.
  if (handle.dptr == nullptr)
    return;

  storage_manager(handle.ctx)->Free(handle);
  profiler_.OnFree(handle);
}

void StorageImpl::DirectFree(Storage::Handle handle) {
  // Do nothing if dtpr is nullptr because the handle may have already
  // been freed or have not been allocated memory yet.
  if (handle.dptr == nullptr)
    return;

  storage_manager(handle.ctx)->DirectFree(handle);
  profiler_.OnFree(handle);
}

void StorageImpl::SharedIncrementRefCount(Storage::Handle handle) {
  CHECK_EQ(handle.ctx.dev_type, Context::kCPUShared);
  auto&& device = storage_managers_.at(Context::kCPUShared);
  auto manager  = device.Get(0, []() {
    LOG(FATAL) << "Cannot increment ref count before allocating any shared memory.";
    return nullptr;
  });
#if !defined(ANDROID) && !defined(__ANDROID__)
  dynamic_cast<CPUSharedStorageManager*>(manager.get())->IncrementRefCount(handle);
#else
  LOG(FATAL) << "Shared memory not implemented on Android";
#endif  // !defined(ANDROID) && !defined(__ANDROID__)
}

const std::string env_var_name(const char* dev_type, env_var_type type) {
  static const std::array<std::string, 5> name = {
      "MEM_POOL_TYPE",
      "POOL_PAGE_SIZE",
      "MEM_LARGE_ALLOC_ROUND_SIZE",
      "MEM_POOL_ROUND_LINEAR_CUTOFF",
      "MEM_POOL_RESERVE",
  };

  return std::string("MXNET_") + dev_type + "_" + name[type];
}

}  // namespace storage

const std::shared_ptr<Storage>& Storage::_GetSharedRef() {
#ifdef __MXNET_JS__
  // dummy code needed for emscripten code to pass
  // do not know why, the new will be NULLPTR
  static int* q = new int();
#endif
  static std::shared_ptr<Storage> inst(new storage::StorageImpl());
  return inst;
}

Storage* Storage::Get() {
  static Storage* ptr = _GetSharedRef().get();
  return ptr;
}
}  // namespace mxnet