Better/more overflow check.

theano/sandbox/cuda/cuda_ndarray.cu

@@ -568,15 +568,27 @@ PyObject* CudaNdarray_ZEROS(int n, int * dims)
 {
 
     size_t total_elements = 1;
-    for(size_t i=0;i<n;i++)
+
+    for(size_t i=0;i<n;i++){
+        // Detect overflow on unsigned integer
+        if (dims[i] != 0 && total_elements > (SIZE_MAX / dims[i])) {
+            PyErr_Format(PyExc_RuntimeError,
+                         "Can't store in size_t for the bytes requested %llu * %llu",
+                         (unsigned long long)total_elements,
+                         (unsigned long long)dims[i]);
+            return NULL;
+        }
         total_elements*=dims[i];
+    }
 
     // total_elements now contains the size of the array, in reals
-    size_t total_size = total_elements * sizeof(real);
-    if (total_elements < 0 or total_size < 0 or total_size < total_elements){
-        PyErr_SetString(PyExc_RuntimeError, "CudaNdarray_ZEROS: size_t overflow, can't allocate that much memory");
+    if (total_elements > (SIZE_MAX / sizeof(real))){
+        PyErr_Format(PyExc_RuntimeError,
+                     "Can't store in size_t for the bytes requested %llu * 4",
+                     (unsigned long long)total_elements);
         return NULL;
     }
+    size_t total_size = total_elements * sizeof(real);
 
     CudaNdarray* rval = (CudaNdarray*)CudaNdarray_New();
     if (!rval)

theano/sandbox/cuda/cuda_ndarray.cuh

@@ -378,8 +378,8 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
             //Detect overflow on unsigned integer
             if (dim[i] != 0 && size > (SIZE_MAX / dim[i])) {
                 PyErr_Format(PyExc_AssertionError,
-                             "Can't store in size_t for the bytes requested %llu",
-                             (unsigned long long)size);
+                             "Can't store in size_t for the bytes requested %llu * %llu",
+                             (unsigned long long)size, (unsigned long long)dim[i]);
                 return -1;
             }
             size = size * dim[i];
@@ -395,7 +395,7 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
             //Detect overflow on unsigned integer
             if (dim[i] != 0 && size > (SIZE_MAX / dim[i])) {
                 PyErr_Format(PyExc_AssertionError,
-                             "Can't store in size_t for the bytes requested %llu",
+                             "Can't store in size_t for the bytes requested %llu * 4",
                              (unsigned long long)size);
                 return -1;
             }
@@ -403,6 +403,14 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
         }
     }
 
+    // Detect overflow on unsigned integer
+    if (size > (SIZE_MAX / sizeof(real))) {
+        PyErr_Format(PyExc_RuntimeError,
+                     "Can't store in size_t for the bytes requested %llu",
+                     (unsigned long long)size);
+        return -1;
+    }
+
     // If the allocated buffer is already of the right size, we don't need to
     // do anything else.
     // Note: self->data_allocated is 0 for a view, so views will fail this