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 <algorithm>
38 #include <memory>
39 #include <algorithm>
40 #if MXNET_USE_MKLDNN == 1
41 #include <mkldnn.hpp>
42 #endif
43 #include "./base.h"
44 #include "./storage.h"
45 #include "./engine.h"
46 // check c++11
47 #if DMLC_USE_CXX11 == 0
48 #error "cxx11 was required for ndarray module"
49 #endif
50 
51 namespace mxnet {
52 // enum for storage types
53 namespace csr {
55 }
56 
57 namespace rowsparse {
59 }
60 
62  kUndefinedStorage = -1, // undefined storage
63  kDefaultStorage, // dense
64  kRowSparseStorage, // row sparse
65  kCSRStorage, // csr
66 };
67 
69  kNormalErr, // normal
70  kCSRShapeErr, // shape mismatch for csr
71  kCSRIndPtrErr, // indptr error for csr
72  kCSRIdxErr, // idx error for csr
73  kRSPShapeErr, // shape mismatch for row sparse
74  kRSPIdxErr, // indices error for row sparse
75 };
76 
77 class MKLDNNMemory;
78 
82 class NDArray {
83  public:
85  NDArray() {
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  }
102  NDArray(const NDArrayStorageType stype, const TShape &shape, Context ctx,
103  bool delay_alloc = true, int dtype = mshadow::default_type_flag,
104  std::vector<int> aux_types = {}, std::vector<TShape> aux_shapes = {},
105  TShape storage_shape = TShape(mshadow::Shape1(0)));
106 
114  NDArray(const TBlob &data, int dev_id)
115  : ptr_(std::make_shared<Chunk>(data, dev_id)), shape_(data.shape_),
116  dtype_(data.type_flag_), storage_type_(kDefaultStorage),
117  entry_({nullptr, 0, 0}) {
118  }
119 
128  NDArray(const TBlob &data, int dev_id, const std::function<void()>& deleter)
129  : ptr_(new Chunk(data, dev_id),
130  [deleter](Chunk *p) {
131  deleter(); // call custom deleter
132  delete p; // delete Chunk object
133  }),
134  shape_(data.shape_),
135  dtype_(data.type_flag_), storage_type_(kDefaultStorage),
136  entry_({nullptr, 0, 0}) {
137  }
138 
140  NDArray(int shared_pid, int shared_id, const TShape& shape, int dtype)
141  : ptr_(std::make_shared<Chunk>(shared_pid, shared_id, shape, dtype)), shape_(shape),
142  dtype_(dtype), storage_type_(kDefaultStorage), entry_({nullptr, 0, 0}) {
143  }
144 
155  NDArray(const NDArrayStorageType stype, const TShape &shape,
156  const TBlob &data, const std::vector<TBlob> &aux_data, int dev_id)
157  : ptr_(std::make_shared<Chunk>(stype, data, aux_data, dev_id)), shape_(shape),
158  dtype_(data.type_flag_), storage_type_(stype), entry_({nullptr, 0, 0}) {
159  }
160 
161  /*
162  * This indicates whether an array is a view of another array (created by
163  * reshape or slice). If an array is a view and the the data is stored in
164  * MKLDNN format, we need to convert the data to the default format when
165  * data in the view is accessed.
166  */
167  inline bool IsView() const {
168  // View only works on the default storage
169  if (storage_type() != kDefaultStorage)
170  return false;
171  // If the array reuses memory, its shape may be different from the storage
172  // shape. However, we shouldn't consider it as a view.
173  if (reuse_)
174  return false;
175  return byte_offset_ > 0 || shape() != ptr_->storage_shape;
176  }
177 
178  /* \brief Check whether the two arrays are the same array */
179  inline bool IsSame(const NDArray& other) const {
180  return ptr_ == other.ptr_ &&
181  shape_ == other.shape_ &&
182  byte_offset_ == other.byte_offset_ &&
183  dtype_ == other.dtype_;
184  }
185 
189  inline const TShape& shape() const {
190  return shape_;
191  }
197  inline const TShape &storage_shape() const {
198  CHECK(ptr_ != nullptr);
199  CHECK_NE(storage_type(), kDefaultStorage)
200  << "storage_shape() is not intended for kDefaultStorage.";
201  return ptr_->storage_shape;
202  }
203 
209  inline const TShape& aux_shape(size_t index) const {
210  CHECK_NE(storage_type(), kDefaultStorage)
211  << "aux_shape() is not intended for kDefaultStorage.";
212  return ptr_->aux_shapes[index];
213  }
214 
215  /* \return the shapes of all aux data */
216  const std::vector<TShape>& aux_shapes() const {
217  CHECK_NE(storage_type(), kDefaultStorage)
218  << "aux_shapes() is not intended for kDefaultStorage.";
219  return ptr_->aux_shapes;
220  }
221 
223  const std::vector<int>& aux_types() const {
224  CHECK_NE(storage_type(), kDefaultStorage)
225  << "aux_types() is not intended for kDefaultStorage.";
226  return ptr_->aux_types;
227  }
228 
236  inline void set_aux_shape(size_t index, const TShape& shape) const {
237  CHECK_NE(storage_type(), kDefaultStorage)
238  << "set_aux_shape() is not intended for kDefaultStorage.";
239  ptr_->set_aux_shape(index, shape);
240  }
241 
245  inline const TBlob& data() const {
246  if (storage_type() == kDefaultStorage) CheckAndAlloc();
247  SetTBlob();
248  return tblob_;
249  }
253  NDArray grad() const;
254 
258  inline TBlob aux_data(size_t i) const {
259  auto stype = storage_type();
260  TBlob res;
261  auto shape = aux_shape(i);
262  auto type = aux_type(i);
263  MSHADOW_TYPE_SWITCH(type, DType, {
264  auto dptr = static_cast<DType*>(ptr_->aux_handles[i].dptr);
265  CHECK(stype == kRowSparseStorage || stype == kCSRStorage)
266  << "Unexpected storage type: " << stype;
267  res = TBlob(dptr, shape, ptr_->aux_handles[i].ctx.dev_mask(), type);
268  });
269  return res;
270  }
274  inline Context ctx() const {
275  CHECK(!is_none());
276  return ptr_->shandle.ctx;
277  }
281  inline int dtype() const {
282  return dtype_;
283  }
284  inline int aux_type(size_t i) const {
285  CHECK(!is_none());
286  return ptr_->aux_types[i];
287  }
288 
290  return storage_type_;
291  }
293  inline bool is_none() const {
294  return ptr_.get() == nullptr;
295  }
297  bool fresh_out_grad() const;
299  void set_fresh_out_grad(bool state) const;
304  inline bool storage_initialized() const {
305  if (is_none()) return false;
306  auto stype = storage_type();
307  CHECK_NE(stype, kDefaultStorage)
308  << "storage_initialized() is not intended for kDefaultStorage.";
309  if (stype == kRowSparseStorage) {
310  CHECK_EQ(aux_shape(rowsparse::kIdx)[0], storage_shape()[0])
311  << "inconsistent storage shape " << storage_shape()
312  << " vs. aux shape " << aux_shape(rowsparse::kIdx);
313  return aux_shape(rowsparse::kIdx).Size() != 0;
314  } else if (stype == kCSRStorage) {
315  CHECK_EQ(aux_shape(csr::kIdx)[0], storage_shape()[0])
316  << "inconsistent storage shape " << storage_shape()
317  << " vs. aux shape " << aux_shape(csr::kIdx);
318  return aux_shape(csr::kIdx).Size() != 0;
319  } else {
320  LOG(FATAL) << "Unknown storage type";
321  }
322  return true;
323  }
326  CHECK(!is_none());
327  CHECK_EQ(storage_type(), kDefaultStorage);
328  CheckAndAlloc();
329  return ptr_->shandle;
330  }
335  inline void WaitToRead() const {
336  if (is_none()) return;
337  Engine::Get()->WaitForVar(ptr_->var);
338  }
343  inline void WaitToWrite() const {
344  if (is_none()) return;
350  [](RunContext, Engine::CallbackOnComplete on_complete) {
351  on_complete();
352  }, Context{}, {}, {ptr_->var});
353  Engine::Get()->WaitForVar(ptr_->var);
354  }
356  inline Engine::VarHandle var() const {
357  return ptr_->var;
358  }
360  inline size_t byte_offset() const {
361  return byte_offset_;
362  }
364  inline size_t version() const {
365  return var()->version();
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  // We can't reuse memory in a view.
538  CHECK(!IsView());
539  NDArray ret = *this;
540  ret.shape_ = shape;
541  ret.dtype_ = dtype;
542  ret.reuse_ = true;
543  return ret;
544  }
545 
551  DLManagedTensor* ToDLPack() const;
552 
564  static NDArray FromDLPack(const DLManagedTensor* tensor);
565 
573  inline void SparseUpdateChunk(const NDArray &arr) const {
574  CHECK(shape_ == arr.shape_) << "ndarray shape is different from the target";
575  CHECK(dtype_ == arr.dtype_) << "ndarray dtype is different from the target";
576  auto stype = arr.storage_type();
577  CHECK(stype == kCSRStorage || stype == kRowSparseStorage)
578  << "Only to be used with CSR and RSP storage types";
579  // swap shandles between src and dst
580  Storage::Handle shandle_dst = arr.ptr_->shandle;
581  arr.ptr_->shandle = ptr_->shandle;
582  ptr_->shandle = shandle_dst;
583 
584  ptr_->storage_shape = arr.ptr_->storage_shape;
585  ptr_->storage_type = arr.ptr_->storage_type;
586  ptr_->ctx = arr.ptr_->ctx;
587 
588  // swap aux_handles between src and dst
589  size_t aux_idx = 0;
590  CHECK(ptr_->aux_handles.size() == arr.ptr_->aux_handles.size())
591  << "ndarray number of aux_handles is different from target";
592  for (auto &aux_handle : arr.ptr_->aux_handles) {
593  Storage::Handle aux_dst = ptr_->aux_handles[aux_idx];
594  ptr_->aux_handles[aux_idx] = aux_handle;
595  aux_handle = aux_dst;
596  aux_idx++;
597  }
598  ptr_->aux_types = arr.ptr_->aux_types;
599  ptr_->aux_shapes = arr.ptr_->aux_shapes;
600  }
601 
607  NDArray Reshape(const TShape &shape) const;
613  NDArray ReshapeWithRecord(const TShape &shape);
617  NDArray Detach() const {
618  NDArray ret(*this);
619  ret.entry_ = nnvm::NodeEntry{nullptr, 0, 0};
620  return ret;
621  }
622 
623  nnvm::Symbol get_autograd_symbol() const;
628  inline void CheckAndAlloc() const {
629  CHECK_EQ(storage_type(), kDefaultStorage);
630  ptr_->CheckAndAlloc();
631  }
632 
642  void ReshapeAndAlloc(const TShape& shape) {
643  CHECK_EQ(storage_type(), kDefaultStorage);
644  CHECK(!is_none());
645  shape_ = shape;
646  ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
647  }
648 
649  /* !
650  * \brief Alloc memory for non-default storage
651  * aux_shape is only known at run time
652  */
653  inline void CheckAndAlloc(const std::vector<TShape> &aux_shapes) const {
654  CHECK_NE(storage_type(), kDefaultStorage)
655  << "CheckAndAlloc(aux_shapes) is not intended for kDefaultStorage";
656  ptr_->CheckAndAlloc(shape_, aux_shapes, dtype_);
657  }
658  inline void CheckAndAllocData(const TShape &storage_shape) const {
659  CHECK_NE(storage_type(), kDefaultStorage)
660  << "CheckAndAllocData is not intended for kDefaultStorage";
661  ptr_->CheckAndAllocData(storage_shape, dtype_);
662  }
663  inline void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const {
664  CHECK_NE(storage_type(), kDefaultStorage)
665  << "CheckAndAllocAuxData is not intended for kDefaultStorage";
666  ptr_->CheckAndAllocAuxData(i, aux_shape);
667  }
668 
669 #if MXNET_USE_MKLDNN == 1
670  /*
671  * Create NDArray from mkldnn memory.
672  * mkldnn_mem The mkldnn memory to be managed.
673  * static_data If true, mkldnn memory won't be freed on destruction.
674  */
675  explicit NDArray(const mkldnn::memory *mkldnn_mem, bool static_data = true);
676  /*
677  * Test if the data is stored in one of special MKLDNN format.
678  */
679  bool IsMKLDNNData() const {
680  return ptr_->IsMKLDNN();
681  }
682  /*
683  * Test if the data is stored in one of default MXNet formats.
684  */
685  bool IsDefaultData() const {
686  return ptr_->IsDefault();
687  }
688  /*
689  * All functions below return a raw pointer to mkldnn memory. Actually there
690  * is a shared pointer that hold the memory either in NDArray or in MKLDNN
691  * stream. As long as we call these functions inside an operator, the return
692  * memory is always valid.
693  */
694 
695  /*
696  * This function returns mkldnn::memory with the default primitive_desc.
697  */
698  const mkldnn::memory *GetMKLDNNData() const;
699  /*
700  * This function returns mkldnn::memory with the given primitive_desc
701  * as long as the array size meets the required size in the given primitive_desc.
702  */
703  const mkldnn::memory *GetMKLDNNData(
704  const mkldnn::memory::primitive_desc &desc) const;
705  /*
706  * This function returns mkldnn::memory with the given primitive_desc.
707  * The returned mkldnn::memory will have the same physical layout as
708  * the given primitive_desc.
709  */
710  const mkldnn::memory *GetMKLDNNDataReorder(
711  const mkldnn::memory::primitive_desc &desc) const;
712 
713  /*
714  * This function copies data from mkldnn memory.
715  */
716  void CopyFrom(const mkldnn::memory &mem);
717  /*
718  * This function allocates memory for array and creates mkldnn memory
719  * with the specified format.
720  */
721  mkldnn::memory *CreateMKLDNNData(
722  const mkldnn::memory::primitive_desc &desc);
723 
724  /*
725  * These are the async version of the methods above.
726  * It changes the layout of this NDArray, but it happens after all accesses to
727  * the array are complete.
728  */
729  void Reorder2DefaultAsync();
730  void MKLDNNDataReorderAsync(const mkldnn::memory::primitive_desc &desc);
731 
732  /*
733  * This creates a new NDArray with the reordered data.
734  * It doesn't affect the data of the original NDArray.
735  */
736  NDArray Reorder2Default() const;
737 
738  void InvalidateMKLDNNData();
739 
740  /*
741  * This function is used inside operators to reshape an array.
742  * It doesn't change the layout of the original array and allocate memory from
743  * the temporary buffer. The returned array is only valid inside the current
744  * invocation of this operator.
745  * This is different from Reshape. Reshape will cause data in the array to be
746  * converted to the default layout and allocate memory from malloc directly,
747  * which can be expensive.
748  * It's used by FullyConnected right now.
749  */
750  NDArray MKLDNNDataReshape(const TShape &shape) const;
751 
755  void UpdateMKLDNNMemDesc();
756 #endif
757 
764  static void Save(dmlc::Stream* fo,
765  const std::vector<NDArray>& data,
766  const std::vector<std::string>& names);
773  static void Load(dmlc::Stream* fi,
774  std::vector<NDArray>* data,
775  std::vector<std::string>* keys);
776 
777  private:
778  friend class Imperative;
780  // shandle is used to store the actual values in the NDArray
781  // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
782  struct Chunk {
786  Storage::Handle shandle;
791  std::vector<Storage::Handle> aux_handles;
792 
793 #if MXNET_USE_MKLDNN == 1
794 
796  std::shared_ptr<MKLDNNMemory> mkl_mem_;
797 #endif
798 
799  Engine::VarHandle var;
805  bool static_data;
808  bool delay_alloc;
809  // the type of the storage. The storage_type is never kUndefinedStorage once the chunk
810  // is constructed.
811  NDArrayStorageType storage_type = kDefaultStorage;
813  std::vector<int> aux_types;
814  // context of data
815  Context ctx;
816  // The shape of the chunk data.
817  // This might not be the same shape as the NDArray, since the storage may be sparse.
818  // The default value for storage_shape is {0} when an empty non-default NDArray is created.
819  TShape storage_shape;
820  // The shape of aux data. The default value for the shape depends on the type of storage.
821  // If aux_shapes[i].Size() is zero, aux data i is empty.
822  std::vector<TShape> aux_shapes;
823 
825  Chunk() : static_data(true), delay_alloc(false) {}
826 
828  Chunk(TShape shape, Context ctx_, bool delay_alloc_, int dtype)
829  : static_data(false), delay_alloc(true), ctx(ctx_) {
830  auto size = shape.Size();
831  storage_shape = shape;
832  var = Engine::Get()->NewVariable();
833  shandle.size = size * mshadow::mshadow_sizeof(dtype);
834  shandle.ctx = ctx_;
835  if (!delay_alloc_) this->CheckAndAlloc();
836  }
837 
838  Chunk(const TBlob &data, int dev_id)
839  : static_data(true), delay_alloc(false) {
840  CHECK(storage_type == kDefaultStorage);
841  var = Engine::Get()->NewVariable();
842  if (data.dev_mask() == cpu::kDevMask) {
843  ctx = Context::CPU();
844  } else {
845  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
846  ctx = Context::GPU(dev_id);
847  }
848  // init shandle
849  shandle.ctx = ctx;
850  shandle.dptr = data.dptr_;
851  shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
852  storage_shape = data.shape_;
853  }
854 
855  Chunk(int shared_pid, int shared_id, const TShape& shape, int dtype)
856  : static_data(false), delay_alloc(false) {
857  var = Engine::Get()->NewVariable();
858  ctx = Context::CPUShared(0);
859  shandle.size = shape.Size() * mshadow::mshadow_sizeof(dtype);
860  shandle.ctx = ctx;
861  shandle.shared_pid = shared_pid;
862  shandle.shared_id = shared_id;
863  Storage::Get()->Alloc(&shandle);
864  storage_shape = shape;
865  }
866  // Constructor for a non-default storage chunk
867  Chunk(NDArrayStorageType storage_type_, const TShape &storage_shape_, Context ctx_,
868  bool delay_alloc_, int dtype, const std::vector<int> &aux_types_,
869  const std::vector<TShape> &aux_shapes_)
870  : static_data(false), delay_alloc(delay_alloc_), storage_type(storage_type_),
871  aux_types(aux_types_), ctx(ctx_), storage_shape(storage_shape_),
872  aux_shapes(aux_shapes_) {
873  shandle.ctx = ctx;
874  var = Engine::Get()->NewVariable();
875  // aux_handles always reflect the correct number of aux data
876  for (size_t i = 0; i < aux_shapes.size(); i++) {
877  CheckAndAllocAuxData(i, aux_shapes[i]);
878  // this line is needed in case when aux_shapes[i].Size() = 0
879  // aux_handles[i] will not be updated and take only default value.
880  aux_handles[i].ctx = ctx;
881  }
882  if (!delay_alloc) {
883  CheckAndAllocData(storage_shape, dtype);
884  }
885  }
886 
887  Chunk(const NDArrayStorageType storage_type_, const TBlob &data,
888  const std::vector<TBlob> &aux_data, int dev_id)
889  : static_data(true), delay_alloc(false), storage_type(storage_type_) {
890  using namespace mshadow;
891  CHECK_NE(storage_type, kDefaultStorage);
892  // init var
893  var = Engine::Get()->NewVariable();
894  // init ctx
895  if (data.dev_mask() == cpu::kDevMask) {
896  ctx = Context::CPU();
897  } else {
898  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
899  ctx = Context::GPU(dev_id);
900  }
901  // init shandle
902  shandle.ctx = ctx;
903  shandle.dptr = data.dptr_;
904  shandle.size = data.shape_.Size() * mshadow_sizeof(data.type_flag_);
905  storage_shape = data.shape_;
906  // init aux handles
907  for (const auto &aux : aux_data) {
908  Storage::Handle aux_handle;
909  aux_handle.ctx = ctx;
910  aux_handle.dptr = aux.dptr_;
911  aux_handle.size = aux.shape_.Size() * mshadow_sizeof(aux.type_flag_);
912  aux_handles.push_back(aux_handle);
913  aux_types.emplace_back(aux.type_flag_);
914  aux_shapes.emplace_back(aux.shape_);
915  }
916  }
917 
919  inline void set_aux_shape(const size_t i, const TShape& shape) {
920  aux_shapes[i] = shape;
921  if (storage_shape.ndim() > 0) {
922  if (storage_type == kRowSparseStorage && i == rowsparse::kIdx) {
923  storage_shape[0] = shape[0];
924  } else if (storage_type == kCSRStorage && i == csr::kIdx) {
925  storage_shape[0] = shape[0];
926  }
927  }
928  }
929 
931  inline void CheckAndAlloc(void) {
932  if (delay_alloc) {
933  shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
934 #if MXNET_USE_MKLDNN == 1
935  mkl_mem_ = nullptr;
936 #endif
937  delay_alloc = false;
938  }
939  }
940 
942  // size is the number of bytes
943  void CheckAndAlloc(uint64_t dbytes) {
944  CHECK_EQ(kDefaultStorage, storage_type)
945  << "CheckAndAlloc(dbytes) is only intended for kDefaultStorage";
946  dbytes = std::max(dbytes, static_cast<uint64_t>(shandle.size));
947  if (delay_alloc) {
948  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
949 #if MXNET_USE_MKLDNN == 1
950  mkl_mem_ = nullptr;
951 #endif
952  delay_alloc = false;
953  } else if (shandle.size < dbytes) {
954  // free storage if necessary and alloc again
955  if (shandle.size > 0) Storage::Get()->Free(shandle);
956  // init storage
957  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
958 #if MXNET_USE_MKLDNN == 1
959  mkl_mem_ = nullptr;
960 #endif
961  }
962  }
963 
964  inline void CheckAndAlloc(const TShape &shape, const std::vector<TShape> &aux_shapes,
965  int dtype) {
966  // calculate size, perform allocation
967  if (kRowSparseStorage == storage_type) {
968  // For row sparse, aux_shape indicates the number of rows to allocate
969  auto aux_shape = aux_shapes[rowsparse::kIdx];
970  CheckAndAllocAuxData(rowsparse::kIdx, aux_shape);
971  TShape storage_shape(shape);
972  storage_shape[0] = aux_shape[0];
973  CheckAndAllocData(storage_shape, dtype);
974  } else if (kCSRStorage == storage_type) {
975  CheckAndAllocAuxData(csr::kIndPtr, aux_shapes[csr::kIndPtr]);
976  CheckAndAllocAuxData(csr::kIdx, aux_shapes[csr::kIdx]);
977  CheckAndAllocData(aux_shapes[csr::kIdx], dtype);
978  } else {
979  LOG(FATAL) << "Storage type " << storage_type << " not implemented for CheckAndAlloc";
980  }
981  }
982  // create storage handle for data based on shape and dtype, assuming ctx is set
983  // storage shape is also updated
984  // if data is already allocated, try reuse the storage. Otherwise, free the current one
985  // and allocate new storage
986  void CheckAndAllocData(const TShape &shape, int dtype);
987 
988 #if MXNET_USE_MKLDNN == 1
989  // Have MKL memory reference to the data in the default storage
990  // or create memory for MKLDNN.
991  void SetMKLMem(const TShape &shape, int dtype);
992  // If the data is stored in MKLDNN layout, we reorder data in mkl_mem_ and
993  // save the result in shandle.
994  void Reorder2Default();
995  // Reroder data to a specified layout.
996  void MKLDNNDataReorder(const mkldnn::memory::primitive_desc &desc);
997  bool IsMKLDNN() const;
998  bool IsDefault() const;
999 #endif
1000 
1001  // create storage handle for aux data based on shape
1002  // this function assumes ctx, aux shapes and aux types are set
1003  // aux shape is also updated
1004  // if aux data is already allocated, try reuse the storage. Otherwise, free the current one
1005  // and allocate new storage
1006  inline void CheckAndAllocAuxData(size_t i, const TShape &shape) {
1007  CHECK_EQ(shape.ndim(), 1) << "shape must be 1D in CheckAndAllocAuxData";
1008  CHECK_NE(storage_type, kUndefinedStorage)
1009  << "storage type cannot be kUndefinedStorage in CheckAndAllocAuxData";
1010  CHECK_NE(storage_type, kDefaultStorage)
1011  << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
1012  if (aux_handles.size() <= i) {
1013  aux_handles.resize(i + 1);
1014  }
1015  size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
1016  if (aux_handles[i].size < aux_bytes) {
1017  // free storage if necessary and alloc again
1018  if (aux_handles[i].size > 0) Storage::Get()->Free(aux_handles[i]);
1019  // init aux storage
1020  aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
1021  }
1022  // init shape
1023  set_aux_shape(i, shape);
1024  }
1026  ~Chunk();
1027  }; // struct Chunk
1028 
1029  void SetTBlob() const;
1030 
1032  std::shared_ptr<Chunk> ptr_{nullptr};
1034  TShape shape_;
1036  size_t byte_offset_ = 0;
1038  int dtype_ = -1;
1040  bool reuse_ = false;
1042  NDArrayStorageType storage_type_ = kUndefinedStorage;
1044  nnvm::NodeEntry entry_;
1052  mutable TBlob tblob_;
1053 }; // class NDArray
1054 
1058 size_t num_aux_data(NDArrayStorageType stype);
1059 
1071 void CopyFromTo(const NDArray &from, const NDArray *to, int priority = 0);
1072 
1086 void CopyFromTo(const NDArray &from, const NDArray& to, int priority = 0, bool is_opr = false);
1087 
1094 void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);
1095 
1102 NDArray operator+(const NDArray &lhs, const NDArray &rhs);
1109 NDArray operator+(const NDArray &lhs, const real_t &rhs);
1116 NDArray operator-(const NDArray &lhs, const NDArray &rhs);
1123 NDArray operator-(const NDArray &lhs, const real_t &rhs);
1130 NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
1137 NDArray operator*(const NDArray &lhs, const real_t &rhs);
1144 NDArray operator/(const NDArray &lhs, const NDArray &rhs);
1151 NDArray operator/(const NDArray &lhs, const real_t &rhs);
1152 
1157 void RandomSeed(uint32_t seed);
1162 void RandomSeed(Context ctx, uint32_t seed);
1169 void SampleUniform(real_t begin, real_t end, NDArray *out);
1176 void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
1183 void SampleGamma(real_t alpha, real_t beta, NDArray *out);
1189 void SampleExponential(real_t lambda, NDArray *out);
1195 void SamplePoisson(real_t lambda, NDArray *out);
1202 void SampleNegBinomial(int32_t k, real_t p, NDArray *out);
1209 void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out);
1210 
1211 
1212 //--------------------------------------------------------------
1213 // The following part are API Registration of NDArray functions.
1214 //--------------------------------------------------------------
1215 
1217 typedef std::function<void (NDArray **used_vars,
1218  real_t *scalars,
1219  NDArray **mutate_vars,
1220  int num_params,
1221  char **param_keys,
1222  char **param_vals)> NDArrayAPIFunction;
1238 };
1241  : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
1242  NDArrayAPIFunction> {
1244  unsigned num_use_vars;
1248  unsigned num_scalars;
1255  : num_use_vars(0),
1256  num_mutate_vars(0),
1257  num_scalars(0),
1258  type_mask(0) {}
1265  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
1266  NDArray *out)) {
1267  body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1268  int num_params, char **param_keys, char **param_vals) {
1269  (*fsetvalue)(s[0], mutate_vars[0]);
1270  };
1271  num_mutate_vars = 1; num_scalars = 1;
1272  this->add_argument("src", "real_t", "Source input to the function.");
1273  return *this;
1274  }
1281  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
1282  const NDArray &mhs,
1283  const NDArray &rhs,
1284  NDArray *out)) {
1285  body = [fternary](NDArray **used_vars,
1286  real_t *s, NDArray **mutate_vars,
1287  int num_params, char **param_keys, char **param_vals) {
1288  (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
1289  };
1290  num_use_vars = 3; num_mutate_vars = 1;
1292  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1293  this->add_argument("mhs", "NDArray", "Middle operand to the function.");
1294  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1295  return *this;
1296  }
1303  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
1304  const NDArray &rhs,
1305  NDArray *out)) {
1306  body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1307  int num_params, char **param_keys, char **param_vals) {
1308  (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
1309  };
1310  num_use_vars = 2; num_mutate_vars = 1;
1312  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1313  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1314  return *this;
1315  }
1322  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
1323  const real_t &rhs,
1324  NDArray *out)) {
1325  body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1326  int num_params, char **param_keys, char **param_vals) {
1327  (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
1328  };
1329  num_use_vars = 1; num_mutate_vars = 1; num_scalars = 1;
1331  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1332  this->add_argument("rhs", "real_t", "Right operand to the function.");
1333  return *this;
1334  }
1341  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
1342  NDArray *out)) {
1343  body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1344  int num_params, char **param_keys, char **param_vals) {
1345  (*funary)(*used_vars[0], mutate_vars[0]);
1346  };
1347  num_use_vars = 1; num_mutate_vars = 1;
1349  this->add_argument("src", "NDArray", "Source input to the function.");
1350  return *this;
1351  }
1359  void (*fgeneric)(NDArray **used_vars,
1360  real_t *s,
1361  NDArray **mutate_vars,
1362  const std::map<std::string, std::string>& param)) {
1363  body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1364  int num_params, char **param_keys, char **param_vals) {
1365  std::map<std::string, std::string> param;
1366  for (int i = 0; i < num_params; ++i) {
1367  param[param_keys[i]] = param_vals[i];
1368  }
1369  fgeneric(used_vars, s, mutate_vars, param);
1370  };
1371  return *this;
1372  }
1378  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
1379  num_use_vars = n; return *this;
1380  }
1387  num_mutate_vars = n; return *this;
1388  }
1394  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
1395  num_scalars = n; return *this;
1396  }
1402  inline NDArrayFunctionReg &set_type_mask(int tmask) {
1403  type_mask = tmask; return *this;
1404  }
1405 }; // NDArrayFunctionReg
1406 
1418 #define MXNET_REGISTER_NDARRAY_FUN(name) \
1419  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)
1420 
1421 } // namespace mxnet
1422 
1423 namespace dmlc {
1425 DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
1426 } // namespace dmlc
1427 #endif // MXNET_NDARRAY_H_
Definition: ndarray.h:74
Definition: ndarray.h:63
NDArrayStorageType
Definition: ndarray.h:61
Definition: ndarray.h:54
NDArrayFunctionReg & set_num_mutate_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:1386
NDArrayFormatErr
Definition: ndarray.h:68
Engine::VarHandle var() const
Definition: ndarray.h:356
void RandomSeed(uint32_t seed)
Seed all random number generator in mxnet.
NDArrayStorageType storage_type() const
Definition: ndarray.h:289
Engine that schedules all the operations according to dependency.
TShape shape_
shape of the tensor
Definition: tensor_blob.h:72
const TShape & storage_shape() const
Definition: ndarray.h:197
NDArrayFunctionReg()
constructor
Definition: ndarray.h:1254
namespace of mxnet
Definition: base.h:118
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:642
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:1378
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:126
static Context GPU(int32_t dev_id=-1)
int type_mask
information on how function should be called from API
Definition: ndarray.h:1250
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:1341
NDArray Detach() const
Return a copy of this NDArray without autograd history.
Definition: ndarray.h:617
int type_flag_
type flag of the tensor blob
Definition: tensor_blob.h:74
NDArrayFunctionReg & set_num_scalars(unsigned n)
set the number of scalar arguments
Definition: ndarray.h:1394
nnvm::TShape TShape
Shape data structure used to record shape information.
Definition: base.h:128
Definition: ndarray.h:72
unsigned num_mutate_vars
number of variable mutated by this function
Definition: ndarray.h:1246
execution time context. The information needed in runtime for actual execution.
Definition: base.h:257
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:1322
base class of engine variables.
Definition: engine.h:45
Definition: ndarray.h:65
Context ctx
Context information about device and ID.
Definition: storage.h:53
Storage::Handle storage_handle() const
get storage handle
Definition: ndarray.h:325
NDArray()
default constructor
Definition: ndarray.h:85
unsigned num_use_vars
number of variable used by this function
Definition: ndarray.h:1244
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:1281
Definition: ndarray.h:62
RowSparseAuxType
Definition: ndarray.h:58
Definition: ndarray.h:70
bool is_none() const
Definition: ndarray.h:293
all the scalar should go before use_vars
Definition: ndarray.h:1228
void SampleExponential(real_t lambda, NDArray *out)
Sample exponential distribution for each elements of out.
void SparseUpdateChunk(const NDArray &arr) const
Update ndarray chunk storage handles using existing ndarray storage handles Also update the aux_handl...
Definition: ndarray.h:573
void * dptr_
pointer to the data
Definition: tensor_blob.h:70
virtual VarHandle NewVariable()=0
Allocate a new variable, the variable can then be used to schedule the operation concurrently via dep...
Definition: ndarray.h:58
whether this function allows the handles in the target to be empty NDArray that are not yet initializ...
Definition: ndarray.h:1237
Definition: ndarray.h:73
static Storage * Get()
const TShape & shape() const
Definition: ndarray.h:189
Definition: ndarray.h:1423
virtual void WaitForVar(VarHandle var)=0
Wait for a variable.
const std::vector< TShape > & aux_shapes() const
Definition: ndarray.h:216
bool IsView() const
Definition: ndarray.h:167
Context ctx() const
Definition: ndarray.h:274
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:54
void SampleGaussian(real_t mu, real_t sigma, NDArray *out)
Sample gaussian distribution for each elements of out.
Definition: ndarray.h:54
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:335
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, bool wait=false)=0
Push an asynchronous operation to the engine.
int dtype() const
Definition: ndarray.h:281
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:304
Storage handle.
Definition: storage.h:41
static Context CPUShared(int32_t dev_id=0)
Definition: ndarray.h:64
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:236
void CheckAndAllocData(const TShape &storage_shape) const
Definition: ndarray.h:658
size_t num_aux_data(NDArrayStorageType stype)
NDArrayFunctionReg & set_type_mask(int tmask)
set type mask
Definition: ndarray.h:1402
void WaitToWrite() const
Block until all the pending read/write operations with respect to current NDArray are finished...
Definition: ndarray.h:343
NDArray(const TBlob &data, int dev_id, const std::function< void()> &deleter)
constructing a static NDArray that shares data with TBlob which is with deleter Use with caution: all...
Definition: ndarray.h:128
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:71
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:1265
NDArray(int shared_pid, int shared_id, const TShape &shape, int dtype)
create ndarray from shared memory
Definition: ndarray.h:140
NDArray operator+(const NDArray &lhs, const NDArray &rhs)
elementwise add
size_t byte_offset() const
Definition: ndarray.h:360
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:1240
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:1303
static Context CPU(int32_t dev_id=0)
runtime functions for NDArray
Definition: imperative.h:39
int aux_type(size_t i) const
Definition: ndarray.h:284
OnComplete Callback to the engine, called by AsyncFn when action completes.
Definition: engine.h:74
configuration of MXNet as well as basic data structure.
all the use_vars should go before scalar
Definition: ndarray.h:1226
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
size_t version() const
return var version of the NDArray
Definition: ndarray.h:364
void CheckAndAlloc(const std::vector< TShape > &aux_shapes) const
Definition: ndarray.h:653
unsigned num_scalars
number of scalars used by this function
Definition: ndarray.h:1248
static Engine * Get()
const TBlob & data() const
Definition: ndarray.h:245
void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const
Definition: ndarray.h:663
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:155
Definition: ndarray.h:69
void CheckAndAlloc() const
Allocate the space if it is delayed allocated. This is an internal function used by system that norma...
Definition: ndarray.h:628
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:124
size_t size
Size of the storage.
Definition: storage.h:49
TBlob aux_data(size_t i) const
Definition: ndarray.h:258
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:133
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:209
ndarray interface
Definition: ndarray.h:82
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:114
int dev_mask() const
device mask of the corresponding device
Definition: tensor_blob.h:242
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:1222
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:1358
bool IsSame(const NDArray &other) const
Definition: ndarray.h:179
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:66
const std::vector< int > & aux_types() const
Definition: ndarray.h:223
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:1224