Enable resource Allocator

include/mxnet/base.h

@@ -1,12 +1,13 @@
 /*!
  *  Copyright (c) 2015 by Contributors
  * \file base.h
- * \brief configuation of mxnet
+ * \brief configuation of mxnet as well as basic data structure.
  */
 #ifndef MXNET_BASE_H_
 #define MXNET_BASE_H_
 
 #include <dmlc/base.h>
+#include <dmlc/io.h>
 #include <dmlc/type_traits.h>
 #include <dmlc/parameter.h>
 #include <mshadow/tensor.h>
@@ -62,6 +63,84 @@ typedef mshadow::default_real_t real_t;
 typedef mshadow::TShape TShape;
 /*! \brief storage container type */
 typedef mshadow::TBlob TBlob;
+
+/*! \brief Context information about the execution enviroment */
+struct Context {
+  /*! \brief the device type we run the op can be cpu::kDevMask or gpu::kDevMask */
+  int32_t dev_mask;
+  /*! \brief device id we are going to run it on */
+  int32_t dev_id;
+  /*! \brief constructor */
+  Context() : dev_mask(cpu::kDevMask), dev_id(0) {}
+  /*!
+   * \brief constructor of context
+   * \param dev_mask the device mask
+   * \param dev_id the device id
+   */
+  Context(int dev_mask, int dev_id)
+      : dev_mask(dev_mask), dev_id(dev_id) {}
+  /*!
+   * \brief check if current context equals another one
+   * \param b another context to compare
+   * \return whether dev mask and id are same
+   */
+  inline bool operator==(const Context &b) const {
+    return dev_mask == b.dev_mask && dev_id == b.dev_id;
+  }
+  /*!
+   * \brief check if current context not equals another one
+   * \param b another context to compare
+   * \return whether they are not the same
+   */
+  inline bool operator!=(const Context &b) const {
+    return !(*this == b);
+  }
+  /*!
+   * \brief save the content into binary stream
+   * \param strm the output stream
+   */
+  void Save(dmlc::Stream *strm) const {
+    strm->Write(&dev_mask, sizeof(dev_mask));
+    strm->Write(&dev_id, sizeof(dev_id));
+  }
+  /*!
+   * \brief load the content from binary stream
+   * \param strm the output stream
+   * \return whether the load is successful
+   */
+  bool Load(dmlc::Stream *strm) {
+    if (strm->Read(&dev_mask, sizeof(int32_t)) != sizeof(int32_t)) return false;
+    if (strm->Read(&dev_id, sizeof(int32_t)) != sizeof(int32_t)) return false;
+    return true;
+  }
+  /*! \brief the maximal device mask, cpu = 1, gpu = 2 */
+  static const int32_t kMaxDevMask = 2;
+  /*!
+   * \brief A dedicate ID for pinned cpu memory.
+   *  Any normal CPU ID should be less than this number.
+   */
+  static const int32_t kPinnedMemoryID = 16;
+};
+
+/*!
+ * \brief execution time context.
+ *  The information needed in runtime for actual execution.
+ */
+struct RunContext {
+  /*!
+   * \brief the stream of the device, can be NULL or Stream<gpu>* in GPU mode
+   */
+  void *stream;
+  /*!
+   * \brief get mshadow stream from Context
+   * \return the mshadow stream
+   * \tparam xpu the device type of the stream
+   */
+  template<typename xpu>
+  inline mshadow::Stream<xpu>* get_stream() const {
+    return static_cast<mshadow::Stream<xpu>*>(stream);
+  }
+};
 }  // namespace mxnet
 
 //! \cond Doxygen_Suppress

include/mxnet/engine.h

@@ -12,7 +12,6 @@
 #endif
 #include <vector>
 #include "./base.h"
-#include "./context.h"
 
 namespace mxnet {
 /*! \brief namespace of engine internal types. */
@@ -28,7 +27,6 @@ typedef Opr* OprHandle;
 }  // namespace engine
 
 #if DMLC_USE_CXX11
