mr.cpp

namespace {

device_memory_resource* cuda_resource() {
  static cuda_memory_resource resource{};
  return &resource;

  std::atomic<device_memory_resource*>& get_default() {
    static std::atomic<device_memory_resource*> res{detail::cuda_resource()};
    return res;
  }
}  // namespace

device_memory_resource* get_default_resource() { return get_default().load(); }

device_memory_resource* set_default_resource(
    device_memory_resource* new_resource) {
  new_resource =
      (new_resource == nullptr) ? detail::cuda_resource() : new_resource;
  return get_default().exchange(new_resource);
}

enum class resource { DEFAULT, POOL, MANAGED, MANAGED_POOL };

/**---------------------------------------------------------------------------*
 * @brief Maintains unique global state for lifetime of library.
*---------------------------------------------------------------------------**/
class state_singleton {
 public:
  state_singleton& Instance() {
    static state_singleton instance;
    return instance;
  }

  /**---------------------------------------------------------------------------*
   * @brief Set the default resource to the requested resource.
   *
   * Does nothing if `initialize()` was already called without then calling
   * `finalize()`.
   *
   * @param mode The requested resource
   *---------------------------------------------------------------------------**/
  void initialize(resource mode) {
    if (not initialized) {
      initialized = true;
      switch (mode) {
        case resource::DEFAULT:
          set_default_resource(nullptr);
        case resource::POOL:
          set_default_resource(pool.reset(new pool_memory_resource{}).get());
        case resource::MANAGED:
          set_default_resource(managed.reset(new managed_memory_resource{}).get());
        case resource::MANAGED_POOL:
          set_default_resource(managed_pool.reset(new managed_pool_memory_resource{}).get());
      }
    }
  }

  /**---------------------------------------------------------------------------*
   * @brief Frees memory held by all resources.
   *
   * Does nothing if `initialize()` was not already called.
   *
   *---------------------------------------------------------------------------**/
  void finalize() {
    if (initiazed) {
      pool.reset();
      managed.reset();
      managed_pool.reset();
      set_default_resource(nullptr);
    }
  }

 private:
  bool initialized{false};
  std::unique_ptr<pool_memory_resource> pool{};
  std::unique_ptr<managed_memory_resource> managed{};
  std::unique_ptr<managed_pool_memory_resource> managed_pool{};
};