mxnet
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 {
251  if (storage_type() == kDefaultStorage) CheckAndAlloc();
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;
312  inline bool storage_initialized() const {
313  if (is_none()) return false;
314  auto stype = storage_type();
315  CHECK_NE(stype, kDefaultStorage)
316  << "storage_initialized() is not intended for kDefaultStorage.";
317  if (stype == kRowSparseStorage) {
318  CHECK_EQ(aux_shape(rowsparse::kIdx)[0], storage_shape()[0])
319  << "inconsistent storage shape " << storage_shape()
320  << " vs. aux shape " << aux_shape(rowsparse::kIdx);
321  return aux_shape(0).Size() != 0;
322  } else if (stype == kCSRStorage) {
323  CHECK_EQ(aux_shape(csr::kIdx)[0], storage_shape()[0])
324  << "inconsistent storage shape " << storage_shape()
325  << " vs. aux shape " << aux_shape(csr::kIdx);
326  return aux_shape(0).Size() != 0;
327  } else {
328  LOG(FATAL) << "Unknown storage type";
329  }
330  return true;
331  }
334  CHECK(!is_none());
335  CHECK_EQ(storage_type(), kDefaultStorage);
336  CheckAndAlloc();
337  return ptr_->shandle;
338  }
343  inline void WaitToRead() const {
344  if (is_none()) return;
345  Engine::Get()->WaitForVar(ptr_->var);
346  }
351  inline void WaitToWrite() const {
352  if (is_none()) return;
358  [](RunContext, Engine::CallbackOnComplete on_complete) {
359  on_complete();
360  }, Context{}, {}, {ptr_->var});
361  Engine::Get()->WaitForVar(ptr_->var);
362  }
364  inline Engine::VarHandle var() const {
365  return ptr_->var;
366  }
371  void Save(dmlc::Stream *strm) const;
377  bool LegacyLoad(dmlc::Stream *strm, const uint32_t magic);
383  bool Load(dmlc::Stream *strm);
389  NDArray &operator=(real_t scalar);
396  NDArray &operator+=(const NDArray &src);
403  NDArray &operator+=(const real_t &src);
410  NDArray &operator-=(const NDArray &src);
417  NDArray &operator-=(const real_t &src);
424  NDArray &operator*=(const NDArray &src);
431  NDArray &operator*=(const real_t &src);
438  NDArray &operator/=(const NDArray &src);
445  NDArray &operator/=(const real_t &src);
451  NDArray Copy(Context ctx) const;
462  void SyncCopyFromCPU(const void *data, size_t size) const;
463 
467  void SyncCopyFromNDArray(const NDArray &src, int i = -1, int j = -1);
468 
479  void SyncCopyToCPU(void *data, size_t size) const;
485  void SyncCheckFormat(const bool full_check) const;
492  NDArray Slice(index_t begin, index_t end) const;
499  NDArray SliceWithRecord(index_t begin, index_t end);
505  NDArray At(index_t idx) const;
511  NDArray AtWithRecord(index_t idx);
516  NDArray aux_ndarray(size_t i) const;
517 
522  NDArray data_ndarray() const;
523 
531  inline NDArray AsArray(const TShape &shape, int dtype) const {
532  CHECK_EQ(storage_type(), kDefaultStorage)
533  << "AsArray is intended only for kDefaultStorage.";
534  CHECK_GE(ptr_->shandle.size,
535  shape.Size() * mshadow::mshadow_sizeof(dtype))
536  << "NDArray.AsArray: target memory size is bigger";
537 #if MKL_EXPERIMENTAL == 1
538  if (Mkl_mem_ != nullptr) {
539  // convert prv to cpu
540  Mkl_mem_->check_and_prv_to_cpu(ptr_->shandle.dptr);
541  }
542 #endif
543  NDArray ret = *this;
544  ret.shape_ = shape;
545  ret.dtype_ = dtype;
546  return ret;
547  }
553  NDArray Reshape(const TShape &shape) const;
559  NDArray ReshapeWithRecord(const TShape &shape);
563  NDArray Detach() const {
564  NDArray ret(*this);
565  ret.entry_ = nnvm::NodeEntry{nullptr, 0, 0};
566  return ret;
567  }
568 
569  nnvm::Symbol get_autograd_symbol() const;
574  inline void CheckAndAlloc() const {
575  CHECK_EQ(storage_type(), kDefaultStorage);
576  ptr_->CheckAndAlloc();
577  }
578 
588  void ReshapeAndAlloc(const TShape& shape) {
589  CHECK_EQ(storage_type(), kDefaultStorage);
590  CHECK(!is_none());
591  shape_ = shape;
592  ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
593  }
594 
595  /* !
596  * \brief Alloc memory for non-default storage
597  * aux_shape is only known at run time
598  */
599  inline void CheckAndAlloc(const std::vector<TShape> &aux_shapes) const {
600  CHECK_NE(storage_type(), kDefaultStorage)
601  << "CheckAndAlloc(aux_shapes) is not intended for kDefaultStorage";
602  ptr_->CheckAndAlloc(shape_, aux_shapes, dtype_);
603  }
604  inline void CheckAndAllocData(const TShape &storage_shape) const {
605  CHECK_NE(storage_type(), kDefaultStorage)
606  << "CheckAndAllocData is not intended for kDefaultStorage";
607  ptr_->CheckAndAllocData(storage_shape, dtype_);
608  }
609  inline void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const {
610  CHECK_NE(storage_type(), kDefaultStorage)
611  << "CheckAndAllocAuxData is not intended for kDefaultStorage";
612  ptr_->CheckAndAllocAuxData(i, aux_shape);
613  }
620  static void Save(dmlc::Stream* fo,
621  const std::vector<NDArray>& data,
622  const std::vector<std::string>& names);
629  static void Load(dmlc::Stream* fi,
630  std::vector<NDArray>* data,
631  std::vector<std::string>* keys);
632 
633  private:
634  friend class Imperative;
636  // shandle is used to store the actual values in the NDArray
637  // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
638  struct Chunk {
642  Storage::Handle shandle;
647  std::vector<Storage::Handle> aux_handles;
649  Engine::VarHandle var;
655  bool static_data;
658  bool delay_alloc;
659  // the type of the storage. The storage_type is never kUndefinedStorage once the chunk
660  // is constructed.
661  NDArrayStorageType storage_type = kDefaultStorage;
663  std::vector<int> aux_types;
664  // context of data
665  Context ctx;
666  // The shape of the chunk data.
667  // This might not be the same shape as the NDArray, since the storage may be sparse.
668  // The default value for storage_shape is {0} when an empty non-default NDArray is created.
669  TShape storage_shape;
670  // The shape of aux data. The default value for the shape depends on the type of storage.
671  // If aux_shapes[i].Size() is zero, aux data i is empty.
672  std::vector<TShape> aux_shapes;
673 
675  Chunk() : static_data(true), delay_alloc(false) {}
676 
678  Chunk(TShape shape, Context ctx_, bool delay_alloc_, int dtype)
679  : static_data(false), delay_alloc(true), ctx(ctx_) {
680  auto size = shape.Size();
681  storage_shape = shape;
682  var = Engine::Get()->NewVariable();
683  shandle.size = size * mshadow::mshadow_sizeof(dtype);
684  shandle.ctx = ctx_;
685  if (!delay_alloc_) this->CheckAndAlloc();
686  }
687 
688  Chunk(const TBlob &data, int dev_id)
689  : static_data(true), delay_alloc(false) {
690  CHECK(storage_type == kDefaultStorage);
691  var = Engine::Get()->NewVariable();
692  if (data.dev_mask() == cpu::kDevMask) {
693  ctx = Context::CPU();
694  } else {
695  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
696  ctx = Context::GPU(dev_id);
697  }
698  // init shandle
699  shandle.ctx = ctx;
700  shandle.dptr = data.dptr_;
701  shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
702  storage_shape = data.shape_;
703  }
704 
705  Chunk(int shared_pid, int shared_id, const TShape& shape, int dtype)
706  : static_data(false), delay_alloc(false) {
707  var = Engine::Get()->NewVariable();
708  ctx = Context::CPUShared(0);
709  shandle.size = shape.Size() * mshadow::mshadow_sizeof(dtype);;
710  shandle.ctx = ctx;
711  shandle.shared_pid = shared_pid;
712  shandle.shared_id = shared_id;
713  Storage::Get()->Alloc(&shandle);
714  storage_shape = shape;
715  }
716  // Constructor for a non-default storage chunk
717  Chunk(NDArrayStorageType storage_type_, const TShape &storage_shape_, Context ctx_,
718  bool delay_alloc_, int dtype, const std::vector<int> &aux_types_,
719  const std::vector<TShape> &aux_shapes_)
720  : static_data(false), delay_alloc(delay_alloc_), storage_type(storage_type_),
721  aux_types(aux_types_), ctx(ctx_), storage_shape(storage_shape_),
722  aux_shapes(aux_shapes_) {
723  shandle.ctx = ctx;
724  var = Engine::Get()->NewVariable();
725  // aux_handles always reflect the correct number of aux data
726  for (size_t i = 0; i < aux_shapes.size(); i++) {
727  CheckAndAllocAuxData(i, aux_shapes[i]);
728  // this line is needed in case when aux_shapes[i].Size() = 0
729  // aux_handles[i] will not be updated and take only default value.
730  aux_handles[i].ctx = ctx;
731  }
732  if (!delay_alloc) {
733  CheckAndAllocData(storage_shape, dtype);
734  }
735  }
736 
737  Chunk(const NDArrayStorageType storage_type_, const TBlob &data,
738  const std::vector<TBlob> &aux_data, int dev_id)
739  : static_data(true), delay_alloc(false), storage_type(storage_type_) {
740  using namespace mshadow;
741  CHECK_NE(storage_type, kDefaultStorage);
742  // init var
743  var = Engine::Get()->NewVariable();
744  // init ctx
745  if (data.dev_mask() == cpu::kDevMask) {
746  ctx = Context::CPU();
747  } else {
748  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
749  ctx = Context::GPU(dev_id);
750  }
751  // init shandle
752  shandle.ctx = ctx;
753  shandle.dptr = data.dptr_;
754  shandle.size = data.shape_.Size() * mshadow_sizeof(data.type_flag_);
755  storage_shape = data.shape_;
756  // init aux handles
757  for (const auto &aux : aux_data) {
758  Storage::Handle aux_handle;
759  aux_handle.ctx = ctx;
760  aux_handle.dptr = aux.dptr_;
761  aux_handle.size = aux.shape_.Size() * mshadow_sizeof(aux.type_flag_);
762  aux_handles.push_back(aux_handle);
763  aux_types.emplace_back(aux.type_flag_);
764  aux_shapes.emplace_back(aux.shape_);
765  }
766  }
767 
769  inline void set_aux_shape(const size_t i, const TShape& shape) {
770  aux_shapes[i] = shape;
771  if (storage_shape.ndim() > 0) {
772  if (storage_type == kRowSparseStorage && i == rowsparse::kIdx) {
773  storage_shape[0] = shape[0];
774  } else if (storage_type == kCSRStorage && i == csr::kIdx) {
775  storage_shape[0] = shape[0];
776  }
777  }
778  }
779 
781  inline void CheckAndAlloc(void) {
782  if (delay_alloc) {
783  shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
784  delay_alloc = false;
785  }
786  }
787 
789  // size is the number of bytes
790  void CheckAndAlloc(uint64_t dbytes) {
791  CHECK_EQ(kDefaultStorage, storage_type)
792  << "CheckAndAlloc(dbytes) is not intended for kDefaultStorage";
793  if (delay_alloc) {
794  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
795  delay_alloc = false;
796  } else if (shandle.size < dbytes) {
797  // free storage if necessary and alloc again
798  if (shandle.size > 0) Storage::Get()->Free(shandle);
799  // init storage
800  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
801  }
802  }
803 
804  inline void CheckAndAlloc(const TShape &shape, const std::vector<TShape> &aux_shapes,
805  int dtype) {
806  // calculate size, perform allocation
807  if (kRowSparseStorage == storage_type) {
808  // For row sparse, aux_shape indicates the number of rows to allocate
809  auto aux_shape = aux_shapes[rowsparse::kIdx];
810  CheckAndAllocAuxData(rowsparse::kIdx, aux_shape);
811  TShape storage_shape(shape);
812  storage_shape[0] = aux_shape[0];
813  CheckAndAllocData(storage_shape, dtype);
814  } else if (kCSRStorage == storage_type) {
815  CheckAndAllocAuxData(csr::kIndPtr, aux_shapes[csr::kIndPtr]);
816  CheckAndAllocAuxData(csr::kIdx, aux_shapes[csr::kIdx]);
817  CheckAndAllocData(aux_shapes[csr::kIdx], dtype);
818  } else {
819  LOG(FATAL) << "Storage type " << storage_type << " not implemented for CheckAndAlloc";
820  }
821  }
822  // create storage handle for data based on shape and dtype, assuming ctx is set
823  // storage shape is also updated
824  // if data is already allocated, try reuse the storage. Otherwise, free the current one
825  // and allocate new storage
826  inline void CheckAndAllocData(const TShape &shape, int dtype) {
827  CHECK_NE(aux_shapes.size(), 0) << "data is expected to be allocated after aux_data";
828  auto dbytes = shape.Size() * mshadow::mshadow_sizeof(dtype);
829  if (shandle.size < dbytes) {
830  // free storage if necessary and alloc again
831  if (shandle.size > 0) Storage::Get()->Free(shandle);
832  // init storage
833  shandle = Storage::Get()->Alloc(dbytes, ctx);
834  }
835  // init shape
836  storage_shape = shape;
837  // delay_alloc is only set when data storage handle is present
838  delay_alloc = false;
839  }
840  // create storage handle for aux data based on shape
841  // this function assumes ctx, aux shapes and aux types are set
842  // aux shape is also updated
843  // if aux data is already allocated, try reuse the storage. Otherwise, free the current one
844  // and allocate new storage
845  inline void CheckAndAllocAuxData(size_t i, const TShape &shape) {
846  CHECK_EQ(shape.ndim(), 1) << "shape must be 1D in CheckAndAllocAuxData";
847  CHECK_NE(storage_type, kUndefinedStorage)
848  << "storage type cannot be kUndefinedStorage in CheckAndAllocAuxData";
849  CHECK_NE(storage_type, kDefaultStorage)
850  << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
851  if (aux_handles.size() <= i) {
852  aux_handles.resize(i + 1);
853  }
854  size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
855  if (aux_handles[i].size < aux_bytes) {
856  // free storage if necessary and alloc again
857  if (aux_handles[i].size > 0) Storage::Get()->Free(aux_handles[i]);
858  // init aux storage
859  aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
860  }
861  // init shape
862  set_aux_shape(i, shape);
863  }
865  ~Chunk() {
866  bool skip_free = static_data || delay_alloc;
867  Storage::Handle h = this->shandle;
868  std::vector<Storage::Handle> aux_h = this->aux_handles;
869  Engine::Get()->DeleteVariable([h, aux_h, skip_free](RunContext s) {
870  if (skip_free == false) {
871  Storage::Get()->Free(h);
872  for (size_t i = 0; i < aux_h.size(); i++) {
873  if (aux_h[i].size > 0) Storage::Get()->Free(aux_h[i]);
874  }
875  }
876  }, shandle.ctx, var);
877  }
878  }; // struct Chunk
879 
880  void SetTBlob() const {
881  CHECK(ptr_ != nullptr);
882  TShape shape = shape_;
883  char *dptr = static_cast<char*>(ptr_->shandle.dptr);
884  auto stype = storage_type();
885  if (stype == kDefaultStorage) {
886  dptr += byte_offset_;
887  } else if (stype == kCSRStorage || stype == kRowSparseStorage) {
888  shape = storage_shape();
889  } else {
890  LOG(FATAL) << "unknown storage type " << stype;
891  }
892  tblob_.dptr_ = dptr;
893  tblob_.shape_ = shape;
894  tblob_.type_flag_ = dtype_;
895  tblob_.SetDLTensor(ptr_->shandle.ctx.dev_mask(), ptr_->shandle.ctx.dev_id);
896 #if MKL_EXPERIMENTAL == 1
897  tblob_.Mkl_mem_ = Mkl_mem_;
898 #endif
899  }
900 
901 #if MKL_EXPERIMENTAL == 1
902  std::shared_ptr<MKLMemHolder> Mkl_mem_;
903 #endif
904 
905  std::shared_ptr<Chunk> ptr_{nullptr};
907  TShape shape_;
909  size_t byte_offset_ = 0;
911  int dtype_ = -1;
913  NDArrayStorageType storage_type_ = kUndefinedStorage;
915  nnvm::NodeEntry entry_;
923  mutable TBlob tblob_;
924 }; // class NDArray
925 
929 size_t num_aux_data(NDArrayStorageType stype);
930 
942 void CopyFromTo(const NDArray &from, const NDArray *to, int priority = 0);
943 
955 void CopyFromTo(const NDArray &from, const NDArray& to, int priority = 0);
956 
963 void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);
964 
971 NDArray operator+(const NDArray &lhs, const NDArray &rhs);
978 NDArray operator+(const NDArray &lhs, const real_t &rhs);
985 NDArray operator-(const NDArray &lhs, const NDArray &rhs);
992 NDArray operator-(const NDArray &lhs, const real_t &rhs);
999 NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
1006 NDArray operator*(const NDArray &lhs, const real_t &rhs);
1013 NDArray operator/(const NDArray &lhs, const NDArray &rhs);
1020 NDArray operator/(const NDArray &lhs, const real_t &rhs);
1021 
1026 void RandomSeed(uint32_t seed);
1033 void SampleUniform(real_t begin, real_t end, NDArray *out);
1040 void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
1047 void SampleGamma(real_t alpha, real_t beta, NDArray *out);
1053 void SampleExponential(real_t lambda, NDArray *out);
1059 void SamplePoisson(real_t lambda, NDArray *out);
1066 void SampleNegBinomial(int32_t k, real_t p, NDArray *out);
1073 void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out);
1074 
1075 
1076 //--------------------------------------------------------------
1077 // The following part are API Registration of NDArray functions.
1078 //--------------------------------------------------------------
1079 
1081 typedef std::function<void (NDArray **used_vars,
1082  real_t *scalars,
1083  NDArray **mutate_vars,
1084  int num_params,
1085  char **param_keys,
1086  char **param_vals)> NDArrayAPIFunction;
1102 };
1105  : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
1106  NDArrayAPIFunction> {
1108  unsigned num_use_vars;
1112  unsigned num_scalars;
1119  : num_use_vars(0),
1120  num_mutate_vars(0),
1121  num_scalars(0),
1122  type_mask(0) {}
1129  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
1130  NDArray *out)) {
1131  body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1132  int num_params, char **param_keys, char **param_vals) {
1133  (*fsetvalue)(s[0], mutate_vars[0]);
1134  };
1135  num_mutate_vars = 1; num_scalars = 1;
1136  this->add_argument("src", "real_t", "Source input to the function.");
1137  return *this;
1138  }
1145  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
1146  const NDArray &mhs,
1147  const NDArray &rhs,
1148  NDArray *out)) {
1149  body = [fternary](NDArray **used_vars,
1150  real_t *s, NDArray **mutate_vars,
1151  int num_params, char **param_keys, char **param_vals) {
1152  (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
1153  };
1154  num_use_vars = 3; num_mutate_vars = 1;
1156  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1157  this->add_argument("mhs", "NDArray", "Middle operand to the function.");
1158  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1159  return *this;
1160  }
1167  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
1168  const NDArray &rhs,
1169  NDArray *out)) {
1170  body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1171  int num_params, char **param_keys, char **param_vals) {
1172  (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
1173  };
1174  num_use_vars = 2; num_mutate_vars = 1;
1176  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1177  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1178  return *this;
1179  }
1186  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
1187  const real_t &rhs,
1188  NDArray *out)) {
1189  body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1190  int num_params, char **param_keys, char **param_vals) {
1191  (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
1192  };
1193  num_use_vars = 1; num_mutate_vars = 1; num_scalars = 1;
1195  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1196  this->add_argument("rhs", "real_t", "Right operand to the function.");
1197  return *this;
1198  }
1205  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
1206  NDArray *out)) {
1207  body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1208  int num_params, char **param_keys, char **param_vals) {
1209  (*funary)(*used_vars[0], mutate_vars[0]);
1210  };
1211  num_use_vars = 1; num_mutate_vars = 1;
1213  this->add_argument("src", "NDArray", "Source input to the function.");
1214  return *this;
1215  }
1223  void (*fgeneric)(NDArray **used_vars,
1224  real_t *s,
1225  NDArray **mutate_vars,
1226  const std::map<std::string, std::string>& param)) {
1227  body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1228  int num_params, char **param_keys, char **param_vals) {
1229  std::map<std::string, std::string> param;
1230  for (int i = 0; i < num_params; ++i) {
1231  param[param_keys[i]] = param_vals[i];
1232  }
1233  fgeneric(used_vars, s, mutate_vars, param);
1234  };
1235  return *this;
1236  }
1242  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
1243  num_use_vars = n; return *this;
1244  }
1251  num_mutate_vars = n; return *this;
1252  }
1258  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
1259  num_scalars = n; return *this;
1260  }
1266  inline NDArrayFunctionReg &set_type_mask(int tmask) {
1267  type_mask = tmask; return *this;
1268  }
1269 }; // NDArrayFunctionReg
1270 
1282 #define MXNET_REGISTER_NDARRAY_FUN(name) \
1283  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)
1284 
1285 } // namespace mxnet
1286 
1287 namespace dmlc {
1289 DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
1290 } // namespace dmlc
1291 #endif // MXNET_NDARRAY_H_
Definition: ndarray.h:72
Definition: ndarray.h:61
NDArrayStorageType
Definition: ndarray.h:59
Definition: ndarray.h:52
NDArrayFunctionReg & set_num_mutate_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:1250
NDArrayFormatErr
Definition: ndarray.h:66
Engine::VarHandle var() const
Definition: ndarray.h:364
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.
TShape shape_
shape of the tensor
Definition: tensor_blob.h:65
const TShape & storage_shape() const
Definition: ndarray.h:204
NDArrayFunctionReg()
constructor
Definition: ndarray.h:1118
namespace of mxnet
Definition: base.h:127
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:588
NDArray operator*(const NDArray &lhs, const NDArray &rhs)
elementwise multiplication
virtual void Free(Handle handle)=0
Free storage.
NDArrayFunctionReg & set_num_use_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:1242
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:1114
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:1205
NDArray Detach() const
Return a copy of this NDArray without autograd history.
Definition: ndarray.h:563
int type_flag_
type flag of the tensor blob
Definition: tensor_blob.h:67
NDArrayFunctionReg & set_num_scalars(unsigned n)
set the number of scalar arguments
Definition: ndarray.h:1258
nnvm::TShape TShape
Shape data structure used to record shape information.
Definition: base.h:137
Definition: ndarray.h:70
unsigned num_mutate_vars
number of variable mutated by this function
Definition: ndarray.h:1110
execution time context. The information needed in runtime for actual execution.
Definition: base.h:253
void * dptr
Pointer to the data.
Definition: storage.h:45
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:1186
Definition: ndarray.h:63
Context ctx
Context information about device and ID.
Definition: storage.h:53
Storage::Handle storage_handle() const
get storage handle
Definition: ndarray.h:333
NDArray()
default constructor
Definition: ndarray.h:82
unsigned num_use_vars
number of variable used by this function
Definition: ndarray.h:1108
int shared_id
Definition: storage.h:58
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:1145
Definition: ndarray.h:60
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:1092
void SampleExponential(real_t lambda, NDArray *out)
Sample exponential distribution for each elements of out.
void * dptr_
pointer to the data
Definition: tensor_blob.h:63
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
whether this function allows the handles in the target to be empty NDArray that are not yet initializ...
Definition: ndarray.h:1101
Definition: ndarray.h:71
static Storage * Get()
const TShape & shape() const
Definition: ndarray.h:196
Definition: ndarray.h:1287
virtual void WaitForVar(VarHandle var)=0
Wait for a variable.
const std::vector< TShape > & aux_shapes() const
Definition: ndarray.h:223
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:343
int dtype() const
Definition: ndarray.h:289
bool storage_initialized() const
Returns true if a sparse ndarray&#39;s aux_data and storage are initialized Throws an exception if the in...
Definition: ndarray.h:312
Storage handle.
Definition: storage.h:41
static Context CPUShared(int32_t dev_id=0)
Definition: ndarray.h:62
void set_aux_shape(size_t index, const TShape &shape) const
For a sparse operation on a csr matrix for example, the size of the column index array is an estimate...
Definition: ndarray.h:243
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:604
size_t num_aux_data(NDArrayStorageType stype)
NDArrayFunctionReg & set_type_mask(int tmask)
set type mask
Definition: ndarray.h:1266
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:351
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
NDArray(const NDArrayStorageType stype, const TShape &shape, Context ctx, bool delay_alloc=true, int dtype=mshadow::default_type_flag, std::vector< int > aux_types={}, std::vector< TShape > aux_shapes={}, TShape storage_shape=TShape(mshadow::Shape1(0)))
constructor for NDArray with storage type
Definition: ndarray.h:105
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:1129
NDArray(int shared_pid, int shared_id, const TShape &shape, int dtype)
create ndarray from shared memory
Definition: ndarray.h:165
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:1104
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:1167
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
OnComplete Callback to the engine, called by AsyncFn when action completes.
Definition: engine.h:56
all the use_vars should go before scalar
Definition: ndarray.h:1090
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:531
void CheckAndAlloc(const std::vector< TShape > &aux_shapes) const
Definition: ndarray.h:599
unsigned num_scalars
number of scalars used by this function
Definition: ndarray.h:1112
static Engine * Get()
const TBlob & data() const
Definition: ndarray.h:250
void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const
Definition: ndarray.h:609
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:574
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:133
size_t size
Size of the storage.
Definition: storage.h:49
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.
const TShape & aux_shape(size_t index) const
get the shape of aux_data(index)
Definition: ndarray.h:216
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
int dev_mask() const
device mask of the corresponding device
Definition: tensor_blob.h:228
Symbol Reshape(const std::string &symbol_name, Symbol data, Shape shape=Shape(), bool reverse=false, Shape target_shape=Shape(), bool keep_highest=false)
Definition: op.h:302
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:1086
void SampleNegBinomial(int32_t k, real_t p, NDArray *out)
Sample negative binomial distribution for each elements of out.
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:1222
int shared_pid
Id for IPC shared memory.
Definition: storage.h:57
tensor blob class that can be used to hold tensor of any dimension, any device and any data type...
Definition: tensor_blob.h:59
const std::vector< int > & aux_types() const
Definition: ndarray.h:230
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:1088