ndarray.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef MXNET_NDARRAY_H_
#define MXNET_NDARRAY_H_
 
#include <dmlc/base.h>
#include <dmlc/io.h>
#include <dmlc/logging.h>
#include <dmlc/registry.h>
#include <dmlc/type_traits.h>
#include <nnvm/node.h>
 
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "./base.h"
#include "./engine.h"
#include "./storage.h"
// check c++11
#if DMLC_USE_CXX11 == 0
#error "cxx11 was required for ndarray module"
#endif
 
namespace dnnl {
struct memory;
}  // namespace dnnl
 
namespace mxnet {
// enum for storage types
namespace csr {
enum CSRAuxType { kIndPtr, kIdx };
}
 
namespace rowsparse {
enum RowSparseAuxType { kIdx };
}
 
enum NDArrayStorageType {
  kUndefinedStorage = -1,  // undefined storage
  kDefaultStorage,         // dense
  kRowSparseStorage,       // row sparse
  kCSRStorage,             // csr
};
 
enum NDArrayFormatErr {
  kNormalErr,     // normal
  kCSRShapeErr,   // shape mismatch for csr
  kCSRIndPtrErr,  // indptr error for csr
  kCSRIdxErr,     // idx error for csr
  kRSPShapeErr,   // shape mismatch for row sparse
  kRSPIdxErr,     // indices error for row sparse
};
 
class DNNLMemory;
 
class NDArray {
 public:
  NDArray() : autograd_entry_(nullptr) {}
  NDArray(const mxnet::TShape& shape,
          Context ctx,
          bool delay_alloc = false,
          int dtype        = mshadow::default_type_flag)
      : ptr_(std::make_shared<Chunk>(shape, ctx, delay_alloc, dtype)),
        shape_(shape),
        dtype_(dtype),
        storage_type_(kDefaultStorage),
        autograd_entry_(nullptr) {}
  NDArray(const NDArrayStorageType stype,
          const mxnet::TShape& shape,
          Context ctx,
          bool delay_alloc                     = true,
          int dtype                            = mshadow::default_type_flag,
          const std::vector<int>& aux_types    = {},
          const mxnet::ShapeVector& aux_shapes = {},
          const mxnet::TShape& storage_shape   = mxnet::TShape(mshadow::Shape1(0))) {
    ReInit(stype, shape, ctx, dtype, delay_alloc, &aux_types, &aux_shapes, &storage_shape);
  }
  explicit NDArray(Context ctx, int dtype = mshadow::default_type_flag)
      : ptr_(std::make_shared<Chunk>(mxnet::TShape(mshadow::Shape1(0)), ctx, true, dtype)),
        shape_(),
        dtype_(dtype),
        storage_type_(kDefaultStorage),
        autograd_entry_(nullptr) {}
  NDArray(const TBlob& data, int dev_id)
      : ptr_(std::make_shared<Chunk>(data, dev_id)),
        shape_(data.shape_),
        dtype_(data.type_flag_),
        storage_type_(kDefaultStorage),
        autograd_entry_(nullptr) {}
 
  NDArray(const TBlob& data, int dev_id, const std::function<void()>& deleter)
      : ptr_(new Chunk(data, dev_id),
             [deleter](Chunk* p) {
               deleter();  // call custom deleter
               delete p;   // delete Chunk object
             }),
        shape_(data.shape_),
        dtype_(data.type_flag_),
        storage_type_(kDefaultStorage),
        autograd_entry_(nullptr) {}
 
  NDArray(int shared_pid, int shared_id, const mxnet::TShape& shape, int dtype)
      : ptr_(std::make_shared<Chunk>(shared_pid, shared_id, shape, dtype)),
        shape_(shape),
        dtype_(dtype),
        storage_type_(kDefaultStorage),
        autograd_entry_(nullptr) {}
 
  NDArray(const NDArrayStorageType stype,
          const mxnet::TShape& shape,
          const TBlob& data,
          const std::vector<TBlob>& aux_data,
          int dev_id)
      : ptr_(std::make_shared<Chunk>(stype, data, aux_data, dev_id)),
        shape_(shape),
        dtype_(data.type_flag_),
        storage_type_(stype),
        autograd_entry_(nullptr) {}
  void Init(const mxnet::TShape& shape) {
    ptr_->Init(shape, this->dtype_);
    this->shape_ = shape;
  }
 
  void InitDetached(const NDArray* src) {
    *this           = *src;
    autograd_entry_ = nnvm::NodeEntry(nullptr);
  }
  inline void ReInit() {
    ptr_ = nullptr;
    Init(kUndefinedStorage, TShape(), -1);
  }
  void ReInit(const NDArrayStorageType stype,
              const mxnet::TShape& shape,
              Context ctx,
              int dtype,
              bool delay_alloc                     = true,
              const std::vector<int>* aux_types    = nullptr,
              const mxnet::ShapeVector* aux_shapes = nullptr,
              const mxnet::TShape* storage_shape   = nullptr);
 
