[MXNET-1185] [WIP] Support large array in several operators

R-package/src/base.h

@@ -354,8 +354,8 @@ inline std::vector<std::string> SafeGetListNames(const Rcpp::List& src) {
  * \param rshape The dimension in R
  * \return A internal vector representation of shapes in mxnet.
  */
-inline std::vector<mx_uint> Dim2InternalShape(const Rcpp::Dimension &rshape) {
-  std::vector<mx_uint> shape(rshape.size());
+inline std::vector<dim_t> Dim2InternalShape(const Rcpp::Dimension &rshape) {
+  std::vector<dim_t> shape(rshape.size());
   for (size_t i = 0; i < rshape.size(); ++i) {
     shape[rshape.size() - i - 1] = rshape[i];
   }

R-package/src/ndarray.cc

@@ -180,7 +180,7 @@ Rcpp::RObject NDArrayPacker::CreateNDArrayPacker() {
 
 Rcpp::Dimension NDArray::dim() const {
   mx_uint ndim;
-  const mx_uint *pshape;
+  const dim_t *pshape;
   MX_CALL(MXNDArrayGetShape(
       ptr_->handle, &ndim, &pshape));
   Rcpp::IntegerVector dat(pshape, pshape + ndim);
@@ -190,7 +190,7 @@ Rcpp::Dimension NDArray::dim() const {
 }
 
 NDArray NDArray::Clone() const {
-  std::vector<mx_uint> shape = Dim2InternalShape(this->dim());
+  std::vector<dim_t> shape = Dim2InternalShape(this->dim());
   Context ctx = this->ctx();
   NDArrayHandle handle;
   MX_CALL(MXNDArrayCreate(dmlc::BeginPtr(shape),
@@ -276,7 +276,7 @@ Rcpp::List NDArray::Load(const std::string& filename) {
 NDArray::RObjectType NDArray::Empty(
     const Rcpp::Dimension& rshape,
     const Context::RObjectType& rctx) {
-  std::vector<mx_uint> shape = Dim2InternalShape(rshape);
+  std::vector<dim_t> shape = Dim2InternalShape(rshape);
   Context ctx(rctx);
   NDArrayHandle handle;
   MX_CALL(MXNDArrayCreate(dmlc::BeginPtr(shape),

R-package/src/symbol.cc

@@ -168,7 +168,7 @@ Symbol::RObjectType Symbol::GetOutput(mx_uint index) const {
 // helper function to convert shape into Rcpp vector
 inline Rcpp::List BuildShapeData(mx_uint shape_size,
                                  const mx_uint *shape_ndim,
-                                 const mx_uint **shape_data,
+                                 const dim_t **shape_data,
                                  const std::vector<std::string> &names) {
   Rcpp::List ret(shape_size);
   for (mx_uint i = 0; i < shape_size; ++i) {
@@ -185,26 +185,26 @@ SEXP Symbol::InferShape(const Rcpp::List& kwargs) const {
       << "Need to pass parameters in key=value style.\n";
   std::vector<std::string> keys = kwargs.names();
   std::vector<mx_uint> arg_ind_ptr(1, 0);
-  std::vector<mx_uint> arg_shape_data;
+  std::vector<dim_t> arg_shape_data;
 
   for (size_t i = 0; i < kwargs.size(); ++i) {
     RCHECK(keys[i].length() != 0)
       << "Need to pass parameters in key=value style.\n";
-    std::vector<mx_uint> dim = Dim2InternalShape(kwargs[i]);
+    std::vector<dim_t> dim = Dim2InternalShape(kwargs[i]);
     arg_shape_data.insert(arg_shape_data.end(), dim.begin(), dim.end());
     arg_ind_ptr.push_back(static_cast<mx_uint>(arg_shape_data.size()));
   }
   std::vector<const char*> c_keys = CKeys(keys);
 
   mx_uint in_shape_size;
   const mx_uint *in_shape_ndim;
-  const mx_uint **in_shape_data;
+  const dim_t **in_shape_data;
   mx_uint out_shape_size;
   const mx_uint *out_shape_ndim;
-  const mx_uint **out_shape_data;
+  const dim_t **out_shape_data;
   mx_uint aux_shape_size;
   const mx_uint *aux_shape_ndim;
-  const mx_uint **aux_shape_data;
+  const dim_t **aux_shape_data;
   int complete;
 
   MX_CALL(MXSymbolInferShape(

cpp-package/include/mxnet-cpp/base.h

@@ -34,7 +34,7 @@
 namespace mxnet {
 namespace cpp {
 
-typedef unsigned index_t;
+typedef int64_t index_t;
 
 enum OpReqType {
   /*! \brief no operation, do not write anything */

cpp-package/include/mxnet-cpp/initializer.h

@@ -167,7 +167,7 @@ class Xavier : public Initializer {
     Shape shape(arr->GetShape());
     float hw_scale = 1.0f;
     if (shape.ndim() > 2) {
-      for (size_t i = 2; i < shape.ndim(); ++i) {
+      for (index_t i = 2; i < shape.ndim(); ++i) {
         hw_scale *= shape[i];
       }
     }

cpp-package/include/mxnet-cpp/ndarray.h

@@ -134,7 +134,7 @@ class NDArray {
   * \param constext context of NDArray
   * \param delay_alloc whether delay the allocation
   */
-  NDArray(const std::vector<mx_uint> &shape, const Context &context,
+  NDArray(const std::vector<index_t> &shape, const Context &context,
           bool delay_alloc = true);
   /*!
   * \brief construct a new dynamic NDArray
@@ -444,7 +444,7 @@ class NDArray {
   /*!
   * \return the shape of current NDArray, in the form of mx_uint vector
   */
-  std::vector<mx_uint> GetShape() const;
+  std::vector<index_t> GetShape() const;
   /*!
   * \return the data type of current NDArray
   */

cpp-package/include/mxnet-cpp/ndarray.hpp

@@ -46,7 +46,7 @@ inline NDArray::NDArray() {
 inline NDArray::NDArray(const NDArrayHandle &handle) {
   blob_ptr_ = std::make_shared<NDBlob>(handle);
 }
-inline NDArray::NDArray(const std::vector<mx_uint> &shape, const Context &context,
+inline NDArray::NDArray(const std::vector<index_t> &shape, const Context &context,
                         bool delay_alloc) {
   NDArrayHandle handle;
   CHECK_EQ(MXNDArrayCreate(shape.data(), shape.size(), context.GetDeviceType(),
@@ -396,11 +396,11 @@ inline size_t NDArray::Size() const {
   return ret;
 }
 
-inline std::vector<mx_uint> NDArray::GetShape() const {
-  const mx_uint *out_pdata;
+inline std::vector<index_t> NDArray::GetShape() const {
+  const index_t *out_pdata;
   mx_uint out_dim;
   MXNDArrayGetShape(blob_ptr_->handle_, &out_dim, &out_pdata);
-  std::vector<mx_uint> ret;
+  std::vector<index_t> ret;
   for (mx_uint i = 0; i < out_dim; ++i) {
     ret.push_back(out_pdata[i]);
   }

cpp-package/include/mxnet-cpp/symbol.h

@@ -161,10 +161,10 @@ class Symbol {
   * \param aux_shapes use to store the infered shapes of auxiliary states
   */
   void InferShape(
-      const std::map<std::string, std::vector<mx_uint> > &arg_shapes,
-      std::vector<std::vector<mx_uint> > *in_shape,
-      std::vector<std::vector<mx_uint> > *aux_shape,
-      std::vector<std::vector<mx_uint> > *out_shape) const;
+      const std::map<std::string, std::vector<index_t> > &arg_shapes,
+      std::vector<std::vector<index_t> > *in_shape,
+      std::vector<std::vector<index_t> > *aux_shape,
+      std::vector<std::vector<index_t> > *out_shape) const;
   /*!
   * \brief List the arguments names.
   *

cpp-package/include/mxnet-cpp/symbol.hpp

@@ -181,14 +181,14 @@ inline std::string Symbol::GetName() const {
 }
 
 inline void Symbol::InferShape(
-    const std::map<std::string, std::vector<mx_uint> > &arg_shapes,
-    std::vector<std::vector<mx_uint> > *in_shape,
-    std::vector<std::vector<mx_uint> > *aux_shape,
-    std::vector<std::vector<mx_uint> > *out_shape) const {
+    const std::map<std::string, std::vector<index_t> > &arg_shapes,
+    std::vector<std::vector<index_t> > *in_shape,
+    std::vector<std::vector<index_t> > *aux_shape,
+    std::vector<std::vector<index_t> > *out_shape) const {
 
   std::vector<const char *> keys;
   std::vector<mx_uint> arg_ind_ptr;
-  std::vector<mx_uint> arg_shape_data;
+  std::vector<index_t> arg_shape_data;
 
   for (const auto &arg : arg_shapes) {
     keys.push_back(arg.first.c_str());
@@ -201,13 +201,13 @@ inline void Symbol::InferShape(
 
   mx_uint in_shape_size;
   const mx_uint *in_shape_ndim;
-  const mx_uint **in_shape_data;
+  const index_t **in_shape_data;
   mx_uint out_shape_size;
   const mx_uint *out_shape_ndim;
-  const mx_uint **out_shape_data;
+  const index_t **out_shape_data;
   mx_uint aux_shape_size;
   const mx_uint *aux_shape_ndim;
-  const mx_uint **aux_shape_data;
+  const index_t **aux_shape_data;
   int complete;
 
   CHECK_EQ(MXSymbolInferShape(GetHandle(), keys.size(), keys.data(),
@@ -220,19 +220,19 @@ inline void Symbol::InferShape(
 
   if (complete) {
     for (mx_uint i = 0; i < in_shape_size; ++i) {
-      in_shape->push_back(std::vector<mx_uint>());
+      in_shape->push_back(std::vector<index_t>());
       for (mx_uint j = 0; j < in_shape_ndim[i]; ++j) {
         (*in_shape)[i].push_back(in_shape_data[i][j]);
       }
     }
     for (mx_uint i = 0; i < aux_shape_size; ++i) {
-      aux_shape->push_back(std::vector<mx_uint>());
+      aux_shape->push_back(std::vector<index_t>());
       for (mx_uint j = 0; j < aux_shape_ndim[i]; ++j) {
         (*aux_shape)[i].push_back(aux_shape_data[i][j]);
       }
     }
     for (mx_uint i = 0; i < out_shape_size; ++i) {
-      out_shape->push_back(std::vector<mx_uint>());
+      out_shape->push_back(std::vector<index_t>());
       for (mx_uint j = 0; j < out_shape_ndim[i]; ++j) {
         (*out_shape)[i].push_back(out_shape_data[i][j]);
       }
@@ -250,8 +250,8 @@ inline void Symbol::InferExecutorArrays(
     const std::map<std::string, NDArray> &aux_map) const {
 
   const auto arg_name_list = ListArguments();
-  std::vector<std::vector<mx_uint> > in_shapes, aux_shapes, out_shapes;
-  std::map<std::string, std::vector<mx_uint> > arg_shapes;
+  std::vector<std::vector<index_t> > in_shapes, aux_shapes, out_shapes;
+  std::map<std::string, std::vector<index_t> > arg_shapes;
 
   for (const auto &arg_name : arg_name_list) {
     auto iter = args_map.find(arg_name);
@@ -307,8 +307,8 @@ inline void Symbol::InferArgsMap(
     const std::map<std::string, NDArray> &known_args) const {
 
   const auto arg_name_list = ListArguments();
-  std::vector<std::vector<mx_uint> > in_shapes, aux_shapes, out_shapes;
-  std::map<std::string, std::vector<mx_uint> > arg_shapes;
+  std::vector<std::vector<index_t> > in_shapes, aux_shapes, out_shapes;
+  std::map<std::string, std::vector<index_t> > arg_shapes;
 
   for (const auto &arg_name : arg_name_list) {
     auto iter = known_args.find(arg_name);

include/mxnet/c_api.h

@@ -487,7 +487,7 @@ MXNET_DLL int MXNDArrayCreateNone(NDArrayHandle *out);
  * \param out the returning handle
  * \return 0 when success, -1 when failure happens
  */
-MXNET_DLL int MXNDArrayCreate(const mx_uint *shape,
+MXNET_DLL int MXNDArrayCreate(const dim_t *shape,
                               mx_uint ndim,
                               int dev_type,
                               int dev_id,
@@ -506,7 +506,7 @@ MXNET_DLL int MXNDArrayCreate(const mx_uint *shape,
  * \param out the returning handle
  * \return 0 when success, -1 when failure happens
  */
-MXNET_DLL int MXNDArrayCreateEx(const mx_uint *shape,
+MXNET_DLL int MXNDArrayCreateEx(const dim_t *shape,
                               mx_uint ndim,
                               int dev_type,
                               int dev_id,
@@ -533,7 +533,7 @@ MXNET_DLL int MXNDArrayCreateEx(const mx_uint *shape,
  * \return 0 when success, -1 when failure happens
  */
 MXNET_DLL int MXNDArrayCreateSparseEx(int storage_type,
-                                      const mx_uint *shape,
+                                      const dim_t *shape,
                                       mx_uint ndim,
                                       int dev_type,
                                       int dev_id,
@@ -542,7 +542,7 @@ MXNET_DLL int MXNDArrayCreateSparseEx(int storage_type,
                                       mx_uint num_aux,
                                       int *aux_type,
                                       mx_uint *aux_ndims,
-                                      const mx_uint *aux_shape,
+                                      const dim_t *aux_shape,
                                       NDArrayHandle *out);
 
 /*!
@@ -650,7 +650,7 @@ MXNET_DLL int MXNDArraySyncCopyToCPU(NDArrayHandle handle,
  */
 MXNET_DLL int MXNDArraySyncCopyFromNDArray(NDArrayHandle handle_dst,
                                            const NDArrayHandle handle_src,
-                                           const int i);
+                                           const dim_t i);
 
 /*!
  * \brief check whether the NDArray format is valid
@@ -693,8 +693,8 @@ MXNET_DLL int MXNDArrayFree(NDArrayHandle handle);
  * \return 0 when success, -1 when failure happens
  */
 MXNET_DLL int MXNDArraySlice(NDArrayHandle handle,
-                             mx_uint slice_begin,
-                             mx_uint slice_end,
+                             dim_t slice_begin,
+                             dim_t slice_end,
                              NDArrayHandle *out);
 
 /*!
@@ -705,7 +705,7 @@ MXNET_DLL int MXNDArraySlice(NDArrayHandle handle,
  * \return 0 when success, -1 when failure happens
  */
 MXNET_DLL int MXNDArrayAt(NDArrayHandle handle,
-                          mx_uint idx,
+                          dim_t idx,
                           NDArrayHandle *out);
 
 /*!
@@ -749,7 +749,7 @@ MXNET_DLL int MXNDArrayReshape64(NDArrayHandle handle,
  */
 MXNET_DLL int MXNDArrayGetShape(NDArrayHandle handle,
                                 mx_uint *out_dim,
-                                const mx_uint **out_pdata);
+                                const dim_t **out_pdata);
 /*!
  * \brief get the content of the data in NDArray
  * \param handle the handle to the ndarray
@@ -1466,16 +1466,16 @@ MXNET_DLL int MXSymbolInferShape(SymbolHandle sym,
                                  mx_uint num_args,
                                  const char** keys,
                                  const mx_uint *arg_ind_ptr,
-                                 const mx_uint *arg_shape_data,
+                                 const dim_t *arg_shape_data,
                                  mx_uint *in_shape_size,
                                  const mx_uint **in_shape_ndim,
-                                 const mx_uint ***in_shape_data,
+                                 const dim_t ***in_shape_data,
                                  mx_uint *out_shape_size,
                                  const mx_uint **out_shape_ndim,
-                                 const mx_uint ***out_shape_data,
+                                 const dim_t ***out_shape_data,
                                  mx_uint *aux_shape_size,
                                  const mx_uint **aux_shape_ndim,
-                                 const mx_uint ***aux_shape_data,
+                                 const dim_t ***aux_shape_data,
                                  int *complete);
 /*!
  * \brief partially infer shape of unknown input shapes given the known one.
@@ -1505,16 +1505,16 @@ MXNET_DLL int MXSymbolInferShapePartial(SymbolHandle sym,
                                         mx_uint num_args,
                                         const char** keys,
                                         const mx_uint *arg_ind_ptr,
-                                        const mx_uint *arg_shape_data,
+                                        const dim_t *arg_shape_data,
                                         mx_uint *in_shape_size,
                                         const mx_uint **in_shape_ndim,
-                                        const mx_uint ***in_shape_data,
+                                        const dim_t ***in_shape_data,
                                         mx_uint *out_shape_size,
                                         const mx_uint **out_shape_ndim,
-                                        const mx_uint ***out_shape_data,
+                                        const dim_t ***out_shape_data,
                                         mx_uint *aux_shape_size,
                                         const mx_uint **aux_shape_ndim,
-                                        const mx_uint ***aux_shape_data,
+                                        const dim_t ***aux_shape_data,
                                         int *complete);
 
 /*!

include/mxnet/ndarray.h

@@ -488,7 +488,7 @@ class NDArray {
   /*!
    * \brief Copy from src.data()/aux_data(i) to this->data()/aux_data(j)
    */
-  void SyncCopyFromNDArray(const NDArray &src, int i = -1, int j = -1);
+  void SyncCopyFromNDArray(const NDArray &src, index_t i = -1, index_t j = -1);
 
   /*!
    * \brief Do a synchronize copy to a continugous CPU memory region.

julia/src/base.jl

@@ -26,6 +26,7 @@ Base.show(io::IO, e::MXError) = print(io, e.msg)
 # Common types used in MXNet API
 ################################################################################
 const MX_uint = Cuint
+const MX_long = Clonglong
 const MX_float = Cfloat
 const MX_handle = Ptr{Void}
 

julia/src/ndarray.jl

@@ -245,8 +245,8 @@ See also the notes on NDArray shapes [`NDArray`](@ref).
 """
 function size(x::NDArray)
   ref_ndim  = Ref{MX_uint}(0)
-  ref_shape = Ref{Ptr{MX_uint}}(0)
-  @mxcall(:MXNDArrayGetShape, (MX_handle, Ref{MX_uint}, Ref{Ptr{MX_uint}}),
+  ref_shape = Ref{Ptr{MX_long}}(0)
+  @mxcall(:MXNDArrayGetShape, (MX_handle, Ref{MX_uint}, Ref{Ptr{MX_long}}),
           x, ref_ndim, ref_shape)
   tuple(map(Int, flipdim(unsafe_wrap(Array, ref_shape[], ref_ndim[]),1))...)
 end
@@ -278,8 +278,8 @@ ndims(x::NDArray) = ndims(x.handle)
 
 function ndims(x::MX_NDArrayHandle)::Int
   ref_ndim  = Ref{MX_uint}(0)
-  ref_shape = Ref{Ptr{MX_uint}}(0)
-  @mxcall(:MXNDArrayGetShape, (MX_handle, Ref{MX_uint}, Ref{Ptr{MX_uint}}),
+  ref_shape = Ref{Ptr{MX_long}}(0)
+  @mxcall(:MXNDArrayGetShape, (MX_handle, Ref{MX_uint}, Ref{Ptr{MX_long}}),
           x, ref_ndim, ref_shape)
   ref_ndim[]
 end