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 "./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) {}
};


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 {
    return get_space_typed<xpu, ndim, real_t>(shape, stream);
  }
  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 {
    CHECK_EQ(req.type, ResourceRequest::kTempSpace);
    return mshadow::Tensor<xpu, ndim, DType>(
        reinterpret_cast<DType*>(get_space_internal(shape.Size() * sizeof(DType))),
        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,
      uint64_t seed) 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;
  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_