-
 /*! \brief Function property, used to hint what action is pushed to engine. */
 enum class FnProperty {
   /*! \brief Normal operation */

include/mxnet/kvstore.h

@@ -10,7 +10,7 @@
 #if DMLC_USE_CXX11
 #include <functional>
 #endif  // DMLC_USE_CXX11
-#include "ndarray.h"
+#include "./ndarray.h"
 
 namespace mxnet {
 

include/mxnet/ndarray.h

@@ -13,10 +13,9 @@
 #include <dmlc/registry.h>
 #include <memory>
 #include "./base.h"
-#include "./context.h"
 #include "./storage.h"
-#include "./context.h"
 #include "./engine.h"
+
 // check c++11
 #if DMLC_USE_CXX11 == 0
 #error "cxx11 was required for ndarray module"

include/mxnet/operator.h

@@ -14,7 +14,7 @@
 #include <string>
 #include <utility>
 #include "./base.h"
-#include "./context.h"
+#include "./resource.h"
 
 namespace mxnet {
 /*! \brief operation request type to Forward and Backward */

include/mxnet/resource.h

@@ -0,0 +1,85 @@
+/*!
+ *  Copyright (c) 2015 by Contributors
+ * \file resource.h
+ * \brief Global resource allocation handling.
+ */
+#ifndef MXNET_RESOURCE_H_
+#define MXNET_RESOURCE_H_
+
+#include <dmlc/logging.h>
+#include "./base.h"
+#include "./engine.h"
+
+namespace mxnet {
+
+/*!
+ * \brief The resources that can be requested by Operator
+ */
+struct ResourceRequest {
+  /*! \brief Resource type, indicating what the pointer type is */
+  enum Type {
+    /*! \brief mshadow::Random<xpu> object */
+    kRandom,
+    /*! \brief Temporal space */
+    kTempSpace
+  };
+  /*! \brief type of resources */
+  Type type;
+  /*! \brief size of space requested, in terms of number of reals */
+  size_t space_num_reals;
+  /*! \brief default constructor */
+  ResourceRequest() {}
+  /*!
+   * \brief constructor, allow implicit conversion
+   * \param type type of resources
+   */
+  ResourceRequest(Type type, size_t space_num_reals = 0)  // NOLINT(*)
+      : type(type), space_num_reals(space_num_reals) {}
+};
+
+
+/*!
+ * \brief Resources used by mxnet operations.
+ *  A resource is something special other than NDArray,
+ *  but will still participate
+ */
+struct Resource {
+  /*! \brief The original request */
+  ResourceRequest req;
+  /*! \brief engine variable */
+  engine::VarHandle var;
+  /*!
+   * \brief pointer to the resource, do not use directly,
+   *  access using member functions
+   */
+  void *ptr_;
+  /*!
+   * \brief Get random number generator.
+   * \return the mshadow random number generator requested.
+   * \tparam xpu the device type of random number generator.
+   */
+  template<typename xpu>
+  inline mshadow::Random<xpu>* get_random() const {
+    CHECK_EQ(req.type, ResourceRequest::kRandom);
+    return static_cast<mshadow::Random<xpu>*>(ptr_);
+  }
+  /*!
+   * \brief Get space requested as mshadow Tensor.
+   *  The resulting tensor must fit in space requsted.
+   * \param shape the Shape of returning tensor.
+   * \param stream the stream of retruning tensor.
+   * \return the mshadow tensor requested.
+   * \tparam xpu the device type of random number generator.
+   * \tparam ndim the number of dimension of the tensor requested.
+   */
+  template<typename xpu, int ndim>
+  inline mshadow::Tensor<xpu, ndim, real_t> get_space(
+      mshadow::Shape<ndim> shape, mshadow::Stream<xpu> *stream) const {
+    CHECK_EQ(req.type, ResourceRequest::kTempSpace);
+    CHECK_GE(req.space_num_reals, shape.Size());
+    return mshadow::Tensor<xpu, ndim, real_t>(
+        static_cast<real_t*>(ptr_), shape, shape[ndim - 1], stream);
+  }
+};
+}  // namespace mxnet
+#endif  // MXNET_RESOURCE_H_

include/mxnet/storage.h

@@ -7,8 +7,7 @@
 #define MXNET_STORAGE_H_
 
 #include <memory>
-#include "base.h"
-#include "context.h"
+#include "./base.h"
 
 namespace mxnet {
 
@@ -64,7 +63,5 @@ class Storage {
   std::unique_ptr<Impl> impl_;
   DISALLOW_COPY_AND_ASSIGN(Storage);
 };  // class Storage
-
 }  // namespace mxnet
-
 #endif  // MXNET_STORAGE_H_

src/engine/engine.cc

@@ -13,7 +13,7 @@ namespace engine {
 inline Engine* CreateEngine() {
   const char *type = getenv("MXNET_ENGINE_TYPE");
   const bool default_engine = (type == nullptr);
-  if (type == nullptr) type = "ThreadedEngine";
+  if (type == nullptr) type = "ThreadedEnginePerDevice";
   std::string stype = type;
 
   Engine *ret = nullptr;