  void SelfReorder2Default();
  void SetShapeFromChunk() const;
  /*
   * This indicates whether an array is a view of another array (created by
   * reshape or slice). If an array is a view and the data is stored in
   * DNNL format, we need to convert the data to the default format when
   * data in the view is accessed.
   */
  inline bool IsView() const {
    // View only works on the default storage
    if (storage_type() != kDefaultStorage)
      return false;
    // If the array reuses memory, its shape may be different from the storage
    // shape. However, we shouldn't consider it as a view.
    if (reuse_)
      return false;
    return byte_offset_ > 0 || shape() != ptr_->storage_shape;
  }
 
  /* \brief Check whether the two arrays are the same array */
  inline bool IsSame(const NDArray& other) const {
    return ptr_ == other.ptr_ && shape_ == other.shape_ && byte_offset_ == other.byte_offset_ &&
           dtype_ == other.dtype_;
  }
 
  inline const mxnet::TShape& shape() const {
    return shape_;
  }
  inline const mxnet::TShape& storage_shape() const {
    CHECK(ptr_ != nullptr);
    CHECK_NE(storage_type(), kDefaultStorage)
        << "storage_shape() is not intended for kDefaultStorage.";
    return ptr_->storage_shape;
  }
 
  inline const mxnet::TShape& aux_shape(size_t index) const {
    CHECK_NE(storage_type(), kDefaultStorage) << "aux_shape() is not intended for kDefaultStorage.";
    return ptr_->aux_shapes[index];
  }
 
  /* \return the shapes of all aux data */
  const mxnet::ShapeVector& aux_shapes() const {
    CHECK_NE(storage_type(), kDefaultStorage)
        << "aux_shapes() is not intended for kDefaultStorage.";
    return ptr_->aux_shapes;
  }
 
  const std::vector<int>& aux_types() const {
    CHECK_NE(storage_type(), kDefaultStorage) << "aux_types() is not intended for kDefaultStorage.";
    return ptr_->aux_types;
  }
 
  inline void set_aux_shape(size_t index, const mxnet::TShape& shape) const {
    CHECK_NE(storage_type(), kDefaultStorage)
        << "set_aux_shape() is not intended for kDefaultStorage.";
    ptr_->set_aux_shape(index, shape);
  }
 
  inline const TBlob& data() const {
    if (storage_type() == kDefaultStorage)
      CheckAndAlloc();
    SetTBlob();
    return tblob_;
  }
  NDArray grad() const;
 
  inline TBlob aux_data(size_t i) const {
    auto stype = storage_type();
    TBlob res;
    auto shape = aux_shape(i);
    auto type  = aux_type(i);
    MSHADOW_TYPE_SWITCH(type, DType, {
      auto dptr = static_cast<DType*>(ptr_->aux_handles[i].dptr);
      CHECK(stype == kRowSparseStorage || stype == kCSRStorage)
          << "Unexpected storage type: " << stype;
      res = TBlob(dptr, shape, ptr_->aux_handles[i].ctx.dev_mask(), type);
    });
    return res;
  }
  inline Context ctx() const {
    CHECK(!is_none());
    return ptr_->shandle.ctx;
  }
  inline int dtype() const {
    return dtype_;
  }
  inline int aux_type(size_t i) const {
    CHECK(!is_none());
    return ptr_->aux_types[i];
  }
 
