mxnet
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros
ndarray.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
25 #ifndef MXNET_NDARRAY_H_
26 #define MXNET_NDARRAY_H_
27 
28 #include <dmlc/base.h>
29 #include <dmlc/logging.h>
30 #include <dmlc/io.h>
31 #include <dmlc/type_traits.h>
32 #include <dmlc/registry.h>
33 #include <nnvm/node.h>
34 #include <vector>
35 #include <map>
36 #include <string>
37 #include <memory>
38 #include "./base.h"
39 #include "./storage.h"
40 #include "./engine.h"
41 #if MKL_EXPERIMENTAL == 1
42 #include <mkl_memory.h>
43 #endif
44 // check c++11
45 #if DMLC_USE_CXX11 == 0
46 #error "cxx11 was required for ndarray module"
47 #endif
48 
49 namespace mxnet {
50 // enum for storage types
51 namespace csr {
53 }
54 
55 namespace rowsparse {
57 }
58 
60  kUndefinedStorage = -1, // undefined storage
61  kDefaultStorage, // dense
62  kRowSparseStorage, // row sparse
63  kCSRStorage, // csr
64 };
65 
67  kNormalErr, // normal
68  kCSRShapeErr, // shape mismatch for csr
69  kCSRIndPtrErr, // indptr error for csr
70  kCSRIdxErr, // idx error for csr
71  kRSPShapeErr, // shape mismatch for row sparse
72  kRSPIdxErr, // indices error for row sparse
73 };
74 
75 
79 class NDArray {
80  public:
82  NDArray() {
83 #if MKL_EXPERIMENTAL == 1
84  Mkl_mem_ = MKLMemHolder::create();
85 #endif
86  }
94  NDArray(const TShape &shape, Context ctx,
95  bool delay_alloc = false, int dtype = mshadow::default_type_flag)
96  : ptr_(std::make_shared<Chunk>(shape, ctx, delay_alloc, dtype)),
97  shape_(shape), dtype_(dtype), storage_type_(kDefaultStorage),
98  entry_({nullptr, 0, 0}) {
99 #if MKL_EXPERIMENTAL == 1
100  Mkl_mem_ = std::make_shared<MKLMemHolder>();
101 #endif
102  }
105  NDArray(const NDArrayStorageType stype, const TShape &shape, Context ctx,
106  bool delay_alloc = true, int dtype = mshadow::default_type_flag,
107  std::vector<int> aux_types = {}, std::vector<TShape> aux_shapes = {},
108  TShape storage_shape = TShape(mshadow::Shape1(0)))
109  : shape_(shape), dtype_(dtype), storage_type_(stype),
110  entry_({nullptr, 0, 0}) {
111  // Assign default aux types if not given
112  if (aux_types.size() == 0) {
113  if (stype == kRowSparseStorage) {
114  aux_types = {mshadow::kInt64};
115  } else if (stype == kCSRStorage) {
116  aux_types = {mshadow::kInt64, mshadow::kInt64};
117  } else {
118  LOG(FATAL) << "Unknown storage type " << stype;
119  }
120  }
121  // Assign default shapes if not given
122  // unknown shapes are intialized as {0} such that Size() would return 0
123  if (aux_shapes.size() == 0) {
124  if (stype == kRowSparseStorage) {
125  aux_shapes = {TShape(mshadow::Shape1(0))};
126  } else if (stype == kCSRStorage) {
127  // aux shapes for indptr and indices
128  aux_shapes = {TShape(mshadow::Shape1(0)), TShape(mshadow::Shape1(0))};
129  } else {
130  LOG(FATAL) << "Unknown storage type " << stype;
131  }
132  }
133  if (storage_shape.Size() == 0) {
134  if (stype == kRowSparseStorage) {
135  storage_shape = shape;
136  storage_shape[0] = aux_shapes[rowsparse::kIdx][0];
137  } else if (stype == kCSRStorage) {
138  storage_shape = aux_shapes[csr::kIdx];
139  } else {
140  LOG(FATAL) << "Unknown storage type " << stype;
141  }
142  }
143  ptr_ = std::make_shared<Chunk>(stype, storage_shape, ctx, delay_alloc,
144  dtype, aux_types, aux_shapes);
145 #if MKL_EXPERIMENTAL == 1
146  Mkl_mem_ = std::make_shared<MKLMemHolder>();
147 #endif
148  }
156  NDArray(const TBlob &data, int dev_id)
157  : ptr_(std::make_shared<Chunk>(data, dev_id)), shape_(data.shape_),
158  dtype_(data.type_flag_), storage_type_(kDefaultStorage),
159  entry_({nullptr, 0, 0}) {
160 #if MKL_EXPERIMENTAL == 1
161  Mkl_mem_ = std::make_shared<MKLMemHolder>();
162 #endif
163  }
165  NDArray(int shared_pid, int shared_id, const TShape& shape, int dtype)
166  : ptr_(std::make_shared<Chunk>(shared_pid, shared_id, shape, dtype)), shape_(shape),
167  dtype_(dtype), storage_type_(kDefaultStorage), entry_({nullptr, 0, 0}) {
168 #if MKL_EXPERIMENTAL == 1
169  Mkl_mem_ = std::make_shared<MKLMemHolder>();
170 #endif
171  }
172 
183  NDArray(const NDArrayStorageType stype, const TShape &shape,
184  const TBlob &data, const std::vector<TBlob> &aux_data, int dev_id)
185  : ptr_(std::make_shared<Chunk>(stype, data, aux_data, dev_id)), shape_(shape),
186  dtype_(data.type_flag_), storage_type_(stype), entry_({nullptr, 0, 0}) {
187 #if MKL_EXPERIMENTAL == 1
188  Mkl_mem_ = std::make_shared<MKLMemHolder>();
189 #endif
190  }
191 
192 
196  inline const TShape& shape() const {
197  return shape_;
198  }
204  inline const TShape &storage_shape() const {
205  CHECK(ptr_ != nullptr);
206  CHECK_NE(storage_type(), kDefaultStorage)
207  << "storage_shape() is not intended for kDefaultStorage.";
208  return ptr_->storage_shape;
209  }
210 
216  inline const TShape& aux_shape(size_t index) const {
217  CHECK_NE(storage_type(), kDefaultStorage)
218  << "aux_shape() is not intended for kDefaultStorage.";
219  return ptr_->aux_shapes[index];
220  }
221 
222  /* \return the shapes of all aux data */
223  const std::vector<TShape>& aux_shapes() const {
224  CHECK_NE(storage_type(), kDefaultStorage)
225  << "aux_shapes() is not intended for kDefaultStorage.";
226  return ptr_->aux_shapes;
227  }
228 
230  const std::vector<int>& aux_types() const {
231  CHECK_NE(storage_type(), kDefaultStorage)
232  << "aux_types() is not intended for kDefaultStorage.";
233  return ptr_->aux_types;
234  }
235 
243  inline void set_aux_shape(size_t index, const TShape& shape) const {
244  ptr_->set_aux_shape(index, shape);
245  }
246 
250  inline const TBlob& data() const {
252  SetTBlob();
253  return tblob_;
254  }
258  NDArray grad() const;
259 
263  inline TBlob aux_data(size_t i) const {
264  auto stype = storage_type();
265  TBlob res;
266  auto shape = aux_shape(i);
267  auto type = aux_type(i);
268  MSHADOW_TYPE_SWITCH(type, DType, {
269  auto dptr = static_cast<DType*>(ptr_->aux_handles[i].dptr);
270  CHECK(stype == kRowSparseStorage || stype == kCSRStorage)
271  << "Unexpected storage type: " << stype;
272  res = TBlob(dptr, shape, ptr_->aux_handles[i].ctx.dev_mask(), type);
273  });
274 #if MKL_EXPERIMENTAL == 1
275  res.Mkl_mem_ = Mkl_mem_;
276 #endif
277  return res;
278  }
282  inline Context ctx() const {
283  CHECK(!is_none());
284  return ptr_->shandle.ctx;
285  }
289  inline int dtype() const {
290  return dtype_;
291  }
292  inline int aux_type(size_t i) const {
293  CHECK(!is_none());
294  return ptr_->aux_types[i];
295  }
296 
298  return storage_type_;
299  }
301  inline bool is_none() const {
302  return ptr_.get() == nullptr;
303  }
305  bool fresh_out_grad() const;
307  void set_fresh_out_grad(bool state) const;
308  // returns true if a sparse ndarray's aux_data and storage are initialized
309  inline bool storage_initialized() const {
310  if (is_none()) return false;
311  auto stype = storage_type();
312  CHECK_NE(stype, kDefaultStorage)
313  << "storage_initialized() is not intended for kDefaultStorage.";
314  if (stype == kRowSparseStorage) {
315  CHECK_EQ(aux_shape(rowsparse::kIdx)[0], storage_shape()[0])
316  << "inconsistent storage shape " << storage_shape()
317  << " vs. aux shape " << aux_shape(rowsparse::kIdx);
318  return aux_shape(0).Size() != 0;
319  } else if (stype == kCSRStorage) {
320  CHECK_EQ(aux_shape(csr::kIdx)[0], storage_shape()[0])
321  << "inconsistent storage shape " << storage_shape()
322  << " vs. aux shape " << aux_shape(csr::kIdx);
323  return aux_shape(0).Size() != 0;
324  } else {
325  LOG(FATAL) << "Unknown storage type";
326  }
327  return true;
328  }
331  CHECK(!is_none());
332  CHECK_EQ(storage_type(), kDefaultStorage);
333  CheckAndAlloc();
334  return ptr_->shandle;
335  }
340  inline void WaitToRead() const {
341  if (is_none()) return;
342  Engine::Get()->WaitForVar(ptr_->var);
343  }
348  inline void WaitToWrite() const {
349  if (is_none()) return;
355  [](RunContext, Engine::CallbackOnComplete on_complete) {
356  on_complete();
357  }, Context{}, {}, {ptr_->var});
358  Engine::Get()->WaitForVar(ptr_->var);
359  }
361  inline Engine::VarHandle var() const {
362  return ptr_->var;
363  }
368  void Save(dmlc::Stream *strm) const;
374  bool LegacyLoad(dmlc::Stream *strm, const uint32_t magic);
380  bool Load(dmlc::Stream *strm);
386  NDArray &operator=(real_t scalar);
393  NDArray &operator+=(const NDArray &src);
400  NDArray &operator+=(const real_t &src);
407  NDArray &operator-=(const NDArray &src);
414  NDArray &operator-=(const real_t &src);
421  NDArray &operator*=(const NDArray &src);
428  NDArray &operator*=(const real_t &src);
435  NDArray &operator/=(const NDArray &src);
442  NDArray &operator/=(const real_t &src);
448  NDArray Copy(Context ctx) const;
459  void SyncCopyFromCPU(const void *data, size_t size) const;
460 
464  void SyncCopyFromNDArray(const NDArray &src, int i = -1, int j = -1);
465 
476  void SyncCopyToCPU(void *data, size_t size) const;
482  void SyncCheckFormat(const bool full_check) const;
489  NDArray Slice(index_t begin, index_t end) const;
502  NDArray At(index_t idx) const;
513  NDArray aux_ndarray(size_t i) const;
514 
519  NDArray data_ndarray() const;
520 
528  inline NDArray AsArray(const TShape &shape, int dtype) const {
529  CHECK_EQ(storage_type(), kDefaultStorage)
530  << "AsArray is intended only for kDefaultStorage.";
531  CHECK_GE(ptr_->shandle.size,
532  shape.Size() * mshadow::mshadow_sizeof(dtype))
533  << "NDArray.AsArray: target memory size is bigger";
534 #if MKL_EXPERIMENTAL == 1
535  if (Mkl_mem_ != nullptr) {
536  // convert prv to cpu
537  Mkl_mem_->check_and_prv_to_cpu(ptr_->shandle.dptr);
538  }
539 #endif
540  NDArray ret = *this;
541  ret.shape_ = shape;
542  ret.dtype_ = dtype;
543  return ret;
544  }
550  NDArray Reshape(const TShape &shape) const;
556  NDArray ReshapeWithRecord(const TShape &shape);
560  NDArray Detach() const {
561  NDArray ret(*this);
562  ret.entry_ = nnvm::NodeEntry{nullptr, 0, 0};
563  return ret;
564  }
565 
566  nnvm::Symbol get_autograd_symbol() const;
571  inline void CheckAndAlloc() const {
572  CHECK_EQ(storage_type(), kDefaultStorage);
573  ptr_->CheckAndAlloc();
574  }
575 
585  void ReshapeAndAlloc(const TShape& shape) {
586  CHECK_EQ(storage_type(), kDefaultStorage);
587  CHECK(!is_none());
588  shape_ = shape;
589  ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
590  }
591 
592  /* !
593  * \brief Alloc memory for non-default storage
594  * aux_shape is only known at run time
595  */
596  inline void CheckAndAlloc(const std::vector<TShape> &aux_shapes) const {
597  CHECK_NE(storage_type(), kDefaultStorage)
598  << "CheckAndAlloc(aux_shapes) is not intended for kDefaultStorage";
599  ptr_->CheckAndAlloc(shape_, aux_shapes, dtype_);
600  }
601  inline void CheckAndAllocData(const TShape &storage_shape) const {
602  CHECK_NE(storage_type(), kDefaultStorage)
603  << "CheckAndAllocData is not intended for kDefaultStorage";
604  ptr_->CheckAndAllocData(storage_shape, dtype_);
605  }
606  inline void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const {
607  CHECK_NE(storage_type(), kDefaultStorage)
608  << "CheckAndAllocAuxData is not intended for kDefaultStorage";
609  ptr_->CheckAndAllocAuxData(i, aux_shape);
610  }
617  static void Save(dmlc::Stream* fo,
618  const std::vector<NDArray>& data,
619  const std::vector<std::string>& names);
626  static void Load(dmlc::Stream* fi,
627  std::vector<NDArray>* data,
628  std::vector<std::string>* keys);
629 
630  private:
631  friend class Imperative;
633  // shandle is used to store the actual values in the NDArray
634  // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
635  struct Chunk {
639  Storage::Handle shandle;
644  std::vector<Storage::Handle> aux_handles;
652  bool static_data;
655  bool delay_alloc;
656  // the type of the storage. The storage_type is never kUndefinedStorage once the chunk
657  // is constructed.
660  std::vector<int> aux_types;
661  // context of data
662  Context ctx;
663  // The shape of the chunk data.
664  // This might not be the same shape as the NDArray, since the storage may be sparse.
665  // The default value for storage_shape is {0} when an empty non-default NDArray is created.
666  TShape storage_shape;
667  // The shape of aux data. The default value for the shape depends on the type of storage.
668  // If aux_shapes[i].Size() is zero, aux data i is empty.
669  std::vector<TShape> aux_shapes;
670 
672  Chunk() : static_data(true), delay_alloc(false) {}
673 
675  Chunk(TShape shape, Context ctx_, bool delay_alloc_, int dtype)
676  : static_data(false), delay_alloc(true), ctx(ctx_) {
677  auto size = shape.Size();
678  storage_shape = shape;
679  var = Engine::Get()->NewVariable();
680  shandle.size = size * mshadow::mshadow_sizeof(dtype);
681  shandle.ctx = ctx_;
682  if (!delay_alloc_) this->CheckAndAlloc();
683  }
684 
685  Chunk(const TBlob &data, int dev_id)
686  : static_data(true), delay_alloc(false) {
687  CHECK(storage_type == kDefaultStorage);
688  var = Engine::Get()->NewVariable();
689  if (data.dev_mask() == cpu::kDevMask) {
690  ctx = Context::CPU();
691  } else {
692  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
693  ctx = Context::GPU(dev_id);
694  }
695  // init shandle
696  shandle.ctx = ctx;
697  shandle.dptr = data.dptr_;
698  shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
699  storage_shape = data.shape_;
700  }
701 
702  Chunk(int shared_pid, int shared_id, const TShape& shape, int dtype)
703  : static_data(false), delay_alloc(false) {
704  var = Engine::Get()->NewVariable();
705  ctx = Context::CPUShared(0);
706  shandle.size = shape.Size() * mshadow::mshadow_sizeof(dtype);;
707  shandle.ctx = ctx;
708  shandle.shared_pid = shared_pid;
709  shandle.shared_id = shared_id;
710  Storage::Get()->Alloc(&shandle);
711  storage_shape = shape;
712  }
713  // Constructor for a non-default storage chunk
714  Chunk(NDArrayStorageType storage_type_, const TShape &storage_shape_, Context ctx_,
715  bool delay_alloc_, int dtype, const std::vector<int> &aux_types_,
716  const std::vector<TShape> &aux_shapes_)
717  : static_data(false), delay_alloc(delay_alloc_), storage_type(storage_type_),
718  aux_types(aux_types_), ctx(ctx_), storage_shape(storage_shape_),
719  aux_shapes(aux_shapes_) {
720  shandle.ctx = ctx;
721  var = Engine::Get()->NewVariable();
722  // aux_handles always reflect the correct number of aux data
723  for (size_t i = 0; i < aux_shapes.size(); i++) {
724  CheckAndAllocAuxData(i, aux_shapes[i]);
725  // this line is needed in case when aux_shapes[i].Size() = 0
726  // aux_handles[i] will not be updated and take only default value.
727  aux_handles[i].ctx = ctx;
728  }
729  if (!delay_alloc) {
730  CheckAndAllocData(storage_shape, dtype);
731  }
732  }
733 
734  Chunk(const NDArrayStorageType storage_type_, const TBlob &data,
735  const std::vector<TBlob> &aux_data, int dev_id)
736  : static_data(true), delay_alloc(false), storage_type(storage_type_) {
737  using namespace mshadow;
738  CHECK_NE(storage_type, kDefaultStorage);
739  // init var
740  var = Engine::Get()->NewVariable();
741  // init ctx
742  if (data.dev_mask() == cpu::kDevMask) {
743  ctx = Context::CPU();
744  } else {
745  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
746  ctx = Context::GPU(dev_id);
747  }
748  // init shandle
749  shandle.ctx = ctx;
750  shandle.dptr = data.dptr_;
751  shandle.size = data.shape_.Size() * mshadow_sizeof(data.type_flag_);
752  storage_shape = data.shape_;
753  // init aux handles
754  for (const auto &aux : aux_data) {
755  Storage::Handle aux_handle;
756  aux_handle.ctx = ctx;
757  aux_handle.dptr = aux.dptr_;
758  aux_handle.size = aux.shape_.Size() * mshadow_sizeof(aux.type_flag_);
759  aux_handles.push_back(aux_handle);
760  aux_types.emplace_back(aux.type_flag_);
761  aux_shapes.emplace_back(aux.shape_);
762  }
763  }
764 
766  inline void set_aux_shape(const size_t i, const TShape& shape) {
767  aux_shapes[i] = shape;
768  if (storage_shape.ndim() > 0) {
770  storage_shape[0] = shape[0];
771  } else if (storage_type == kCSRStorage && i == csr::kIdx) {
772  storage_shape[0] = shape[0];
773  }
774  }
775  }
776 
778  inline void CheckAndAlloc(void) {
779  if (delay_alloc) {
780  shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
781  delay_alloc = false;
782  }
783  }
784 
786  // size is the number of bytes
787  void CheckAndAlloc(uint64_t dbytes) {
788  CHECK_EQ(kDefaultStorage, storage_type)
789  << "CheckAndAlloc(dbytes) is not intended for kDefaultStorage";
790  if (delay_alloc) {
791  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
792  delay_alloc = false;
793  } else if (shandle.size < dbytes) {
794  // free storage if necessary and alloc again
795  if (shandle.size > 0) Storage::Get()->Free(shandle);
796  // init storage
797  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
798  }
799  }
800 
801  inline void CheckAndAlloc(const TShape &shape, const std::vector<TShape> &aux_shapes,
802  int dtype) {
803  // calculate size, perform allocation
805  // For row sparse, aux_shape indicates the number of rows to allocate
806  auto aux_shape = aux_shapes[rowsparse::kIdx];
808  TShape storage_shape(shape);
809  storage_shape[0] = aux_shape[0];
810  CheckAndAllocData(storage_shape, dtype);
811  } else if (kCSRStorage == storage_type) {
814  CheckAndAllocData(aux_shapes[csr::kIdx], dtype);
815  } else {
816  LOG(FATAL) << "Storage type " << storage_type << " not implemented for CheckAndAlloc";
817  }
818  }
819  // create storage handle for data based on shape and dtype, assuming ctx is set
820  // storage shape is also updated
821  // if data is already allocated, try reuse the storage. Otherwise, free the current one
822  // and allocate new storage
823  inline void CheckAndAllocData(const TShape &shape, int dtype) {
824  CHECK_NE(aux_shapes.size(), 0) << "data is expected to be allocated after aux_data";
825  auto dbytes = shape.Size() * mshadow::mshadow_sizeof(dtype);
826  if (shandle.size < dbytes) {
827  // free storage if necessary and alloc again
828  if (shandle.size > 0) Storage::Get()->Free(shandle);
829  // init storage
830  shandle = Storage::Get()->Alloc(dbytes, ctx);
831  }
832  // init shape
833  storage_shape = shape;
834  // delay_alloc is only set when data storage handle is present
835  delay_alloc = false;
836  }
837  // create storage handle for aux data based on shape
838  // this function assumes ctx, aux shapes and aux types are set
839  // aux shape is also updated
840  // if aux data is already allocated, try reuse the storage. Otherwise, free the current one
841  // and allocate new storage
842  inline void CheckAndAllocAuxData(size_t i, const TShape &shape) {
843  CHECK_EQ(shape.ndim(), 1) << "shape must be 1D in CheckAndAllocAuxData";
845  << "storage type cannot be kUndefinedStorage in CheckAndAllocAuxData";
846  CHECK_NE(storage_type, kDefaultStorage)
847  << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
848  if (aux_handles.size() <= i) {
849  aux_handles.resize(i + 1);
850  }
851  size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
852  if (aux_handles[i].size < aux_bytes) {
853  // free storage if necessary and alloc again
854  if (aux_handles[i].size > 0) Storage::Get()->Free(aux_handles[i]);
855  // init aux storage
856  aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
857  }
858  // init shape
859  set_aux_shape(i, shape);
860  }
862  ~Chunk() {
863  bool skip_free = static_data || delay_alloc;
864  Storage::Handle h = this->shandle;
865  std::vector<Storage::Handle> aux_h = this->aux_handles;
866  Engine::Get()->DeleteVariable([h, aux_h, skip_free](RunContext s) {
867  if (skip_free == false) {
868  Storage::Get()->Free(h);
869  for (size_t i = 0; i < aux_h.size(); i++) {
870  if (aux_h[i].size > 0) Storage::Get()->Free(aux_h[i]);
871  }
872  }
873  }, shandle.ctx, var);
874  }
875  }; // struct Chunk
876 
877  void SetTBlob() const {
878  CHECK(ptr_ != nullptr);
879  TShape shape = shape_;
880  char *dptr = static_cast<char*>(ptr_->shandle.dptr);
881  auto stype = storage_type();
882  if (stype == kDefaultStorage) {
883  dptr += byte_offset_;
884  } else if (stype == kCSRStorage || stype == kRowSparseStorage) {
885  shape = storage_shape();
886  } else {
887  LOG(FATAL) << "unknown storage type " << stype;
888  }
889  tblob_.dptr_ = dptr;
890  tblob_.shape_ = shape;
891  tblob_.type_flag_ = dtype_;
892  tblob_.SetDLTensor(ptr_->shandle.ctx.dev_mask(), ptr_->shandle.ctx.dev_id);
893 #if MKL_EXPERIMENTAL == 1
894  tblob_.Mkl_mem_ = Mkl_mem_;
895 #endif
896  }
897 
898 #if MKL_EXPERIMENTAL == 1
899  std::shared_ptr<MKLMemHolder> Mkl_mem_;
900 #endif
901 
902  std::shared_ptr<Chunk> ptr_{nullptr};
904  TShape shape_;
906  size_t byte_offset_ = 0;
908  int dtype_ = -1;
910  NDArrayStorageType storage_type_ = kUndefinedStorage;
912  nnvm::NodeEntry entry_;
920  mutable TBlob tblob_;
921 }; // class NDArray
922 
926 size_t num_aux_data(NDArrayStorageType stype);
927 
939 void CopyFromTo(const NDArray &from, const NDArray *to, int priority = 0);
940 
952 void CopyFromTo(const NDArray &from, const NDArray& to, int priority = 0);
953 
960 void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);
961 
968 NDArray operator+(const NDArray &lhs, const NDArray &rhs);
975 NDArray operator+(const NDArray &lhs, const real_t &rhs);
982 NDArray operator-(const NDArray &lhs, const NDArray &rhs);
989 NDArray operator-(const NDArray &lhs, const real_t &rhs);
996 NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
1003 NDArray operator*(const NDArray &lhs, const real_t &rhs);
1010 NDArray operator/(const NDArray &lhs, const NDArray &rhs);
1017 NDArray operator/(const NDArray &lhs, const real_t &rhs);
1018 
1023 void RandomSeed(uint32_t seed);
1030 void SampleUniform(real_t begin, real_t end, NDArray *out);
1037 void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
1044 void SampleGamma(real_t alpha, real_t beta, NDArray *out);
1050 void SampleExponential(real_t lambda, NDArray *out);
1056 void SamplePoisson(real_t lambda, NDArray *out);
1063 void SampleNegBinomial(int32_t k, real_t p, NDArray *out);
1070 void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out);
1071 
1072 
1073 //--------------------------------------------------------------
1074 // The following part are API Registration of NDArray functions.
1075 //--------------------------------------------------------------
1076 
1078 typedef std::function<void (NDArray **used_vars,
1079  real_t *scalars,
1080  NDArray **mutate_vars,
1081  int num_params,
1082  char **param_keys,
1083  char **param_vals)> NDArrayAPIFunction;
1099 };
1102  : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
1103  NDArrayAPIFunction> {
1105  unsigned num_use_vars;
1109  unsigned num_scalars;
1116  : num_use_vars(0),
1117  num_mutate_vars(0),
1118  num_scalars(0),
1119  type_mask(0) {}
1126  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
1127  NDArray *out)) {
1128  body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1129  int num_params, char **param_keys, char **param_vals) {
1130  (*fsetvalue)(s[0], mutate_vars[0]);
1131  };
1132  num_mutate_vars = 1; num_scalars = 1;
1133  this->add_argument("src", "real_t", "Source input to the function.");
1134  return *this;
1135  }
1142  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
1143  const NDArray &mhs,
1144  const NDArray &rhs,
1145  NDArray *out)) {
1146  body = [fternary](NDArray **used_vars,
1147  real_t *s, NDArray **mutate_vars,
1148  int num_params, char **param_keys, char **param_vals) {
1149  (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
1150  };
1151  num_use_vars = 3; num_mutate_vars = 1;
1153  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1154  this->add_argument("mhs", "NDArray", "Middle operand to the function.");
1155  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1156  return *this;
1157  }
1164  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
1165  const NDArray &rhs,
1166  NDArray *out)) {
1167  body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1168  int num_params, char **param_keys, char **param_vals) {
1169  (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
1170  };
1171  num_use_vars = 2; num_mutate_vars = 1;
1173  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1174  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1175  return *this;
1176  }
1183  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
1184  const real_t &rhs,
1185  NDArray *out)) {
1186  body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1187  int num_params, char **param_keys, char **param_vals) {
1188  (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
1189  };
1190  num_use_vars = 1; num_mutate_vars = 1; num_scalars = 1;
1192  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1193  this->add_argument("rhs", "real_t", "Right operand to the function.");
1194  return *this;
1195  }
1202  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
1203  NDArray *out)) {
1204  body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1205  int num_params, char **param_keys, char **param_vals) {
1206  (*funary)(*used_vars[0], mutate_vars[0]);
1207  };
1208  num_use_vars = 1; num_mutate_vars = 1;
1210  this->add_argument("src", "NDArray", "Source input to the function.");
1211  return *this;
1212  }
1220  void (*fgeneric)(NDArray **used_vars,
1221  real_t *s,
1222  NDArray **mutate_vars,
1223  const std::map<std::string, std::string>& param)) {
1224  body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1225  int num_params, char **param_keys, char **param_vals) {
1226  std::map<std::string, std::string> param;
1227  for (int i = 0; i < num_params; ++i) {
1228  param[param_keys[i]] = param_vals[i];
1229  }
1230  fgeneric(used_vars, s, mutate_vars, param);
1231  };
1232  return *this;
1233  }
1239  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
1240  num_use_vars = n; return *this;
1241  }
1248  num_mutate_vars = n; return *this;
1249  }
1255  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
1256  num_scalars = n; return *this;
1257  }
1263  inline NDArrayFunctionReg &set_type_mask(int tmask) {
1264  type_mask = tmask; return *this;
1265  }
1266 }; // NDArrayFunctionReg
1267 
1279 #define MXNET_REGISTER_NDARRAY_FUN(name) \
1280  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)
1281 
1282 } // namespace mxnet
1283 
1284 namespace dmlc {
1286 DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
1287 } // namespace dmlc
1288 #endif // MXNET_NDARRAY_H_
Definition: ndarray.h:72
Definition: ndarray.h:61
NDArray ReshapeWithRecord(const TShape &shape)
Get an reshaped NDArray. Supports autograd recording.
NDArrayStorageType
Definition: ndarray.h:59
bool Load(dmlc::Stream *strm)
load the content from binary stream
Definition: ndarray.h:52
NDArrayFunctionReg & set_num_mutate_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:1247
NDArrayFormatErr
Definition: ndarray.h:66
Engine::VarHandle var() const
Definition: ndarray.h:361
void RandomSeed(uint32_t seed)
Seed the random number generator.
NDArrayStorageType storage_type() const
Definition: ndarray.h:297
Engine that schedules all the operations according to dependency.
NDArray & operator/=(const NDArray &src)
elementwise division from current ndarray this mutate the current NDArray
NDArray AtWithRecord(index_t idx)
Index a NDArray.
void SyncCopyFromCPU(const void *data, size_t size) const
Do a synchronize copy from a continugous CPU memory region.
NDArrayFunctionReg()
constructor
Definition: ndarray.h:1115
void ReshapeAndAlloc(const TShape &shape)
Allocate the space if the allocation has been delayed or the requested size is bigger than the availa...
Definition: ndarray.h:585
NDArray operator*(const NDArray &lhs, const NDArray &rhs)
elementwise multiplication
virtual void Free(Handle handle)=0
Free storage.
NDArray Slice(index_t begin, index_t end) const
Slice a NDArray.
NDArrayFunctionReg & set_num_use_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:1239
DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true)
traits
mshadow::default_real_t real_t
data type that will be used to store ndarray
Definition: base.h:135
static Context GPU(int32_t dev_id=-1)
int type_mask
information on how function should be called from API
Definition: ndarray.h:1111
NDArrayFunctionReg & set_function(void(*funary)(const NDArray &src, NDArray *out))
set the function body to a unary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:1202
NDArray Detach() const
Return a copy of this NDArray without autograd history.
Definition: ndarray.h:560
NDArrayFunctionReg & set_num_scalars(unsigned n)
set the number of scalar arguments
Definition: ndarray.h:1255
nnvm::TShape TShape
Shape data structure used to record shape information.
Definition: base.h:137
void SyncCopyFromNDArray(const NDArray &src, int i=-1, int j=-1)
Copy from src.data()/aux_data(i) to this->data()/aux_data(j)
NDArray & operator+=(const NDArray &src)
elementwise add to current space this mutate the current NDArray
Definition: ndarray.h:70
unsigned num_mutate_vars
number of variable mutated by this function
Definition: ndarray.h:1107
execution time context. The information needed in runtime for actual execution.
Definition: base.h:253
NDArray aux_ndarray(size_t i) const
Generate a deep copy of aux_data(i) returned as a default storage type NDArray.
NDArrayFunctionReg & set_function(void(*fscalar)(const NDArray &lhs, const real_t &rhs, NDArray *out))
set the function body to a binary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:1183
void SyncCopyToCPU(void *data, size_t size) const
Do a synchronize copy to a continugous CPU memory region.
Definition: ndarray.h:63
nnvm::Symbol get_autograd_symbol() const
Storage::Handle storage_handle() const
get storage handle
Definition: ndarray.h:330
NDArray()
default constructor
Definition: ndarray.h:82
unsigned num_use_vars
number of variable used by this function
Definition: ndarray.h:1105
bool LegacyLoad(dmlc::Stream *strm, const uint32_t magic)
load ndarrays before supporting sparse ndarrays
NDArrayFunctionReg & set_function(void(*fternary)(const NDArray &lhs, const NDArray &mhs, const NDArray &rhs, NDArray *out))
set the function body to a ternary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:1142
Definition: ndarray.h:60
NDArray Reshape(const TShape &shape) const
Get an reshaped NDArray.
bool fresh_out_grad() const
NDArray SliceWithRecord(index_t begin, index_t end)
Slice a NDArray. Supports recording with autograd.
RowSparseAuxType
Definition: ndarray.h:56
Definition: ndarray.h:68
bool is_none() const
Definition: ndarray.h:301
all the scalar should go before use_vars
Definition: ndarray.h:1089
void SampleExponential(real_t lambda, NDArray *out)
Sample exponential distribution for each elements of out.
NDArray data_ndarray() const
Generate a deep copy of data() returned as a default storage type NDArray.
virtual VarHandle NewVariable()=0
Allocate a new variable, the variable can then be used to schedule the operation concurrently via dep...
Definition: ndarray.h:56
NDArray & operator=(real_t scalar)
set all the elements in ndarray to be scalar
whether this function allows the handles in the target to be empty NDArray that are not yet initializ...
Definition: ndarray.h:1098
Definition: ndarray.h:71
static Storage * Get()
virtual void WaitForVar(VarHandle var)=0
Wait for a variable.
NDArray & operator-=(const NDArray &src)
elementwise subtract from current ndarray this mutate the current NDArray
Context ctx() const
Definition: ndarray.h:282
void CopyFromTo(const NDArray &from, const NDArray *to, int priority=0)
issue an copy operation from one NDArray to another the two ndarray can sit on different devices this...
CSRAuxType
Definition: ndarray.h:52
void SampleGaussian(real_t mu, real_t sigma, NDArray *out)
Sample gaussian distribution for each elements of out.
Definition: ndarray.h:52
Storage manager across multiple devices.
void WaitToRead() const
Block until all the pending write operations with respect to current NDArray are finished, and read can be performed.
Definition: ndarray.h:340
int dtype() const
Definition: ndarray.h:289
bool storage_initialized() const
Definition: ndarray.h:309
Storage handle.
Definition: storage.h:41
static Context CPUShared(int32_t dev_id=0)
Definition: ndarray.h:62
virtual void DeleteVariable(SyncFn delete_fn, Context exec_ctx, VarHandle var)=0
Schedule the deletion of a variable.
void CheckAndAllocData(const TShape &storage_shape) const
Definition: ndarray.h:601
size_t num_aux_data(NDArrayStorageType stype)
NDArrayFunctionReg & set_type_mask(int tmask)
set type mask
Definition: ndarray.h:1263
engine::VarHandle VarHandle
Variable pointer.
Definition: engine.h:105
virtual void PushAsync(AsyncFn exec_fun, Context exec_ctx, std::vector< VarHandle > const &const_vars, std::vector< VarHandle > const &mutable_vars, FnProperty prop=FnProperty::kNormal, int priority=0, const char *opr_name=nullptr)=0
Push an asynchronous operation to the engine.
void WaitToWrite() const
Block until all the pending read/write operations with respect to current NDArray are finished...
Definition: ndarray.h:348
Handle Alloc(size_t size, Context ctx)
Allocate a new contiguous memory for a given size.
Definition: storage.h:66
NDArray operator-(const NDArray &lhs, const NDArray &rhs)
elementwise subtraction
Definition: ndarray.h:69
NDArrayFunctionReg & set_function(void(*fsetvalue)(const real_t &rhs, NDArray *out))
set the function body to a NDArray setvalue function this will also auto set the parameters correctly...
Definition: ndarray.h:1126
NDArray operator+(const NDArray &lhs, const NDArray &rhs)
elementwise add
void SampleUniform(real_t begin, real_t end, NDArray *out)
Sample uniform distribution for each elements of out.
Registry entry for NDArrayFunction.
Definition: ndarray.h:1101
NDArray At(index_t idx) const
Index a NDArray.
entry_({nullptr, 0, 0})
Definition: ndarray.h:110
NDArrayFunctionReg & set_function(void(*fbinary)(const NDArray &lhs, const NDArray &rhs, NDArray *out))
set the function body to a binary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:1164
static Context CPU(int32_t dev_id=0)
runtime functions for NDArray
Definition: imperative.h:56
int aux_type(size_t i) const
Definition: ndarray.h:292
NDArray Copy(Context ctx) const
return a new copy this NDArray
OnComplete Callback to the engine, called by AsyncFn when action completes.
Definition: engine.h:56
configuation of mxnet as well as basic data structure.
NDArray & operator*=(const NDArray &src)
elementwise multiplication to current ndarray this mutate the current NDArray
all the use_vars should go before scalar
Definition: ndarray.h:1087
NDArray AsArray(const TShape &shape, int dtype) const
Create a NDArray that shares memory with current one The new array must have smaller memory size than...
Definition: ndarray.h:528
void CheckAndAlloc(const std::vector< TShape > &aux_shapes) const
Definition: ndarray.h:596
unsigned num_scalars
number of scalars used by this function
Definition: ndarray.h:1109
static Engine * Get()
void Save(dmlc::Stream *strm) const
save the content into binary stream
void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const
Definition: ndarray.h:606
NDArray(const NDArrayStorageType stype, const TShape &shape, const TBlob &data, const std::vector< TBlob > &aux_data, int dev_id)
constructing a static NDArray of non-default storage that shares data with TBlob Use with caution: al...
Definition: ndarray.h:183
Definition: ndarray.h:67
void CheckAndAlloc() const
Allocate the space if it is delayed allocated. This is an internal function used by system that norma...
Definition: ndarray.h:571
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:133
TBlob aux_data(size_t i) const
Definition: ndarray.h:263
void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out)
Sample generalized negative binomial distribution for each elements of out.
Context information about the execution environment.
Definition: base.h:142
void SamplePoisson(real_t lambda, NDArray *out)
Sample Poisson distribution for each elements of out.
ndarray interface
Definition: ndarray.h:79
NDArray(const TBlob &data, int dev_id)
constructing a static NDArray that shares data with TBlob Use with caution: allocate ONLY ONE NDArray...
Definition: ndarray.h:156
void ElementwiseSum(const std::vector< NDArray > &source, NDArray *out, int priority=0)
Perform elementwise sum over each data from source, store result into out.
std::function< void(NDArray **used_vars, real_t *scalars, NDArray **mutate_vars, int num_params, char **param_keys, char **param_vals)> NDArrayAPIFunction
definition of NDArray function
Definition: ndarray.h:1083
void SampleNegBinomial(int32_t k, real_t p, NDArray *out)
Sample negative binomial distribution for each elements of out.
void set_fresh_out_grad(bool state) const
NDArrayFunctionReg & set_function(void(*fgeneric)(NDArray **used_vars, real_t *s, NDArray **mutate_vars, const std::map< std::string, std::string > &param))
set the function body to a unary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:1219
void SyncCheckFormat(const bool full_check) const
check whether the NDArray format is valid
tensor blob class that can be used to hold tensor of any dimension, any device and any data type...
Definition: tensor_blob.h:59
void SampleGamma(real_t alpha, real_t beta, NDArray *out)
Sample gamma distribution for each elements of out.
NDArray(const TShape &shape, Context ctx, bool delay_alloc=false, int dtype=mshadow::default_type_flag)
constructs a new dynamic NDArray
Definition: ndarray.h:94
NDArray operator/(const NDArray &lhs, const NDArray &rhs)
elementwise division
NDArrayFunctionTypeMask
mask information on how functions can be exposed
Definition: ndarray.h:1085