resource.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#ifndef MXNET_RESOURCE_H_
#define MXNET_RESOURCE_H_
 
#include <dmlc/logging.h>
#include <string>
#include "./base.h"
#include "./engine.h"
#include "./random_generator.h"
 
namespace mxnet {
 
struct ResourceRequest {
  enum Type {
    kRandom,
    kTempSpace,
    kParallelRandom
#if MXNET_USE_CUDNN == 1
    ,
    kCuDNNDropoutDesc
#endif  // MXNET_USE_CUDNN == 1
  };
  Type type;
  ResourceRequest() {}
  ResourceRequest(Type type)  // NOLINT(*)
      : type(type) {}
};
 
namespace {
inline std::string __extract_fname(const std::string& path) {
  std::size_t last_dir_pos = path.find_last_of("/\\");
  if (last_dir_pos == std::string::npos) {
    return path;
  }
  return path.substr(last_dir_pos + 1);
}
}  // anonymous namespace
 
#if (defined(__GNUC__) || defined(__GNUG__)) && !defined(__clang__)
#define MXNET_RESOURCE_DEFAULT_NAME_FARG(tag)                          \
  std::string(tag) + " (" + __extract_fname(__builtin_FILE()) + " +" + \
      std::to_string(__builtin_LINE()) + ")"
#else  // !__GNUC__ || __clang__
#define MXNET_RESOURCE_DEFAULT_NAME_FARG(tag) \
  std::string(tag) + " (" + __extract_fname(__FILE__) + " +" + std::to_string(__LINE__) + ")"
#endif  // __GNUC__ && !__clang__
 
struct Resource {
  ResourceRequest req;
  engine::VarHandle var;
  int32_t id;
  void* ptr_;
  Resource() : id(0) {}
  template <typename xpu, typename DType>
  inline mshadow::Random<xpu, DType>* get_random(mshadow::Stream<xpu>* stream) const {
    CHECK_EQ(req.type, ResourceRequest::kRandom);
    mshadow::Random<xpu, DType>* ret = static_cast<mshadow::Random<xpu, DType>*>(ptr_);
    ret->set_stream(stream);
    return ret;
  }
 
  template <typename xpu, typename DType>
  inline common::random::RandGenerator<xpu, DType>* get_parallel_random() const {
    CHECK_EQ(req.type, ResourceRequest::kParallelRandom);
    return static_cast<common::random::RandGenerator<xpu, DType>*>(ptr_);
  }
 
  template <typename xpu, int ndim>
  inline mshadow::Tensor<xpu, ndim, real_t> get_space(
      mshadow::Shape<ndim> shape,
      mshadow::Stream<xpu>* stream,
      const std::string& name = MXNET_RESOURCE_DEFAULT_NAME_FARG("temp_space")) const {
    return get_space_typed<xpu, ndim, real_t>(shape, stream, name);
  }
  template <int ndim>
  inline mshadow::Tensor<cpu, ndim, real_t> get_host_space(mshadow::Shape<ndim> shape) const {
    return get_host_space_typed<cpu, ndim, real_t>(shape);
  }
  template <typename xpu, int ndim, typename DType>
  inline mshadow::Tensor<xpu, ndim, DType> get_space_typed(
      mshadow::Shape<ndim> shape,
      mshadow::Stream<xpu>* stream,
      const std::string& name = MXNET_RESOURCE_DEFAULT_NAME_FARG("temp_space")) const {
    CHECK_EQ(req.type, ResourceRequest::kTempSpace);
    return mshadow::Tensor<xpu, ndim, DType>(
        reinterpret_cast<DType*>(get_space_internal(shape.Size() * sizeof(DType), name)),
        shape,
        shape[ndim - 1],
        stream);
  }
#if MXNET_USE_CUDNN == 1
 
  void get_cudnn_dropout_desc(
      cudnnDropoutDescriptor_t* dropout_desc,
      mshadow::Stream<gpu>* stream,
      const float dropout,
      const std::string& name = MXNET_RESOURCE_DEFAULT_NAME_FARG("cudnn_dropout_state")) const;
#endif  // MXNET_USE_CUDNN == 1
 
  template <int ndim, typename DType>
  inline mshadow::Tensor<cpu, ndim, DType> get_host_space_typed(mshadow::Shape<ndim> shape) const {
    return mshadow::Tensor<cpu, ndim, DType>(
        reinterpret_cast<DType*>(get_host_space_internal(shape.Size() * sizeof(DType))),
        shape,
        shape[ndim - 1],
        nullptr);
  }
  void* get_space_internal(size_t size, const std::string& name) const;
  void* get_host_space_internal(size_t size) const;
};
 
class ResourceManager {
 public:
  virtual Resource Request(Context ctx, const ResourceRequest& req) = 0;
  virtual void SeedRandom(uint32_t seed) = 0;
  virtual void SeedRandom(Context ctx, uint32_t seed) = 0;
  virtual ~ResourceManager() DMLC_THROW_EXCEPTION {}
  static ResourceManager* Get();
};
}  // namespace mxnet
#endif  // MXNET_RESOURCE_H_