  inline NDArrayStorageType storage_type() const {
    return storage_type_;
  }
  inline bool is_none() const {
    return ptr_.get() == nullptr;
  }
  bool fresh_out_grad() const;
  void set_fresh_out_grad(bool state) const;
  inline bool storage_initialized() const {
    if (is_none())
      return false;
    auto stype = storage_type();
    CHECK_NE(stype, kDefaultStorage)
        << "storage_initialized() is not intended for kDefaultStorage.";
    if (stype == kRowSparseStorage) {
      CHECK_EQ(aux_shape(rowsparse::kIdx)[0], storage_shape()[0])
          << "inconsistent storage shape " << storage_shape() << " vs. aux shape "
          << aux_shape(rowsparse::kIdx);
      return aux_shape(rowsparse::kIdx).Size() != 0;
    } else if (stype == kCSRStorage) {
      CHECK_EQ(aux_shape(csr::kIdx)[0], storage_shape()[0])
          << "inconsistent storage shape " << storage_shape() << " vs. aux shape "
          << aux_shape(csr::kIdx);
      return aux_shape(csr::kIdx).Size() != 0;
    } else {
      LOG(FATAL) << "Unknown storage type";
    }
    return true;
  }
  inline Storage::Handle storage_handle() const {
    CHECK(!is_none());
    CHECK_EQ(storage_type(), kDefaultStorage);
    CheckAndAlloc();
    return ptr_->shandle;
  }
  void AssignStorageInfo(const std::string& profiler_scope, const std::string& name);
  void WaitToRead() const;
  void WaitToWrite() const;
  void StreamSync(int stream) const;
  inline Engine::VarHandle var() const {
    return ptr_->var;
  }
  inline size_t byte_offset() const {
    return byte_offset_;
  }
  inline size_t version() const {
    return var()->version();
  }
  void Save(dmlc::Stream* strm) const;
  bool LegacyLoad(dmlc::Stream* strm, const uint32_t magic);
  bool Load(dmlc::Stream* strm);
  NDArray& operator=(real_t scalar);
  NDArray& operator+=(const NDArray& src);
  NDArray& operator+=(const real_t& src);
  NDArray& operator-=(const NDArray& src);
  NDArray& operator-=(const real_t& src);
  NDArray& operator*=(const NDArray& src);
  NDArray& operator*=(const real_t& src);
  NDArray& operator/=(const NDArray& src);
  NDArray& operator/=(const real_t& src);
  NDArray Copy(Context ctx) const;
  void SyncCopyFromCPU(const void* data, size_t size) const;
 
  void SyncCopyFromNDArray(const NDArray& src, int i = -1, int j = -1);
 
  void SyncCopyToCPU(void* data, size_t size) const;
  void SyncCheckFormat(const bool full_check) const;
  NDArray Slice(index_t begin, index_t end) const;
  NDArray SliceWithRecord(index_t begin, index_t end);
  NDArray At(index_t idx) const;
  NDArray AtWithRecord(index_t idx);
  NDArray aux_ndarray(size_t i) const;
 
  NDArray data_ndarray() const;
 
  inline NDArray AsArray(const mxnet::TShape& shape, int dtype) const {
    CHECK_EQ(storage_type(), kDefaultStorage) << "AsArray is intended only for kDefaultStorage.";
    CHECK_GE(ptr_->shandle.size, shape.Size() * mshadow::mshadow_sizeof(dtype))
        << "NDArray.AsArray: target memory size is bigger";
    // We can't reuse memory in a view.
    CHECK(!IsView());
    NDArray ret = *this;
    ret.shape_  = shape;
    ret.dtype_  = dtype;
    ret.reuse_  = true;
    return ret;
  }
 
  inline void InitAsArray(const NDArray& src, const mxnet::TShape& shape, int dtype) {
    CHECK_EQ(src.storage_type(), kDefaultStorage)
        << "AsArray is intended only for kDefaultStorage.";
    CHECK_GE(src.ptr_->shandle.size, shape.Size() * mshadow::mshadow_sizeof(dtype))
        << "NDArray.AsArray: target memory size is bigger than what was allocated.";
    // We can't reuse memory in a view.
    CHECK(!src.IsView());
    *this  = src;
    shape_ = shape;
    dtype_ = dtype;
    reuse_ = true;
  }
 
  DLManagedTensor* ToDLPack() const;
 
  static NDArray FromDLPack(const DLManagedTensor* tensor, bool transient_handle);
 
  inline void SparseUpdateChunk(const NDArray& arr) const {
    CHECK(shape_ == arr.shape_) << "ndarray shape is different from the target";
    CHECK(dtype_ == arr.dtype_) << "ndarray dtype is different from the target";
    auto stype = arr.storage_type();
    CHECK(stype == kCSRStorage || stype == kRowSparseStorage)
        << "Only to be used with CSR and RSP storage types";
    // swap shandles between src and dst
    Storage::Handle shandle_dst = arr.ptr_->shandle;
    arr.ptr_->shandle           = ptr_->shandle;
    ptr_->shandle               = shandle_dst;
 
    ptr_->storage_shape = arr.ptr_->storage_shape;
    ptr_->storage_type  = arr.ptr_->storage_type;
    ptr_->ctx           = arr.ptr_->ctx;
 
    // swap aux_handles between src and dst
    size_t aux_idx = 0;
    CHECK(ptr_->aux_handles.size() == arr.ptr_->aux_handles.size())
        << "ndarray number of aux_handles is different from target";
    for (auto& aux_handle : arr.ptr_->aux_handles) {
      Storage::Handle aux_dst    = ptr_->aux_handles[aux_idx];
      ptr_->aux_handles[aux_idx] = aux_handle;
      aux_handle                 = aux_dst;
      aux_idx++;
    }
    ptr_->aux_types  = arr.ptr_->aux_types;
    ptr_->aux_shapes = arr.ptr_->aux_shapes;
  }
 
