ndarray.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_NDARRAY_H_
#define MXNET_NDARRAY_H_

#include <dmlc/base.h>
#include <dmlc/logging.h>
#include <dmlc/io.h>
#include <dmlc/type_traits.h>
#include <dmlc/registry.h>
#include <nnvm/node.h>
#include <vector>
#include <map>
#include <string>
#include <memory>
#include "./base.h"
#include "./storage.h"
#include "./engine.h"
#if MKL_EXPERIMENTAL == 1
#include <mkl_memory.h>
#endif
// check c++11
#if DMLC_USE_CXX11 == 0
#error "cxx11 was required for ndarray module"
#endif

namespace mxnet {

// forward declaration
namespace autograd {
class AGNode;

using AGNodePtr = std::shared_ptr<AGNode>;

class AGNodeEntry {
 public:
  AGNodePtr ag_node;
  uint32_t index;
  uint32_t version;

  void clear() {
    ag_node.reset();
    index = version = 0;
  }

  nnvm::NodeEntry nn_entry() const;
  bool is_none() const;
};

class AutogradRuntime;
}  // namespace autograd

class NDArray {
 public:
  NDArray() {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = MKLMemHolder::create();
#endif
  }
  NDArray(const TShape &shape, Context ctx,
          bool delay_alloc = false, int dtype = mshadow::default_type_flag)
      : ptr_(std::make_shared<Chunk>(shape.Size(), ctx, delay_alloc, dtype)),
        shape_(shape), dtype_(dtype), entry_({nullptr, 0, 0}) {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = std::make_shared<MKLMemHolder>();
#endif
  }
  NDArray(const TBlob &data, int dev_id)
      : ptr_(std::make_shared<Chunk>(data, dev_id)), shape_(data.shape_),
        dtype_(data.type_flag_), entry_({nullptr, 0, 0}) {
#if MKL_EXPERIMENTAL == 1
    Mkl_mem_ = std::make_shared<MKLMemHolder>();
#endif
  }
  inline const TShape& shape() const {
    return shape_;
  }
  inline const TBlob& data() const {
    CheckAndAlloc();
    SetTBlob();
    return tblob_;
  }
  NDArray grad() const;
  inline Context ctx() const {
    return ptr_->shandle.ctx;
  }
  inline int dtype() const {
    return dtype_;
  }
  inline bool is_none() const {
    return ptr_.get() == nullptr;
  }
  bool fresh_out_grad() const;
  void set_fresh_out_grad(bool state) const;
  inline void WaitToRead() const {
    if (is_none()) return;
    Engine::Get()->WaitForVar(ptr_->var);
  }
  inline void WaitToWrite() const {
    if (is_none()) return;
    Engine::Get()->PushSync([](RunContext) {}, Context{}, {}, {ptr_->var});
    Engine::Get()->WaitForVar(ptr_->var);
  }
  inline Engine::VarHandle var() const {
    return ptr_->var;
  }
  void Save(dmlc::Stream *strm) const;
  bool Load(dmlc::Stream *strm);
  NDArray &operator=(real_t scalar);
  NDArray &operator+=(const NDArray &src);
  NDArray &operator+=(const real_t &src);
  NDArray &operator-=(const NDArray &src);
  NDArray &operator-=(const real_t &src);
  NDArray &operator*=(const NDArray &src);
  NDArray &operator*=(const real_t &src);
  NDArray &operator/=(const NDArray &src);
  NDArray &operator/=(const real_t &src);
  NDArray T() const;
  NDArray Copy(Context ctx) const;
  void SyncCopyFromCPU(const void *data, size_t size) const;
  void SyncCopyToCPU(void *data, size_t size) const;
  NDArray Slice(index_t begin, index_t end) const;
  NDArray At(index_t idx) const;
  inline NDArray AsArray(const TShape &shape, int dtype) const {
    CHECK_GE(shape_.Size() * mshadow::mshadow_sizeof(dtype_),
             shape.Size() * mshadow::mshadow_sizeof(dtype))
        << "NDArray.AsArray: target memory size is bigger";
#if MKL_EXPERIMENTAL == 1
    if (Mkl_mem_ != nullptr) {
      // convert prv to cpu
      Mkl_mem_->check_and_prv_to_cpu(ptr_->shandle.dptr);
    }
#endif
    NDArray ret = *this;
    ret.shape_ = shape;
    ret.dtype_ = dtype;
    return ret;
  }
  NDArray Reshape(const TShape &shape) const;
  NDArray Detach() const {
    NDArray ret(*this);
    ret.entry_ = autograd::AGNodeEntry{nullptr, 0, 0};
    return ret;
  }

