Merge pull request #2224 from Ericgig/bug.memory_overflow

Add insufficient memory checks

doc/changes/2224.bugfix

@@ -0,0 +1 @@
+Raise on insufficient memory.

qutip/core/data/add.pyx

@@ -18,12 +18,6 @@ from qutip.core.data cimport csr, dense, dia
 
 cnp.import_array()
 
-cdef extern from *:
-    void *PyDataMem_NEW(size_t size)
-    void *PyDataMem_NEW_ZEROED(size_t size, size_t elsize)
-    void PyDataMem_FREE(void *ptr)
-
-
 __all__ = [
     'add', 'add_csr', 'add_dense', 'iadd_dense', 'add_dia',
     'sub', 'sub_csr', 'sub_dense', 'sub_dia',

qutip/core/data/csr.pyx

@@ -530,6 +530,9 @@ cpdef CSR empty(base.idxint rows, base.idxint cols, base.idxint size):
         <base.idxint *> PyDataMem_NEW(size * sizeof(base.idxint))
     out.row_index =\
         <base.idxint *> PyDataMem_NEW(row_size * sizeof(base.idxint))
+    if not out.data: raise MemoryError()
+    if not out.col_index: raise MemoryError()
+    if not out.row_index: raise MemoryError()
     # Set the number of non-zero elements to 0.
     out.row_index[rows] = 0
     return out

qutip/core/data/dense.pyx

@@ -120,6 +120,7 @@ cdef class Dense(base.Data):
         cdef Dense out = Dense.__new__(Dense)
         cdef size_t size = self.shape[0]*self.shape[1]*sizeof(double complex)
         cdef double complex *ptr = <double complex *> PyDataMem_NEW(size)
+        if not ptr: raise MemoryError()
         memcpy(ptr, self.data, size)
         out.shape = self.shape
         out.data = ptr
@@ -162,6 +163,7 @@ cdef class Dense(base.Data):
         """
         cdef size_t size = self.shape[0]*self.shape[1]*sizeof(double complex)
         cdef double complex *ptr = <double complex *> PyDataMem_NEW(size)
+        if not ptr: raise MemoryError()
         memcpy(ptr, self.data, size)
         cdef object out =\
             cnp.PyArray_SimpleNewFromData(2, [self.shape[0], self.shape[1]],
@@ -244,6 +246,7 @@ cpdef Dense empty(base.idxint rows, base.idxint cols, bint fortran=True):
     cdef Dense out = Dense.__new__(Dense)
     out.shape = (rows, cols)
     out.data = <double complex *> PyDataMem_NEW(rows * cols * sizeof(double complex))
+    if not out.data: raise MemoryError()
     out._deallocate = True
     out.fortran = fortran
     return out
@@ -264,6 +267,7 @@ cpdef Dense zeros(base.idxint rows, base.idxint cols, bint fortran=True):
     out.shape = (rows, cols)
     out.data =\
         <double complex *> PyDataMem_NEW_ZEROED(rows * cols, sizeof(double complex))
+    if not out.data: raise MemoryError()
     out.fortran = fortran
     out._deallocate = True
     return out
@@ -290,6 +294,7 @@ cpdef Dense from_csr(CSR matrix, bint fortran=False):
         <double complex *>
         PyDataMem_NEW_ZEROED(out.shape[0]*out.shape[1], sizeof(double complex))
     )
+    if not out.data: raise MemoryError()
     out.fortran = fortran
     out._deallocate = True
     cdef size_t row, ptr_in, ptr_out, row_stride, col_stride

qutip/core/data/dia.pyx

@@ -271,6 +271,8 @@ cpdef Dia empty(base.idxint rows, base.idxint cols, base.idxint num_diag):
         <double complex *> PyDataMem_NEW(cols * num_diag * sizeof(double complex))
     out.offsets =\
         <base.idxint *> PyDataMem_NEW(num_diag * sizeof(base.idxint))
+    if not out.data: raise MemoryError()
+    if not out.offsets: raise MemoryError()
     return out
 
 

qutip/core/data/kron.pyx

@@ -10,6 +10,7 @@ from qutip.core.data.dense cimport Dense
 from .adjoint import transpose
 from qutip.core.data.dia cimport Dia
 from qutip.core.data cimport csr, dia
+from qutip.core.data.convert import to as _to
 import numpy
 
 __all__ = [
@@ -112,36 +113,48 @@ cpdef Dia kron_dia(Dia left, Dia right):
                     right.shape[1]
                 )
                 num_diag += 1
+        out.num_diag = num_diag
 
     else:
-        out = dia.empty(nrows, ncols, _mul_checked(max_diag, ncols_l))
-        delta = right.shape[0] - right.shape[1]
-        for diag_left in range(left.num_diag):
-            for diag_right in range(right.num_diag):
-                start_left = max(0, left.offsets[diag_left])
-                end_left = min(left.shape[1], left.shape[0] + left.offsets[diag_left])
-                for col_left in range(start_left, end_left):
-                    memset(
-                        out.data + (num_diag * out.shape[1]), 0,
-                        out.shape[1] * sizeof(double complex)
-                    )
-
-                    out.offsets[num_diag] = (
-                        left.offsets[diag_left] * right.shape[0]
-                        + right.offsets[diag_right]
-                        - col_left * delta
-                    )
-
-                    start_right = max(0, right.offsets[diag_right])
-                    end_right = min(right.shape[1], right.shape[0] + right.offsets[diag_right])
-                    for col_right in range(start_right, end_right):
-                        out.data[num_diag * out.shape[1] + col_left * right.shape[1] + col_right] = (
-                            right.data[diag_right * right.shape[1] + col_right]
-                            * left.data[diag_left * left.shape[1] + col_left]
+        max_diag = _mul_checked(max_diag, ncols_l)
+        if max_diag < nrows:
+            out = dia.empty(nrows, ncols, max_diag)
+            delta = right.shape[0] - right.shape[1]
+            for diag_left in range(left.num_diag):
+                for diag_right in range(right.num_diag):
+                    start_left = max(0, left.offsets[diag_left])
+                    end_left = min(left.shape[1], left.shape[0] + left.offsets[diag_left])
+                    for col_left in range(start_left, end_left):
+                        memset(
+                            out.data + (num_diag * out.shape[1]), 0,
+                            out.shape[1] * sizeof(double complex)
                         )
-                    num_diag += 1
 
-    out.num_diag = num_diag
+                        out.offsets[num_diag] = (
+                            left.offsets[diag_left] * right.shape[0]
+                            + right.offsets[diag_right]
+                            - col_left * delta
+                        )
+
+                        start_right = max(0, right.offsets[diag_right])
+                        end_right = min(right.shape[1], right.shape[0] + right.offsets[diag_right])
+                        for col_right in range(start_right, end_right):
+                            out.data[num_diag * out.shape[1] + col_left * right.shape[1] + col_right] = (
+                                right.data[diag_right * right.shape[1] + col_right]
+                                * left.data[diag_left * left.shape[1] + col_left]
+                            )
+                        num_diag += 1
+            out.num_diag = num_diag
+
+        else:
+            # The output is not sparse enough ant the empty data array would be
+            # larger than the dense array.
+            # Fallback on dense operation
+            left_dense = _to(Dense, left)
+            right_dense = _to(Dense, right)
+            out_dense = kron_dense(left_dense, right_dense)
+            out = _to(Dia, out_dense)
+
     out = dia.clean_dia(out, True)
     return out