Merge pull request #1817 from nouiz/crash

theano/sandbox/cuda/cuda_ndarray.cu

@@ -915,12 +915,14 @@ __global__ void k_take_3(const int d0, const int d1, const int d2,
         npy_int64 idx = indices[i0];
         if (idx<0)
             idx += dB0; // To allow negative indexing.
-        if ((idx < 0) || (idx >= dB0))
+        if ((idx < 0) || (idx >= dB0)){
             // Any value other the 0 probably work. But to be more safe, I want
             // to change all bits to prevent problem with concurrent write that
             // could cross cache line. But this should not happen with the
             // current code and driver.
             *err = 0xFFFF;
+            continue;
+        }
         for (int i1 = threadIdx.x; i1 < d1; i1 += blockDim.x){
             for (int i2 = threadIdx.y; i2 < d2; i2 += blockDim.y){
                 int a_idx = i0*sA0 + i1*sA1 + i2*sA2;

theano/sandbox/cuda/cuda_ndarray.cuh

@@ -34,6 +34,11 @@
 
 #include <numpy/arrayobject.h>
 #include <stdio.h>
+#include <stdint.h>
+#ifndef SIZE_MAX
+    #define SIZE_MAX ((size_t)-1)
+#endif
+
 
 #include <cublas.h>
 
@@ -342,7 +347,7 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
 {
     // allocate an empty ndarray with c_contiguous access
     // return 0 on success
-    int size = 1; //set up the strides for contiguous tensor
+    size_t size = 1; //set up the strides for contiguous tensor
     assert (nd >= 0);
 
     // Here we modify the host structure to have the desired shape and
@@ -357,6 +362,13 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
         {
             CudaNdarray_set_stride(self, i, (dim[i] == 1) ? 0 : size);
             CudaNdarray_set_dim(self, i, dim[i]);
+            //Detect overflow on unsigned integer
+            if (size > (SIZE_MAX / dim[i])) {
+                PyErr_Format(PyExc_AssertionError,
+                             "Can't store in size_t the bytes resquested",
+                             size);
+                return -1;
+            }
             size = size * dim[i];
         }
     }
@@ -366,6 +378,14 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
         {
             CudaNdarray_set_stride(self, i, (dim[i] == 1) ? 0 : size);
             CudaNdarray_set_dim(self, i, dim[i]);
+
+            //Detect overflow on unsigned integer
+            if (size > (SIZE_MAX / dim[i])) {
+                PyErr_Format(PyExc_AssertionError,
+                             "Can't store in size_t the bytes resquested",
+                             size);
+                return -1;
+            }
             size = size * dim[i];
         }
     }
@@ -393,14 +413,6 @@ static int CudaNdarray_alloc_contiguous(CudaNdarray *self, const int nd,
         return -1;
     }
 
-    if (size < 0)
-    {
-        PyErr_Format(PyExc_AssertionError,
-                     "size (%i) < 0",
-                     size);
-        return -1;
-    }
-
     self->devdata = (float*)device_malloc(size*sizeof(real));
     if (size && !self->devdata)
     {