  nnvm::Symbol get_autograd_symbol() {
    CHECK(!entry_.is_none())
      << "NDArray is not part of a computation graph. Did you forget to turn on recording?";
    nnvm::Symbol ret;
    ret.outputs.emplace_back(entry_.nn_entry());
    return ret;
  }
  inline void CheckAndAlloc() const {
    ptr_->CheckAndAlloc();
  }
  static void Save(dmlc::Stream* fo,
                   const std::vector<NDArray>& data,
                   const std::vector<std::string>& names);
  static void Load(dmlc::Stream* fi,
                   std::vector<NDArray>* data,
                   std::vector<std::string>* keys);

 private:
  friend class autograd::AutogradRuntime;
  struct Chunk {
    Storage::Handle shandle;
    Engine::VarHandle var;
    bool static_data;
    bool delay_alloc;
    Chunk() : static_data(true), delay_alloc(false) {
      var  = Engine::Get()->NewVariable();
    }
    Chunk(const TBlob &data, int dev_id)
        : static_data(true),
          delay_alloc(false) {
      var = Engine::Get()->NewVariable();
      if (data.dev_mask() == cpu::kDevMask) {
        shandle.ctx = Context::CPU();
      } else {
        CHECK_EQ(data.dev_mask(), gpu::kDevMask);
        shandle.ctx = Context::GPU(dev_id);
      }
      shandle.dptr = data.dptr_;
      shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
    }
    Chunk(uint64_t size, Context ctx, bool delay_alloc_, int dtype)
        : static_data(false), delay_alloc(true) {
      var = Engine::Get()->NewVariable();
      shandle.size = size * mshadow::mshadow_sizeof(dtype);
      shandle.ctx = ctx;
      if (!delay_alloc_) this->CheckAndAlloc();
    }
    inline void CheckAndAlloc(void) {
      if (delay_alloc) {
        shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
        delay_alloc = false;
      }
    }
    ~Chunk() {
      if (static_data || delay_alloc) {
        Engine::Get()->DeleteVariable([](RunContext s) {}, shandle.ctx, var);
      } else {
        Storage::Handle h = this->shandle;
        Engine::Get()->DeleteVariable([h](RunContext s) {
            Storage::Get()->Free(h);
          }, shandle.ctx, var);
      }
    }
  };

  void SetTBlob() const {
    tblob_.dptr_ = static_cast<char*>(ptr_->shandle.dptr) + byte_offset_;
    tblob_.shape_ = shape_;
    tblob_.type_flag_ = dtype_;
    tblob_.SetDLTensor(ptr_->shandle.ctx.dev_mask(), ptr_->shandle.ctx.dev_id);
#if MKL_EXPERIMENTAL == 1
    tblob_.Mkl_mem_ = Mkl_mem_;
#endif
  }

#if MKL_EXPERIMENTAL == 1
  std::shared_ptr<MKLMemHolder> Mkl_mem_;
#endif

  std::shared_ptr<Chunk> ptr_;
  TShape shape_;
  size_t byte_offset_ = 0;
  int dtype_ = -1;
  autograd::AGNodeEntry entry_;
  mutable TBlob tblob_;
};

void CopyFromTo(const NDArray &from, NDArray *to, int priority = 0);


void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);

NDArray operator+(const NDArray &lhs, const NDArray &rhs);
NDArray operator+(const NDArray &lhs, const real_t &rhs);
NDArray operator-(const NDArray &lhs, const NDArray &rhs);
NDArray operator-(const NDArray &lhs, const real_t &rhs);
NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
NDArray operator*(const NDArray &lhs, const real_t &rhs);
NDArray operator/(const NDArray &lhs, const NDArray &rhs);
NDArray operator/(const NDArray &lhs, const real_t &rhs);