  NDArray Reshape(const mxnet::TShape& shape) const;
  NDArray ReshapeWithRecord(const mxnet::TShape& shape);
  NDArray Detach() const {
    NDArray ret(*this);
    ret.autograd_entry_        = nnvm::NodeEntry(nullptr);
    ret.deferredcompute_entry_ = nnvm::NodeEntry(nullptr);
    return ret;
  }
 
  nnvm::Symbol get_autograd_symbol() const;
  inline void CheckAndAlloc() const {
    CHECK_EQ(storage_type(), kDefaultStorage);
    ptr_->CheckAndAlloc();
  }
 
  void ReshapeAndAlloc(const mxnet::TShape& shape) {
    CHECK_EQ(storage_type(), kDefaultStorage);
    CHECK(!is_none());
    shape_ = shape;
    ptr_->CheckAndAlloc(shape.Size() * mshadow::mshadow_sizeof(dtype_));
  }
 
  /* !
   * \brief Alloc memory for non-default storage
   * aux_shape is only known at run time
   */
  inline void CheckAndAlloc(const mxnet::ShapeVector& aux_shapes) const {
    CHECK_NE(storage_type(), kDefaultStorage)
        << "CheckAndAlloc(aux_shapes) is not intended for kDefaultStorage";
    ptr_->CheckAndAlloc(shape_, aux_shapes, dtype_);
  }
  inline void CheckAndAllocData(const mxnet::TShape& storage_shape) const {
    CHECK_NE(storage_type(), kDefaultStorage)
        << "CheckAndAllocData is not intended for kDefaultStorage";
    ptr_->CheckAndAllocData(storage_shape, dtype_);
  }
  inline void CheckAndAllocAuxData(size_t i, const mxnet::TShape& aux_shape) const {
    CHECK_NE(storage_type(), kDefaultStorage)
        << "CheckAndAllocAuxData is not intended for kDefaultStorage";
    ptr_->CheckAndAllocAuxData(i, aux_shape);
  }
 
#if MXNET_USE_ONEDNN == 1
  /*
   * Create NDArray from dnnl memory.
   * dnnl_mem The dnnl memory to be managed.
   */
  explicit NDArray(const std::shared_ptr<dnnl::memory>& dnnl_mem);
  /*
   * Create NDArray from dnnl memory descriptor.
   * mem_pd The dnnl memory descriptor to be created.
   */
  explicit NDArray(const void* md);
  /*
   * Test if the data is stored in one of special DNNL formats.
   */
  bool IsDNNLData() const {
    return ptr_->IsDNNL();
  }
  /*
   * Test if the data is stored in one of default MXNet formats.
   */
  bool IsDefaultData() const {
    return ptr_->IsDefault();
  }
  /*
   * All functions below return a raw pointer to dnnl memory. Actually there
   * is a shared pointer that hold the memory either in NDArray or in DNNL
   * stream. As long as we call these functions inside an operator, the return
   * memory is always valid.
   */
 
  /*
   * This function returns dnnl::memory with the default primitive_desc.
   */
  const dnnl::memory* GetDNNLData() const;
  /*
   * This function returns dnnl::memory with the given primitive_desc
   * as long as the array size meets the required size in the given primitive_desc.
   */
  const dnnl::memory* GetDNNLData(const void* md) const;
  /*
   * This function returns dnnl::memory with the given primitive_desc.
   * The returned dnnl::memory will have the same physical layout as
   * the given primitive_desc.
   */
  const dnnl::memory* GetDNNLDataReorder(const void* md) const;
 
  /*
   * This function copies data from dnnl memory.
   */
  void CopyFrom(const dnnl::memory& mem);
  /*
   * This function allocates memory for array and creates dnnl memory
   * with the specified format.
   */
  dnnl::memory* CreateDNNLData(const void* md);
 
  /*
   * These are the async version of the methods above.
   * It changes the layout of this NDArray, but it happens after all accesses to
   * the array are complete.
   */
  void Reorder2DefaultAsync() const;
  void DNNLDataReorderAsync(const void* md) const;
 
  /*
   * This creates a new NDArray with the reordered data.
   * It doesn't affect the data of the original NDArray.
   */
  NDArray Reorder2Default() const;
 
  /*
   * This creates a new NDArray using f32 with the reordered data.
   * It doesn't affect the data of the original NDArray.
   */
  NDArray Reorder2DefaultFloatFormat() const;
 
  void InvalidateDNNLData();
 
  /*
   * This function is used inside operators to reshape an array.
   * It doesn't change the layout of the original array and allocate memory from
   * the temporary buffer. The returned array is only valid inside the current
   * invocation of this operator.
   * This is different from Reshape. Reshape will cause data in the array to be
   * converted to the default layout and allocate memory from malloc directly,
   * which can be expensive.
   * It's used by FullyConnected right now.
   */
  NDArray DNNLDataReshape(const mxnet::TShape& shape) const;
 
