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  }
120  NDArray(int shared_pid, int shared_id, const TShape& shape, int dtype)
121  : ptr_(std::make_shared<Chunk>(shared_pid, shared_id, shape, dtype)), shape_(shape),
122  dtype_(dtype), storage_type_(kDefaultStorage), entry_({nullptr, 0, 0}) {
123  }
124 
135  NDArray(const NDArrayStorageType stype, const TShape &shape,
136  const TBlob &data, const std::vector<TBlob> &aux_data, int dev_id)
137  : ptr_(std::make_shared<Chunk>(stype, data, aux_data, dev_id)), shape_(shape),
138  dtype_(data.type_flag_), storage_type_(stype), entry_({nullptr, 0, 0}) {
139  }
140 
141  /*
142  * This indicates whether an array is a view of another array (created by
143  * reshape or slice). If an array is a view and the the data is stored in
144  * MKLDNN format, we need to convert the data to the default format when
145  * data in the view is accessed.
146  */
147  inline bool IsView() const {
148  // View only works on the default storage
149  if (storage_type() != kDefaultStorage)
150  return false;
151  // If the array reuses memory, its shape may be different from the storage
152  // shape. However, we shouldn't consider it as a view.
153  if (reuse_)
154  return false;
155  return byte_offset_ > 0 || shape() != ptr_->storage_shape;
156  }
157 
158  /* \brief Check whether the two arrays are the same array */
159  inline bool IsSame(const NDArray& other) const {
160  return ptr_ == other.ptr_ &&
161  shape_ == other.shape_ &&
162  byte_offset_ == other.byte_offset_ &&
163  dtype_ == other.dtype_;
164  }
165 
169  inline const TShape& shape() const {
170  return shape_;
171  }
177  inline const TShape &storage_shape() const {
178  CHECK(ptr_ != nullptr);
179  CHECK_NE(storage_type(), kDefaultStorage)
180  << "storage_shape() is not intended for kDefaultStorage.";
181  return ptr_->storage_shape;
182  }
183 
189  inline const TShape& aux_shape(size_t index) const {
190  CHECK_NE(storage_type(), kDefaultStorage)
191  << "aux_shape() is not intended for kDefaultStorage.";
192  return ptr_->aux_shapes[index];
193  }
194 
195  /* \return the shapes of all aux data */
196  const std::vector<TShape>& aux_shapes() const {
197  CHECK_NE(storage_type(), kDefaultStorage)
198  << "aux_shapes() is not intended for kDefaultStorage.";
199  return ptr_->aux_shapes;
200  }
201 
203  const std::vector<int>& aux_types() const {
204  CHECK_NE(storage_type(), kDefaultStorage)
205  << "aux_types() is not intended for kDefaultStorage.";
206  return ptr_->aux_types;
207  }
208 
216  inline void set_aux_shape(size_t index, const TShape& shape) const {
217  CHECK_NE(storage_type(), kDefaultStorage)
218  << "set_aux_shape() is not intended for kDefaultStorage.";
219  ptr_->set_aux_shape(index, shape);
220  }
221 
225  inline const TBlob& data() const {
226  if (storage_type() == kDefaultStorage) CheckAndAlloc();
227  SetTBlob();
228  return tblob_;
229  }
233  NDArray grad() const;
234 
238  inline TBlob aux_data(size_t i) const {
239  auto stype = storage_type();
240  TBlob res;
241  auto shape = aux_shape(i);
242  auto type = aux_type(i);
243  MSHADOW_TYPE_SWITCH(type, DType, {
244  auto dptr = static_cast<DType*>(ptr_->aux_handles[i].dptr);
245  CHECK(stype == kRowSparseStorage || stype == kCSRStorage)
246  << "Unexpected storage type: " << stype;
247  res = TBlob(dptr, shape, ptr_->aux_handles[i].ctx.dev_mask(), type);
248  });
249  return res;
250  }
254  inline Context ctx() const {
255  CHECK(!is_none());
256  return ptr_->shandle.ctx;
257  }
261  inline int dtype() const {
262  return dtype_;
263  }
264  inline int aux_type(size_t i) const {
265  CHECK(!is_none());
266  return ptr_->aux_types[i];
267  }
268 
270  return storage_type_;
271  }
273  inline bool is_none() const {
274  return ptr_.get() == nullptr;
275  }
277  bool fresh_out_grad() const;
279  void set_fresh_out_grad(bool state) const;
284  inline bool storage_initialized() const {
285  if (is_none()) return false;
286  auto stype = storage_type();
287  CHECK_NE(stype, kDefaultStorage)
288  << "storage_initialized() is not intended for kDefaultStorage.";
289  if (stype == kRowSparseStorage) {
290  CHECK_EQ(aux_shape(rowsparse::kIdx)[0], storage_shape()[0])
291  << "inconsistent storage shape " << storage_shape()
292  << " vs. aux shape " << aux_shape(rowsparse::kIdx);
293  return aux_shape(rowsparse::kIdx).Size() != 0;
294  } else if (stype == kCSRStorage) {
295  CHECK_EQ(aux_shape(csr::kIdx)[0], storage_shape()[0])
296  << "inconsistent storage shape " << storage_shape()
297  << " vs. aux shape " << aux_shape(csr::kIdx);
298  return aux_shape(csr::kIdx).Size() != 0;
299  } else {
300  LOG(FATAL) << "Unknown storage type";
301  }
302  return true;
303  }
306  CHECK(!is_none());
307  CHECK_EQ(storage_type(), kDefaultStorage);
308  CheckAndAlloc();
309  return ptr_->shandle;
310  }
315  inline void WaitToRead() const {
316  if (is_none()) return;
317  Engine::Get()->WaitForVar(ptr_->var);
318  }
323  inline void WaitToWrite() const {
324  if (is_none()) return;
330  [](RunContext, Engine::CallbackOnComplete on_complete) {
331  on_complete();
332  }, Context{}, {}, {ptr_->var});
333  Engine::Get()->WaitForVar(ptr_->var);
334  }
336  inline Engine::VarHandle var() const {
337  return ptr_->var;
338  }
340  inline size_t byte_offset() const {
341  return byte_offset_;
342  }
347  void Save(dmlc::Stream *strm) const;
353  bool LegacyLoad(dmlc::Stream *strm, const uint32_t magic);
359  bool Load(dmlc::Stream *strm);
365  NDArray &operator=(real_t scalar);
372  NDArray &operator+=(const NDArray &src);
379  NDArray &operator+=(const real_t &src);
386  NDArray &operator-=(const NDArray &src);
393  NDArray &operator-=(const real_t &src);
400  NDArray &operator*=(const NDArray &src);
407  NDArray &operator*=(const real_t &src);
414  NDArray &operator/=(const NDArray &src);
421  NDArray &operator/=(const real_t &src);
427  NDArray Copy(Context ctx) const;
438  void SyncCopyFromCPU(const void *data, size_t size) const;
439 
443  void SyncCopyFromNDArray(const NDArray &src, int i = -1, int j = -1);
444 
455  void SyncCopyToCPU(void *data, size_t size) const;
461  void SyncCheckFormat(const bool full_check) const;
468  NDArray Slice(index_t begin, index_t end) const;
475  NDArray SliceWithRecord(index_t begin, index_t end);
481  NDArray At(index_t idx) const;
487  NDArray AtWithRecord(index_t idx);
492  NDArray aux_ndarray(size_t i) const;
493 
498  NDArray data_ndarray() const;
499 
507  inline NDArray AsArray(const TShape &shape, int dtype) const {
508  CHECK_EQ(storage_type(), kDefaultStorage)
509  << "AsArray is intended only for kDefaultStorage.";
510  CHECK_GE(ptr_->shandle.size,
511  shape.Size() * mshadow::mshadow_sizeof(dtype))
512  << "NDArray.AsArray: target memory size is bigger";
513  // We can't reuse memory in a view.
514  CHECK(!IsView());
515  NDArray ret = *this;
516  ret.shape_ = shape;
517  ret.dtype_ = dtype;
518  ret.reuse_ = true;
519  return ret;
520  }
521 
529  inline void SparseUpdateChunk(const NDArray &arr) const {
530  CHECK(shape_ == arr.shape_) << "ndarray shape is different from the target";
531  CHECK(dtype_ == arr.dtype_) << "ndarray dtype is different from the target";
532  auto stype = arr.storage_type();
533  CHECK(stype == kCSRStorage || stype == kRowSparseStorage)
534  << "Only to be used with CSR and RSP storage types";
535  // swap shandles between src and dst
536  Storage::Handle shandle_dst = arr.ptr_->shandle;
537  arr.ptr_->shandle = ptr_->shandle;
538  ptr_->shandle = shandle_dst;
539 
540  ptr_->storage_shape = arr.ptr_->storage_shape;
541  ptr_->storage_type = arr.ptr_->storage_type;
542  ptr_->ctx = arr.ptr_->ctx;
543 
544  // swap aux_handles between src and dst
545  size_t aux_idx = 0;
546  CHECK(ptr_->aux_handles.size() == arr.ptr_->aux_handles.size())
547  << "ndarray number of aux_handles is different from target";
548  for (auto &aux_handle : arr.ptr_->aux_handles) {
549  Storage::Handle aux_dst = ptr_->aux_handles[aux_idx];
550  ptr_->aux_handles[aux_idx] = aux_handle;
551  aux_handle = aux_dst;
552  aux_idx++;
553  }
554  ptr_->aux_types = arr.ptr_->aux_types;
555  ptr_->aux_shapes = arr.ptr_->aux_shapes;
556  }
557 
563  NDArray Reshape(const TShape &shape) const;
569  NDArray ReshapeWithRecord(const TShape &shape);
573  NDArray Detach() const {
574  NDArray ret(*this);
575  ret.entry_ = nnvm::NodeEntry{nullptr, 0, 0};
576  return ret;
577  }
578 
579  nnvm::Symbol get_autograd_symbol() const;
584  inline void CheckAndAlloc() const {
585  CHECK_EQ(storage_type(), kDefaultStorage);
586  ptr_->CheckAndAlloc();
587  }
588 
598  void ReshapeAndAlloc(const TShape& shape) {
599  CHECK_EQ(storage_type(), kDefaultStorage);
600  CHECK(!is_none());
601  shape_ = shape;
602  ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
603  }
604 
605  /* !
606  * \brief Alloc memory for non-default storage
607  * aux_shape is only known at run time
608  */
609  inline void CheckAndAlloc(const std::vector<TShape> &aux_shapes) const {
610  CHECK_NE(storage_type(), kDefaultStorage)
611  << "CheckAndAlloc(aux_shapes) is not intended for kDefaultStorage";
612  ptr_->CheckAndAlloc(shape_, aux_shapes, dtype_);
613  }
614  inline void CheckAndAllocData(const TShape &storage_shape) const {
615  CHECK_NE(storage_type(), kDefaultStorage)
616  << "CheckAndAllocData is not intended for kDefaultStorage";
617  ptr_->CheckAndAllocData(storage_shape, dtype_);
618  }
619  inline void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const {
620  CHECK_NE(storage_type(), kDefaultStorage)
621  << "CheckAndAllocAuxData is not intended for kDefaultStorage";
622  ptr_->CheckAndAllocAuxData(i, aux_shape);
623  }
624 
625 #if MXNET_USE_MKLDNN == 1
626  /*
627  * Test if the data is stored in one of special MKLDNN format.
628  */
629  bool IsMKLDNNData() const {
630  return ptr_->IsMKLDNN();
631  }
632  /*
633  * Test if the data is stored in one of default MXNet formats.
634  */
635  bool IsDefaultData() const {
636  return ptr_->IsDefault();
637  }
638  /*
639  * All functions below return a raw pointer to mkldnn memory. Actually there
640  * is a shared pointer that hold the memory either in NDArray or in MKLDNN
641  * stream. As long as we call these functions inside an operator, the return
642  * memory is always valid.
643  */
644 
645  /*
646  * This function returns mkldnn::memory with the default primitive_desc.
647  */
648  const mkldnn::memory *GetMKLDNNData() const;
649  /*
650  * This function returns mkldnn::memory with the given primitive_desc
651  * as long as the array size meets the required size in the given primitive_desc.
652  */
653  const mkldnn::memory *GetMKLDNNData(
654  const mkldnn::memory::primitive_desc &desc) const;
655  /*
656  * This function returns mkldnn::memory with the given primitive_desc.
657  * The returned mkldnn::memory will have the same physical layout as
658  * the given primitive_desc.
659  */
660  const mkldnn::memory *GetMKLDNNDataReorder(
661  const mkldnn::memory::primitive_desc &desc) const;
662 
663  /*
664  * This function copies data from mkldnn memory.
665  */
666  void CopyFrom(const mkldnn::memory &mem);
667  /*
668  * This function allocates memory for array and creates mkldnn memory
669  * with the specified format.
670  */
671  mkldnn::memory *CreateMKLDNNData(
672  const mkldnn::memory::primitive_desc &desc);
673 
674  /*
675  * These are the async version of the methods above.
676  * It changes the layout of this NDArray, but it happens after all accesses to
677  * the array are complete.
678  */
679  void Reorder2DefaultAsync();
680  void MKLDNNDataReorderAsync(const mkldnn::memory::primitive_desc &desc);
681 
682  /*
683  * This creates a new NDArray with the reordered data.
684  * It doesn't affect the data of the original NDArray.
685  */
686  NDArray Reorder2Default() const;
687 
688  void InvalidateMKLDNNData();
689 
690  /*
691  * This function is used inside operators to reshape an array.
692  * It doesn't change the layout of the original array and allocate memory from
693  * the temporary buffer. The returned array is only valid inside the current
694  * invocation of this operator.
695  * This is different from Reshape. Reshape will cause data in the array to be
696  * converted to the default layout and allocate memory from malloc directly,
697  * which can be expensive.
698  * It's used by FullyConnected right now.
699  */
700  NDArray MKLDNNDataReshape(const TShape &shape) const;
701 #endif
702 
709  static void Save(dmlc::Stream* fo,
710  const std::vector<NDArray>& data,
711  const std::vector<std::string>& names);
718  static void Load(dmlc::Stream* fi,
719  std::vector<NDArray>* data,
720  std::vector<std::string>* keys);
721 
722  private:
723  friend class Imperative;
725  // shandle is used to store the actual values in the NDArray
726  // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
727  struct Chunk {
731  Storage::Handle shandle;
736  std::vector<Storage::Handle> aux_handles;
737 
738 #if MXNET_USE_MKLDNN == 1
739 
741  std::shared_ptr<MKLDNNMemory> mkl_mem_;
742 #endif
743 
744  Engine::VarHandle var;
750  bool static_data;
753  bool delay_alloc;
754  // the type of the storage. The storage_type is never kUndefinedStorage once the chunk
755  // is constructed.
756  NDArrayStorageType storage_type = kDefaultStorage;
758  std::vector<int> aux_types;
759  // context of data
760  Context ctx;
761  // The shape of the chunk data.
762  // This might not be the same shape as the NDArray, since the storage may be sparse.
763  // The default value for storage_shape is {0} when an empty non-default NDArray is created.
764  TShape storage_shape;
765  // The shape of aux data. The default value for the shape depends on the type of storage.
766  // If aux_shapes[i].Size() is zero, aux data i is empty.
767  std::vector<TShape> aux_shapes;
768 
770  Chunk() : static_data(true), delay_alloc(false) {}
771 
773  Chunk(TShape shape, Context ctx_, bool delay_alloc_, int dtype)
774  : static_data(false), delay_alloc(true), ctx(ctx_) {
775  auto size = shape.Size();
776  storage_shape = shape;
777  var = Engine::Get()->NewVariable();
778  shandle.size = size * mshadow::mshadow_sizeof(dtype);
779  shandle.ctx = ctx_;
780  if (!delay_alloc_) this->CheckAndAlloc();
781  }
782 
783  Chunk(const TBlob &data, int dev_id)
784  : static_data(true), delay_alloc(false) {
785  CHECK(storage_type == kDefaultStorage);
786  var = Engine::Get()->NewVariable();
787  if (data.dev_mask() == cpu::kDevMask) {
788  ctx = Context::CPU();
789  } else {
790  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
791  ctx = Context::GPU(dev_id);
792  }
793  // init shandle
794  shandle.ctx = ctx;
795  shandle.dptr = data.dptr_;
796  shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
797  storage_shape = data.shape_;
798  }
799 
800  Chunk(int shared_pid, int shared_id, const TShape& shape, int dtype)
801  : static_data(false), delay_alloc(false) {
802  var = Engine::Get()->NewVariable();
803  ctx = Context::CPUShared(0);
804  shandle.size = shape.Size() * mshadow::mshadow_sizeof(dtype);
805  shandle.ctx = ctx;
806  shandle.shared_pid = shared_pid;
807  shandle.shared_id = shared_id;
808  Storage::Get()->Alloc(&shandle);
809  storage_shape = shape;
810  }
811  // Constructor for a non-default storage chunk
812  Chunk(NDArrayStorageType storage_type_, const TShape &storage_shape_, Context ctx_,
813  bool delay_alloc_, int dtype, const std::vector<int> &aux_types_,
814  const std::vector<TShape> &aux_shapes_)
815  : static_data(false), delay_alloc(delay_alloc_), storage_type(storage_type_),
816  aux_types(aux_types_), ctx(ctx_), storage_shape(storage_shape_),
817  aux_shapes(aux_shapes_) {
818  shandle.ctx = ctx;
819  var = Engine::Get()->NewVariable();
820  // aux_handles always reflect the correct number of aux data
821  for (size_t i = 0; i < aux_shapes.size(); i++) {
822  CheckAndAllocAuxData(i, aux_shapes[i]);
823  // this line is needed in case when aux_shapes[i].Size() = 0
824  // aux_handles[i] will not be updated and take only default value.
825  aux_handles[i].ctx = ctx;
826  }
827  if (!delay_alloc) {
828  CheckAndAllocData(storage_shape, dtype);
829  }
830  }
831 
832  Chunk(const NDArrayStorageType storage_type_, const TBlob &data,
833  const std::vector<TBlob> &aux_data, int dev_id)
834  : static_data(true), delay_alloc(false), storage_type(storage_type_) {
835  using namespace mshadow;
836  CHECK_NE(storage_type, kDefaultStorage);
837  // init var
838  var = Engine::Get()->NewVariable();
839  // init ctx
840  if (data.dev_mask() == cpu::kDevMask) {
841  ctx = Context::CPU();
842  } else {
843  CHECK_EQ(data.dev_mask(), gpu::kDevMask);
844  ctx = Context::GPU(dev_id);
845  }
846  // init shandle
847  shandle.ctx = ctx;
848  shandle.dptr = data.dptr_;
849  shandle.size = data.shape_.Size() * mshadow_sizeof(data.type_flag_);
850  storage_shape = data.shape_;
851  // init aux handles
852  for (const auto &aux : aux_data) {
853  Storage::Handle aux_handle;
854  aux_handle.ctx = ctx;
855  aux_handle.dptr = aux.dptr_;
856  aux_handle.size = aux.shape_.Size() * mshadow_sizeof(aux.type_flag_);
857  aux_handles.push_back(aux_handle);
858  aux_types.emplace_back(aux.type_flag_);
859  aux_shapes.emplace_back(aux.shape_);
860  }
861  }
862 
864  inline void set_aux_shape(const size_t i, const TShape& shape) {
865  aux_shapes[i] = shape;
866  if (storage_shape.ndim() > 0) {
867  if (storage_type == kRowSparseStorage && i == rowsparse::kIdx) {
868  storage_shape[0] = shape[0];
869  } else if (storage_type == kCSRStorage && i == csr::kIdx) {
870  storage_shape[0] = shape[0];
871  }
872  }
873  }
874 
876  inline void CheckAndAlloc(void) {
877  if (delay_alloc) {
878  shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
879 #if MXNET_USE_MKLDNN == 1
880  mkl_mem_ = nullptr;
881 #endif
882  delay_alloc = false;
883  }
884  }
885 
887  // size is the number of bytes
888  void CheckAndAlloc(uint64_t dbytes) {
889  CHECK_EQ(kDefaultStorage, storage_type)
890  << "CheckAndAlloc(dbytes) is only intended for kDefaultStorage";
891  dbytes = std::max(dbytes, static_cast<uint64_t>(shandle.size));
892  if (delay_alloc) {
893  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
894 #if MXNET_USE_MKLDNN == 1
895  mkl_mem_ = nullptr;
896 #endif
897  delay_alloc = false;
898  } else if (shandle.size < dbytes) {
899  // free storage if necessary and alloc again
900  if (shandle.size > 0) Storage::Get()->Free(shandle);
901  // init storage
902  shandle = Storage::Get()->Alloc(dbytes, shandle.ctx);
903 #if MXNET_USE_MKLDNN == 1
904  mkl_mem_ = nullptr;
905 #endif
906  }
907  }
908 
909  inline void CheckAndAlloc(const TShape &shape, const std::vector<TShape> &aux_shapes,
910  int dtype) {
911  // calculate size, perform allocation
912  if (kRowSparseStorage == storage_type) {
913  // For row sparse, aux_shape indicates the number of rows to allocate
914  auto aux_shape = aux_shapes[rowsparse::kIdx];
915  CheckAndAllocAuxData(rowsparse::kIdx, aux_shape);
916  TShape storage_shape(shape);
917  storage_shape[0] = aux_shape[0];
918  CheckAndAllocData(storage_shape, dtype);
919  } else if (kCSRStorage == storage_type) {
920  CheckAndAllocAuxData(csr::kIndPtr, aux_shapes[csr::kIndPtr]);
921  CheckAndAllocAuxData(csr::kIdx, aux_shapes[csr::kIdx]);
922  CheckAndAllocData(aux_shapes[csr::kIdx], dtype);
923  } else {
924  LOG(FATAL) << "Storage type " << storage_type << " not implemented for CheckAndAlloc";
925  }
926  }
927  // create storage handle for data based on shape and dtype, assuming ctx is set
928  // storage shape is also updated
929  // if data is already allocated, try reuse the storage. Otherwise, free the current one
930  // and allocate new storage
931  void CheckAndAllocData(const TShape &shape, int dtype);
932 
933 #if MXNET_USE_MKLDNN == 1
934  // Have MKL memory reference to the data in the default storage
935  // or create memory for MKLDNN.
936  void SetMKLMem(const TShape &shape, int dtype);
937  // If the data is stored in MKLDNN layout, we reorder data in mkl_mem_ and
938  // save the result in shandle.
939  void Reorder2Default();
940  // Reroder data to a specified layout.
941  void MKLDNNDataReorder(const mkldnn::memory::primitive_desc &desc);
942  bool IsMKLDNN() const;
943  bool IsDefault() const;
944 #endif
945 
946  // create storage handle for aux data based on shape
947  // this function assumes ctx, aux shapes and aux types are set
948  // aux shape is also updated
949  // if aux data is already allocated, try reuse the storage. Otherwise, free the current one
950  // and allocate new storage
951  inline void CheckAndAllocAuxData(size_t i, const TShape &shape) {
952  CHECK_EQ(shape.ndim(), 1) << "shape must be 1D in CheckAndAllocAuxData";
953  CHECK_NE(storage_type, kUndefinedStorage)
954  << "storage type cannot be kUndefinedStorage in CheckAndAllocAuxData";
955  CHECK_NE(storage_type, kDefaultStorage)
956  << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
957  if (aux_handles.size() <= i) {
958  aux_handles.resize(i + 1);
959  }
960  size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
961  if (aux_handles[i].size < aux_bytes) {
962  // free storage if necessary and alloc again
963  if (aux_handles[i].size > 0) Storage::Get()->Free(aux_handles[i]);
964  // init aux storage
965  aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
966  }
967  // init shape
968  set_aux_shape(i, shape);
969  }
971  ~Chunk();
972  }; // struct Chunk
973 
974  void SetTBlob() const;
975 
977  std::shared_ptr<Chunk> ptr_{nullptr};
979  TShape shape_;
981  size_t byte_offset_ = 0;
983  int dtype_ = -1;
985  bool reuse_ = false;
987  NDArrayStorageType storage_type_ = kUndefinedStorage;
989  nnvm::NodeEntry entry_;
997  mutable TBlob tblob_;
998 }; // class NDArray
999 
1003 size_t num_aux_data(NDArrayStorageType stype);
1004 
1016 void CopyFromTo(const NDArray &from, const NDArray *to, int priority = 0);
1017 
1031 void CopyFromTo(const NDArray &from, const NDArray& to, int priority = 0, bool is_opr = false);
1032 
1039 void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);
1040 
1047 NDArray operator+(const NDArray &lhs, const NDArray &rhs);
1054 NDArray operator+(const NDArray &lhs, const real_t &rhs);
1061 NDArray operator-(const NDArray &lhs, const NDArray &rhs);
1068 NDArray operator-(const NDArray &lhs, const real_t &rhs);
1075 NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
1082 NDArray operator*(const NDArray &lhs, const real_t &rhs);
1089 NDArray operator/(const NDArray &lhs, const NDArray &rhs);
1096 NDArray operator/(const NDArray &lhs, const real_t &rhs);
1097 
1102 void RandomSeed(uint32_t seed);
1107 void RandomSeed(Context ctx, uint32_t seed);
1114 void SampleUniform(real_t begin, real_t end, NDArray *out);
1121 void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
1128 void SampleGamma(real_t alpha, real_t beta, NDArray *out);
1134 void SampleExponential(real_t lambda, NDArray *out);
1140 void SamplePoisson(real_t lambda, NDArray *out);
1147 void SampleNegBinomial(int32_t k, real_t p, NDArray *out);
1154 void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray *out);
1155 
1156 
1157 //--------------------------------------------------------------
1158 // The following part are API Registration of NDArray functions.
1159 //--------------------------------------------------------------
1160 
1162 typedef std::function<void (NDArray **used_vars,
1163  real_t *scalars,
1164  NDArray **mutate_vars,
1165  int num_params,
1166  char **param_keys,
1167  char **param_vals)> NDArrayAPIFunction;
1183 };
1186  : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
1187  NDArrayAPIFunction> {
1189  unsigned num_use_vars;
1193  unsigned num_scalars;
1200  : num_use_vars(0),
1201  num_mutate_vars(0),
1202  num_scalars(0),
1203  type_mask(0) {}
1210  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
1211  NDArray *out)) {
1212  body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1213  int num_params, char **param_keys, char **param_vals) {
1214  (*fsetvalue)(s[0], mutate_vars[0]);
1215  };
1216  num_mutate_vars = 1; num_scalars = 1;
1217  this->add_argument("src", "real_t", "Source input to the function.");
1218  return *this;
1219  }
1226  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
1227  const NDArray &mhs,
1228  const NDArray &rhs,
1229  NDArray *out)) {
1230  body = [fternary](NDArray **used_vars,
1231  real_t *s, NDArray **mutate_vars,
1232  int num_params, char **param_keys, char **param_vals) {
1233  (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
1234  };
1235  num_use_vars = 3; num_mutate_vars = 1;
1237  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1238  this->add_argument("mhs", "NDArray", "Middle operand to the function.");
1239  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1240  return *this;
1241  }
1248  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
1249  const NDArray &rhs,
1250  NDArray *out)) {
1251  body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1252  int num_params, char **param_keys, char **param_vals) {
1253  (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
1254  };
1255  num_use_vars = 2; num_mutate_vars = 1;
1257  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1258  this->add_argument("rhs", "NDArray", "Right operand to the function.");
1259  return *this;
1260  }
1267  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
1268  const real_t &rhs,
1269  NDArray *out)) {
1270  body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1271  int num_params, char **param_keys, char **param_vals) {
1272  (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
1273  };
1274  num_use_vars = 1; num_mutate_vars = 1; num_scalars = 1;
1276  this->add_argument("lhs", "NDArray", "Left operand to the function.");
1277  this->add_argument("rhs", "real_t", "Right operand to the function.");
1278  return *this;
1279  }
1286  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
1287  NDArray *out)) {
1288  body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1289  int num_params, char **param_keys, char **param_vals) {
1290  (*funary)(*used_vars[0], mutate_vars[0]);
1291  };
1292  num_use_vars = 1; num_mutate_vars = 1;
1294  this->add_argument("src", "NDArray", "Source input to the function.");
1295  return *this;
1296  }
1304  void (*fgeneric)(NDArray **used_vars,
1305  real_t *s,
1306  NDArray **mutate_vars,
1307  const std::map<std::string, std::string>& param)) {
1308  body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
1309  int num_params, char **param_keys, char **param_vals) {
1310  std::map<std::string, std::string> param;
1311  for (int i = 0; i < num_params; ++i) {
1312  param[param_keys[i]] = param_vals[i];
1313  }
1314  fgeneric(used_vars, s, mutate_vars, param);
1315  };
1316  return *this;
1317  }
1323  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
1324  num_use_vars = n; return *this;
1325  }
1332  num_mutate_vars = n; return *this;
1333  }
1339  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
1340  num_scalars = n; return *this;
1341  }
1347  inline NDArrayFunctionReg &set_type_mask(int tmask) {
1348  type_mask = tmask; return *this;
1349  }
1350 }; // NDArrayFunctionReg
1351 
1363 #define MXNET_REGISTER_NDARRAY_FUN(name) \
1364  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)
1365 
1366 } // namespace mxnet
1367 
1368 namespace dmlc {
1370 DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
1371 } // namespace dmlc
1372 #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:1331
NDArrayFormatErr
Definition: ndarray.h:68
Engine::VarHandle var() const
Definition: ndarray.h:336
void RandomSeed(uint32_t seed)
Seed all random number generator in mxnet.
NDArrayStorageType storage_type() const
Definition: ndarray.h:269
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:177
NDArrayFunctionReg()
constructor
Definition: ndarray.h:1199
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:598
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:1323
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:1195
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:1286
NDArray Detach() const
Return a copy of this NDArray without autograd history.
Definition: ndarray.h:573
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:1339
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:1191
execution time context. The information needed in runtime for actual execution.
Definition: base.h:249
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:1267
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:305
NDArray()
default constructor
Definition: ndarray.h:85
unsigned num_use_vars
number of variable used by this function
Definition: ndarray.h:1189
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:1226
Definition: ndarray.h:62
RowSparseAuxType
Definition: ndarray.h:58
Definition: ndarray.h:70
bool is_none() const
Definition: ndarray.h:273
all the scalar should go before use_vars
Definition: ndarray.h:1173
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:529
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:1182
Definition: ndarray.h:73
static Storage * Get()
const TShape & shape() const
Definition: ndarray.h:169
Definition: ndarray.h:1368
virtual void WaitForVar(VarHandle var)=0
Wait for a variable.
const std::vector< TShape > & aux_shapes() const
Definition: ndarray.h:196
bool IsView() const
Definition: ndarray.h:147
Context ctx() const
Definition: ndarray.h:254
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:315
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:261
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:284
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:216
void CheckAndAllocData(const TShape &storage_shape) const
Definition: ndarray.h:614
size_t num_aux_data(NDArrayStorageType stype)
NDArrayFunctionReg & set_type_mask(int tmask)
set type mask
Definition: ndarray.h:1347
engine::VarHandle VarHandle
Variable pointer.
Definition: engine.h:107
void WaitToWrite() const
Block until all the pending read/write operations with respect to current NDArray are finished...
Definition: ndarray.h:323
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:1210
NDArray(int shared_pid, int shared_id, const TShape &shape, int dtype)
create ndarray from shared memory
Definition: ndarray.h:120
NDArray operator+(const NDArray &lhs, const NDArray &rhs)
elementwise add
size_t byte_offset() const
Definition: ndarray.h:340
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:1185
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:1248
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:264
OnComplete Callback to the engine, called by AsyncFn when action completes.
Definition: engine.h:56
configuration of MXNet as well as basic data structure.
all the use_vars should go before scalar
Definition: ndarray.h:1171
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:507
void CheckAndAlloc(const std::vector< TShape > &aux_shapes) const
Definition: ndarray.h:609
unsigned num_scalars
number of scalars used by this function
Definition: ndarray.h:1193
static Engine * Get()
const TBlob & data() const
Definition: ndarray.h:225
void CheckAndAllocAuxData(size_t i, const TShape &aux_shape) const
Definition: ndarray.h:619
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:135
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:584
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:238
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:189
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:209
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:1167
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:1303
bool IsSame(const NDArray &other) const
Definition: ndarray.h:159
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:203
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:1169