void RandomSeed(uint32_t seed);
void SampleUniform(real_t begin, real_t end, NDArray *out);
void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
void SampleGamma(real_t alpha, real_t beta, NDArray *out);
void SampleExponential(real_t lambda, NDArray *out);
void SamplePoisson(real_t lambda, NDArray *out);
void SampleNegBinomial(int32_t k, real_t p, NDArray *out);
void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out);


//--------------------------------------------------------------
// The following part are API Registration of NDArray functions.
//--------------------------------------------------------------

typedef std::function<void (NDArray **used_vars,
                            real_t *scalars,
                            NDArray **mutate_vars,
                            int num_params,
                            char **param_keys,
                            char **param_vals)> NDArrayAPIFunction;
enum NDArrayFunctionTypeMask {
  kNDArrayArgBeforeScalar = 1,
  kScalarArgBeforeNDArray = 1 << 1,
  kAcceptEmptyMutateTarget = 1 << 2
};
struct NDArrayFunctionReg
    : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
                                        NDArrayAPIFunction> {
  unsigned num_use_vars;
  unsigned num_mutate_vars;
  unsigned num_scalars;
  int type_mask;
  NDArrayFunctionReg()
      : num_use_vars(0),
        num_mutate_vars(0),
        num_scalars(0),
        type_mask(0) {}
  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
                                                            NDArray *out)) {
    body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                        int num_params, char **param_keys, char **param_vals) {
      (*fsetvalue)(s[0], mutate_vars[0]);
    };
    num_mutate_vars = 1; num_scalars = 1;
    this->add_argument("src", "real_t", "Source input to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
                                                          const NDArray &mhs,
                                                          const NDArray &rhs,
                                                                NDArray *out)) {
    body = [fternary](NDArray **used_vars,
      real_t *s, NDArray **mutate_vars,
      int num_params, char **param_keys, char **param_vals) {
      (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
    };
    num_use_vars = 3; num_mutate_vars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("mhs", "NDArray", "Middle operand to the function.");
    this->add_argument("rhs", "NDArray", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
                                                          const NDArray &rhs,
                                                          NDArray *out)) {
    body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                      int num_params, char **param_keys, char **param_vals) {
      (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
    };
    num_use_vars = 2; num_mutate_vars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("rhs", "NDArray", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
                                                          const real_t &rhs,
                                                          NDArray *out)) {
    body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                      int num_params, char **param_keys, char **param_vals) {
      (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
    };
    num_use_vars = 1; num_mutate_vars = 1; num_scalars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("rhs", "real_t", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
                                                         NDArray *out)) {
    body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                     int num_params, char **param_keys, char **param_vals) {
      (*funary)(*used_vars[0], mutate_vars[0]);
    };
    num_use_vars = 1; num_mutate_vars = 1;
    type_mask = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("src", "NDArray", "Source input to the function.");
    return *this;
  }
  inline NDArrayFunctionReg &set_function(
    void (*fgeneric)(NDArray **used_vars,
                     real_t *s,
                     NDArray **mutate_vars,
                     const std::map<std::string, std::string>& param)) {
    body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
                       int num_params, char **param_keys, char **param_vals) {
      std::map<std::string, std::string> param;
      for (int i = 0; i < num_params; ++i) {
        param[param_keys[i]] = param_vals[i];
      }
      fgeneric(used_vars, s, mutate_vars, param);
    };
    return *this;
  }
  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
    num_use_vars = n; return *this;
  }
  inline NDArrayFunctionReg &set_num_mutate_vars(unsigned n) {
    num_mutate_vars = n; return *this;
  }
  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
    num_scalars = n; return *this;
  }
  inline NDArrayFunctionReg &set_type_mask(int tmask) {
    type_mask = tmask; return *this;
  }
};  // NDArrayFunctionReg

#define MXNET_REGISTER_NDARRAY_FUN(name)                                 \
  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)

}  // namespace mxnet

namespace dmlc {
DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
}  // namespace dmlc
#endif  // MXNET_NDARRAY_H_