  void UpdateDNNLMemDesc(const void* desc);
#endif
 
  static void Save(dmlc::Stream* fo,
                   const std::vector<NDArray>& data,
                   const std::vector<std::string>& names);
  static void Load(dmlc::Stream* fi, std::vector<NDArray>* data, std::vector<std::string>* keys);
 
 private:
  friend class Imperative;
  // shandle is used to store the actual values in the NDArray
  // aux_handles store the aux data(such as indices) if it's needed by non-default storage.
  struct Chunk {
    Storage::Handle shandle;
    std::vector<Storage::Handle> aux_handles;
 
#if MXNET_USE_ONEDNN == 1
 
    std::shared_ptr<DNNLMemory> dnnl_mem_;
#endif
 
    Engine::VarHandle var;
    bool static_data;
    bool delay_alloc;
    // the type of the storage. The storage_type is never kUndefinedStorage once the chunk
    // is constructed.
    NDArrayStorageType storage_type = kDefaultStorage;
    std::vector<int> aux_types;
    // context of data
    Context ctx;
    // The shape of the chunk data.
    // This might not be the same shape as the NDArray, since the storage may be sparse.
    // The default value for storage_shape is {0} when an empty non-default NDArray is created.
    mxnet::TShape storage_shape;
    // The shape of aux data. The default value for the shape depends on the type of storage.
    // If aux_shapes[i].Size() is zero, aux data i is empty.
    mxnet::ShapeVector aux_shapes;
    std::shared_ptr<Storage> storage_ref_;
    std::weak_ptr<Engine> engine_ref_;
 
    Chunk()
        : static_data(true),
          delay_alloc(false),
          storage_ref_(Storage::_GetSharedRef()),
          engine_ref_(Engine::_GetSharedRef()) {}
 
    Chunk(mxnet::TShape shape, Context ctx_, bool delay_alloc_, int dtype)
        : static_data(false),
          delay_alloc(true),
          ctx(ctx_),
          storage_ref_(Storage::_GetSharedRef()),
          engine_ref_(Engine::_GetSharedRef()) {
      storage_shape = shape;
      if (shape_is_known(storage_shape)) {
        shandle.size = shape.Size() * mshadow::mshadow_sizeof(dtype);
      }
      var         = Engine::Get()->NewVariable();
      shandle.ctx = ctx_;
      if (!delay_alloc_) {
        this->CheckAndAlloc();
      }
    }
 
    Chunk(const TBlob& data, int dev_id)
        : static_data(true),
          delay_alloc(false),
          storage_ref_(Storage::_GetSharedRef()),
          engine_ref_(Engine::_GetSharedRef()) {
      CHECK(storage_type == kDefaultStorage);
      var = Engine::Get()->NewVariable();
      if (data.dev_mask() == cpu::kDevMask) {
        ctx = Context::CPU();
      } else {
        CHECK_EQ(data.dev_mask(), gpu::kDevMask);
        ctx = Context::GPU(dev_id);
      }
      // init shandle
      shandle.ctx   = ctx;
      shandle.dptr  = data.dptr_;
      shandle.size  = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
      storage_shape = data.shape_;
    }
 
    Chunk(int shared_pid, int shared_id, const mxnet::TShape& shape, int dtype)
        : static_data(false),
          delay_alloc(false),
          storage_ref_(Storage::_GetSharedRef()),
          engine_ref_(Engine::_GetSharedRef()) {
      var                = Engine::Get()->NewVariable();
      ctx                = Context::CPUShared(0);
      shandle.size       = shape.Size() * mshadow::mshadow_sizeof(dtype);
      shandle.ctx        = ctx;
      shandle.shared_pid = shared_pid;
      shandle.shared_id  = shared_id;
      Storage::Get()->Alloc(&shandle);
      storage_shape = shape;
    }
    // Constructor for a non-default storage chunk
    Chunk(NDArrayStorageType storage_type_,
          const mxnet::TShape& storage_shape_,
          Context ctx_,
          bool delay_alloc_,
          int dtype,
          const std::vector<int>& aux_types_,
          const mxnet::ShapeVector& aux_shapes_)
        : static_data(false),
          delay_alloc(delay_alloc_),
          storage_type(storage_type_),
          aux_types(aux_types_),
          ctx(ctx_),
          storage_shape(storage_shape_),
          aux_shapes(aux_shapes_),
          storage_ref_(Storage::_GetSharedRef()),
          engine_ref_(Engine::_GetSharedRef()) {
      shandle.ctx = ctx;
      var         = Engine::Get()->NewVariable();
      // aux_handles always reflect the correct number of aux data
      for (size_t i = 0; i < aux_shapes.size(); i++) {
        CheckAndAllocAuxData(i, aux_shapes[i]);
        // this line is needed in case when aux_shapes[i].Size() = 0
        // aux_handles[i] will not be updated and take only default value.
        aux_handles[i].ctx = ctx;
      }
      if (!delay_alloc) {
        CheckAndAllocData(storage_shape, dtype);
      }
    }
 
