bug fix for 1>>31 -> uint32_t{1}<<31 and showing error message when u…

…ser attempts to pass large ndarrays as inputs to operators or large shape of inputs for memory allocation.

python/mxnet/ndarray/ndarray.py

@@ -155,6 +155,11 @@ def _new_alloc_handle(shape, ctx, delay_alloc, dtype=mx_real_t):
             ctypes.c_int(int(_DTYPE_NP_TO_MX[np.dtype(dtype).type])),
             ctypes.byref(hdl)))
     else:
+        size = 1
+        for idx in shape:
+            size = size * idx
+        if size > 2**32:
+            print("Size of tensor you are trying to allocate is larger than 2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1")
         check_call(_LIB.MXNDArrayCreateEx(
             c_array_buf(mx_uint, native_array('I', shape)),
             mx_uint(len(shape)),

src/c_api/c_api.cc

@@ -198,7 +198,7 @@ void CreateNDArray(const DataType* shape,
                    int dtype,
                    NDArrayHandle* out) {
   mxnet::TShape requested_shape = mxnet::TShape(shape, shape + ndim);
-  CHECK_LT(requested_shape.Size(), 1 >> 31)
+  CHECK_LT(requested_shape.Size(), (uint32_t{1} << 31) - 1)
            << "Size of tensor you are trying to allocate is larger than "
               "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
   *out = new NDArray(requested_shape,

src/c_api/c_api_ndarray.cc

@@ -54,7 +54,10 @@ void SetNDInputsOutputs(const nnvm::Op* op,
   ndinputs->clear();
   ndinputs->reserve(num_inputs);
   for (int i = 0; i < num_inputs; ++i) {
-    ndinputs->emplace_back(reinterpret_cast<NDArray*>(inputs[i]));
+    NDArray* inp = reinterpret_cast<NDArray*>(inputs[i]);
+    CHECK_LT(inp->shape().Size(), (uint32_t{1} << 31) - 1) << "Size of tensor you are trying to allocate is larger than "
+                               "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
+    ndinputs->emplace_back(inp);
   }
 
   ndoutputs->clear();

src/ndarray/ndarray.cc

@@ -1879,7 +1879,7 @@ NDArray NDArray::Copy(Context ctx) const {
 
 void NDArray::SyncCopyFromCPU(const void *data, size_t size) const {
   mxnet::TShape dshape = this->shape();
-  CHECK_LT(size, 1 >> 31) << "Size of tensor you are trying to allocate is larger than "
+  CHECK_LT(size, (uint32_t{1} << 31) - 1) << "Size of tensor you are trying to allocate is larger than "
                              "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
   CHECK_EQ(dshape.Size(), size)
       << "Memory size do not match";
@@ -2016,7 +2016,7 @@ void NDArray::SyncCopyFromNDArray(const NDArray& src, int i, int j) {
 
 void NDArray::SyncCopyToCPU(void *data, size_t size) const {
   mxnet::TShape dshape = this->shape();
-  CHECK_LT(size, 1 >> 31) << "Size of tensor you are trying to allocate is larger than "
+  CHECK_LT(size, (uint32_t{1} << 31) - 1) << "Size of tensor you are trying to allocate is larger than "
                              "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
   CHECK_EQ(dshape.Size(), size)
       << "Memory size do not match";

src/operator/tensor/init_op.h

@@ -336,6 +336,8 @@ inline bool InitStorageType(const nnvm::NodeAttrs& attrs,
 template <bool is_integer = false, typename ValueType, typename xpu>
 void Fill(mshadow::Stream<xpu> *s, const TBlob& b, const OpReqType req, ValueType val) {
   // If b is a zero-size tensor, do nothing.
+  CHECK_LT(b.Size(), (uint32_t{1} << 31) - 1) << "Size of tensor you are trying to allocate is larger than "
+                                 "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
   if (b.Size() == 0) return;
   if (req != kNullOp) {
     const size_t size = b.Size();