    Chunk(const NDArrayStorageType storage_type_,
          const TBlob& data,
          const std::vector<TBlob>& aux_data,
          int dev_id)
        : static_data(true),
          delay_alloc(false),
          storage_type(storage_type_),
          storage_ref_(Storage::_GetSharedRef()),
          engine_ref_(Engine::_GetSharedRef()) {
      using namespace mshadow;
      CHECK_NE(storage_type, kDefaultStorage);
      // init var
      var = Engine::Get()->NewVariable();
      // init ctx
      if (data.dev_mask() == cpu::kDevMask) {
        ctx = Context::CPU();
      } else {
        CHECK_EQ(data.dev_mask(), gpu::kDevMask);
        ctx = Context::GPU(dev_id);
      }
      // init shandle
      shandle.ctx   = ctx;
      shandle.dptr  = data.dptr_;
      shandle.size  = data.shape_.Size() * mshadow_sizeof(data.type_flag_);
      storage_shape = data.shape_;
      // init aux handles
      for (const auto& aux : aux_data) {
        Storage::Handle aux_handle;
        aux_handle.ctx  = ctx;
        aux_handle.dptr = aux.dptr_;
        aux_handle.size = aux.shape_.Size() * mshadow_sizeof(aux.type_flag_);
        aux_handles.push_back(aux_handle);
        aux_types.emplace_back(aux.type_flag_);
        aux_shapes.emplace_back(aux.shape_);
      }
    }
 
    inline void set_aux_shape(const size_t i, const mxnet::TShape& shape) {
      aux_shapes[i] = shape;
      if (storage_shape.ndim() >= 0) {
        if (storage_type == kRowSparseStorage && i == rowsparse::kIdx) {
          storage_shape[0] = shape[0];
        } else if (storage_type == kCSRStorage && i == csr::kIdx) {
          storage_shape[0] = shape[0];
        }
      }
    }
 
    inline void CheckAndAlloc(void) {
      if (delay_alloc) {
        Storage::Get()->Alloc(&shandle);
#if MXNET_USE_ONEDNN == 1
        dnnl_mem_ = nullptr;
#endif
        delay_alloc = false;
      }
    }
 
    // size is the number of bytes
    void CheckAndAlloc(uint64_t dbytes) {
      CHECK_EQ(kDefaultStorage, storage_type)
          << "CheckAndAlloc(dbytes) is only intended for kDefaultStorage";
      dbytes = std::max(dbytes, static_cast<uint64_t>(shandle.size));
      if (delay_alloc) {
        shandle.size = dbytes;
        Storage::Get()->Alloc(&shandle);
#if MXNET_USE_ONEDNN == 1
        dnnl_mem_ = nullptr;
#endif
        delay_alloc = false;
      } else if (shandle.size < dbytes) {
        // free storage
        Storage::Get()->Free(shandle);
        // init storage
        shandle.size = dbytes;
        Storage::Get()->Alloc(&shandle);
#if MXNET_USE_ONEDNN == 1
        dnnl_mem_ = nullptr;
#endif
      }
    }
    void Init(const mxnet::TShape& shape, int dtype) {
      auto size     = shape.Size();
      storage_shape = shape;
      shandle.size  = size * mshadow::mshadow_sizeof(dtype);
      this->CheckAndAlloc();
    }
    inline void CheckAndAlloc(const mxnet::TShape& shape,
                              const mxnet::ShapeVector& aux_shapes,
                              int dtype) {
      // calculate size, perform allocation
      if (kRowSparseStorage == storage_type) {
        // For row sparse, aux_shape indicates the number of rows to allocate
        auto aux_shape = aux_shapes[rowsparse::kIdx];
        CheckAndAllocAuxData(rowsparse::kIdx, aux_shape);
        mxnet::TShape storage_shape(shape);
        storage_shape[0] = aux_shape[0];
        CheckAndAllocData(storage_shape, dtype);
      } else if (kCSRStorage == storage_type) {
        CheckAndAllocAuxData(csr::kIndPtr, aux_shapes[csr::kIndPtr]);
        CheckAndAllocAuxData(csr::kIdx, aux_shapes[csr::kIdx]);
        CheckAndAllocData(aux_shapes[csr::kIdx], dtype);
      } else {
        LOG(FATAL) << "Storage type " << storage_type << " not implemented for CheckAndAlloc";
      }
    }
    // create storage handle for data based on shape and dtype, assuming ctx is set
    // storage shape is also updated
    // if data is already allocated, try reuse the storage. Otherwise, free the current one
    // and allocate new storage
    void CheckAndAllocData(const mxnet::TShape& shape, int dtype);
 
#if MXNET_USE_ONEDNN == 1
    // Have DNNL memory reference to the data in the default storage
    // or create memory for DNNL.
    void SetDNNLMem(const mxnet::TShape& shape, int dtype);
    // If the data is stored in DNNL layout, we reorder data in dnnl_mem_ and
    // save the result in shandle.
    void Reorder2Default();
    // Reroder data to a specified layout.
    void DNNLDataReorder(const void* md);
    bool IsDNNL() const;
    bool IsDefault() const;
#endif
 
    // create storage handle for aux data based on shape
    // this function assumes ctx, aux shapes and aux types are set
    // aux shape is also updated
    // if aux data is already allocated, try reuse the storage. Otherwise, free the current one
    // and allocate new storage
    inline void CheckAndAllocAuxData(size_t i, const mxnet::TShape& shape) {
      CHECK_EQ(shape.ndim(), 1) << "shape must be 1D in CheckAndAllocAuxData";
      CHECK_NE(storage_type, kUndefinedStorage)
          << "storage type cannot be kUndefinedStorage in CheckAndAllocAuxData";
      CHECK_NE(storage_type, kDefaultStorage)
          << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
      if (aux_handles.size() <= i) {
        aux_handles.resize(i + 1);
      }
      size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
      if (aux_handles[i].size < aux_bytes) {
        // free storage
        Storage::Get()->Free(aux_handles[i]);
        // init aux storage
        aux_handles[i] = Storage::Get()->Alloc(aux_bytes, ctx);
      }
      // init shape
      set_aux_shape(i, shape);
    }
    ~Chunk();
  };  // struct Chunk
 
  inline void Init(const NDArrayStorageType stype, const mxnet::TShape& shape, int dtype) {
    shape_          = shape;
    dtype_          = dtype;
    storage_type_   = stype;
    reuse_          = false;
    byte_offset_    = 0;
    autograd_entry_ = nnvm::NodeEntry(nullptr);
  }
 
  void SetTBlob() const;
 
  std::shared_ptr<Chunk> ptr_{nullptr};
  mutable mxnet::TShape shape_;
  size_t byte_offset_ = 0;
  int dtype_ = -1;
  bool reuse_ = false;
  NDArrayStorageType storage_type_ = kUndefinedStorage;
  nnvm::NodeEntry autograd_entry_;
  nnvm::NodeEntry deferredcompute_entry_;
  mutable TBlob tblob_;
};  // class NDArray
 
size_t num_aux_data(NDArrayStorageType stype);
 
void CopyFromTo(const NDArray& from, const NDArray* to, int priority = 0);
 
void CopyFromTo(const NDArray& from, const NDArray& to, int priority = 0, bool is_opr = false);
 
void ElementwiseSum(const std::vector<NDArray>& source, NDArray* out, int priority = 0);
 
NDArray operator+(const NDArray& lhs, const NDArray& rhs);
NDArray operator+(const NDArray& lhs, const real_t& rhs);
NDArray operator-(const NDArray& lhs, const NDArray& rhs);
NDArray operator-(const NDArray& lhs, const real_t& rhs);
NDArray operator*(const NDArray& lhs, const NDArray& rhs);
NDArray operator*(const NDArray& lhs, const real_t& rhs);
NDArray operator/(const NDArray& lhs, const NDArray& rhs);
NDArray operator/(const NDArray& lhs, const real_t& rhs);
 
void RandomSeed(uint32_t seed);
void RandomSeed(Context ctx, uint32_t seed);
void SampleUniform(real_t begin, real_t end, NDArray* out);
void SampleGaussian(real_t mu, real_t sigma, NDArray* out);
void SampleGamma(real_t alpha, real_t beta, NDArray* out);
void SampleExponential(real_t lambda, NDArray* out);
void SamplePoisson(real_t lambda, NDArray* out);
void SampleNegBinomial(int32_t k, real_t p, NDArray* out);
void SampleGenNegBinomial(real_t mu, real_t alpha, NDArray* out);
 
//--------------------------------------------------------------
// The following part are API Registration of NDArray functions.
//--------------------------------------------------------------
 
typedef std::function<void(NDArray** used_vars,
                           real_t* scalars,
                           NDArray** mutate_vars,
                           int num_params,
                           char** param_keys,
                           char** param_vals)>
    NDArrayAPIFunction;
enum NDArrayFunctionTypeMask {
  kNDArrayArgBeforeScalar = 1,
  kScalarArgBeforeNDArray = 1 << 1,
  kAcceptEmptyMutateTarget = 1 << 2
};
struct NDArrayFunctionReg
    : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg, NDArrayAPIFunction> {
  unsigned num_use_vars;
  unsigned num_mutate_vars;
  unsigned num_scalars;
  int type_mask;
  NDArrayFunctionReg() : num_use_vars(0), num_mutate_vars(0), num_scalars(0), type_mask(0) {}
  inline NDArrayFunctionReg& set_function(void (*fsetvalue)(const real_t& rhs, NDArray* out)) {
    body            = [fsetvalue](NDArray** used_vars,
                       real_t* s,
                       NDArray** mutate_vars,
                       int num_params,
                       char** param_keys,
                       char** param_vals) { (*fsetvalue)(s[0], mutate_vars[0]); };
    num_mutate_vars = 1;
    num_scalars     = 1;
    this->add_argument("src", "real_t", "Source input to the function.");
    return *this;
  }
  inline NDArrayFunctionReg& set_function(
      void (*fternary)(const NDArray& lhs, const NDArray& mhs, const NDArray& rhs, NDArray* out)) {
    body = [fternary](NDArray** used_vars,
                      real_t* s,
                      NDArray** mutate_vars,
                      int num_params,
                      char** param_keys,
                      char** param_vals) {
      (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
    };
    num_use_vars    = 3;
    num_mutate_vars = 1;
    type_mask       = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("mhs", "NDArray", "Middle operand to the function.");
    this->add_argument("rhs", "NDArray", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg& set_function(void (*fbinary)(const NDArray& lhs,
                                                          const NDArray& rhs,
                                                          NDArray* out)) {
    body = [fbinary](NDArray** used_vars,
                     real_t* s,
                     NDArray** mutate_vars,
                     int num_params,
                     char** param_keys,
                     char** param_vals) {
      (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
    };
    num_use_vars    = 2;
    num_mutate_vars = 1;
    type_mask       = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("rhs", "NDArray", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg& set_function(void (*fscalar)(const NDArray& lhs,
                                                          const real_t& rhs,
                                                          NDArray* out)) {
    body            = [fscalar](NDArray** used_vars,
                     real_t* s,
                     NDArray** mutate_vars,
                     int num_params,
                     char** param_keys,
                     char** param_vals) { (*fscalar)(*used_vars[0], s[0], mutate_vars[0]); };
    num_use_vars    = 1;
    num_mutate_vars = 1;
    num_scalars     = 1;
    type_mask       = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("lhs", "NDArray", "Left operand to the function.");
    this->add_argument("rhs", "real_t", "Right operand to the function.");
    return *this;
  }
  inline NDArrayFunctionReg& set_function(void (*funary)(const NDArray& src, NDArray* out)) {
    body            = [funary](NDArray** used_vars,
                    real_t* s,
                    NDArray** mutate_vars,
                    int num_params,
                    char** param_keys,
                    char** param_vals) { (*funary)(*used_vars[0], mutate_vars[0]); };
    num_use_vars    = 1;
    num_mutate_vars = 1;
    type_mask       = kNDArrayArgBeforeScalar | kAcceptEmptyMutateTarget;
    this->add_argument("src", "NDArray", "Source input to the function.");
    return *this;
  }
  inline NDArrayFunctionReg& set_function(
      void (*fgeneric)(NDArray** used_vars,
                       real_t* s,
                       NDArray** mutate_vars,
                       const std::map<std::string, std::string>& param)) {
    body = [fgeneric](NDArray** used_vars,
                      real_t* s,
                      NDArray** mutate_vars,
                      int num_params,
                      char** param_keys,
                      char** param_vals) {
      std::map<std::string, std::string> param;
      for (int i = 0; i < num_params; ++i) {
        param[param_keys[i]] = param_vals[i];
      }
      fgeneric(used_vars, s, mutate_vars, param);
    };
    return *this;
  }
  inline NDArrayFunctionReg& set_num_use_vars(unsigned n) {
    num_use_vars = n;
    return *this;
  }
  inline NDArrayFunctionReg& set_num_mutate_vars(unsigned n) {
    num_mutate_vars = n;
    return *this;
  }
  inline NDArrayFunctionReg& set_num_scalars(unsigned n) {
    num_scalars = n;
    return *this;
  }
  inline NDArrayFunctionReg& set_type_mask(int tmask) {
    type_mask = tmask;
    return *this;
  }
};  // NDArrayFunctionReg
 
#define MXNET_REGISTER_NDARRAY_FUN(name) \
  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)
 
}  // namespace mxnet
 
namespace dmlc {
DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
}  // namespace dmlc
#endif  // MXNET_NDARRAY_H_