Repository: numba/numba

Issue:
Tricky numpy.dot memory issue
Hi all, I encounter a memory issue when running a piece of my code with numpy.dot in a jit decorated function. The memory is unusually consumed even though I run the np.dot on a contiguous array. Here is the code to repeat this problem.

```
import numpy as np
import numba

@numba.jit('complex128(complex128[:,:,:], int32, int32, int32, int32)', nopython=True, nogil=True)
def _cubeFactorToElement_bug(cubeFactor, cubei0,cubej0, cubei1,cubej1):
    return np.dot(cubeFactor[cubei0,cubej0,:], np.conj(cubeFactor[cubei1,cubej1,:]))

@numba.jit('void(complex128[:,:,:], complex128[:], int32[:],int32[:])', nopython=True, nogil=True)
def f(tensor, ret, row, col):
    for idx0 in range(len(row)):
        for idx1 in range(len(row)):
            ret[0] += _cubeFactorToElement_bug(tensor, row[idx0], col[idx0] ,row[idx1] ,col[idx1])

np.random.seed((1,2,3))
rows = np.array([np.random.randint(95) for _ in range(2000)], np.int32)
cols = np.array([np.random.randint(95) for _ in range(2000)], np.int32)

A = np.zeros((95,95,50),np.complex128)
res = np.zeros(1, np.complex128)
```

The above is the setup and the following will induce the memory consuming problem.

```
f(A,res,rows, cols)
```

Each time I run this code, roughly 3GB of the ram is consumed. I believe this is a bug due to numba implementation of np.dot. If I replace `_cubeFactorToElement_bug` by the following code:

```
@numba.jit('complex128(complex128[:,:,:], int32, int32, int32, int32)', nopython=True, nogil=True)
def _cubeFactorToElement_ok(cubeFactor, cubei0,cubej0, cubei1,cubej1):
    res = cubeFactor[cubei0,cubej0,0] * np.conj(cubeFactor[cubei1,cubej1,0])
    for k in range(1,cubeFactor.shape[0]):
        res += cubeFactor[cubei0,cubej0,k] * np.conj(cubeFactor[cubei1,cubej1,k])
    return res
```

the memory issue is resolved.



Discussion:
Thanks for reporting this! This is an actual bug indeed.


Resolving patch:
diff --git a/numba/targets/linalg.py b/numba/targets/linalg.py
--- a/numba/targets/linalg.py
+++ b/numba/targets/linalg.py
@@ -4,6 +4,8 @@
 
 from __future__ import print_function, absolute_import, division
 
+import contextlib
+
 from llvmlite import ir
 import numpy
 
@@ -45,14 +47,17 @@ def make_constant_slot(context, builder, ty, val):
     const = context.get_constant_generic(builder, ty, val)
     return cgutils.alloca_once_value(builder, const)
 
+
+@contextlib.contextmanager
 def make_contiguous(context, builder, sig, args):
     """
     Ensure that all array arguments are contiguous, if necessary by
     copying them.
-    A new (sig, args) tuple is returned.
+    A new (sig, args) tuple is yielded.
     """
     newtys = []
     newargs = []
+    copies = []
     for ty, val in zip(sig.args, args):
         if not isinstance(ty, types.Array) or ty.layout in 'CF':
             newty, newval = ty, val
@@ -60,9 +65,12 @@ def make_contiguous(context, builder, sig, args):
             newty = ty.copy(layout='C')
             copysig = signature(newty, ty)
             newval = array_copy(context, builder, copysig, (val,))
+            copies.append((newty, newval))
         newtys.append(newty)
         newargs.append(newval)
-    return signature(sig.return_type, *newtys), tuple(newargs)
+    yield signature(sig.return_type, *newtys), tuple(newargs)
+    for ty, val in copies:
+        context.nrt_decref(builder, ty, val)
 
 
 def check_c_int(context, builder, n):
@@ -286,19 +294,19 @@ def dot_2(context, builder, sig, args):
     a @ b
     """
     ensure_blas()
-    sig, args = make_contiguous(context, builder, sig, args)
-
-    ndims = [x.ndim for x in sig.args[:2]]
-    if ndims == [2, 2]:
-        return dot_2_mm(context, builder, sig, args)
-    elif ndims == [2, 1]:
-        return dot_2_mv(context, builder, sig, args)
-    elif ndims == [1, 2]:
-        return dot_2_vm(context, builder, sig, args)
-    elif ndims == [1, 1]:
-        return dot_2_vv(context, builder, sig, args)
-    else:
-        assert 0
+
+    with make_contiguous(context, builder, sig, args) as (sig, args):
+        ndims = [x.ndim for x in sig.args[:2]]
+        if ndims == [2, 2]:
+            return dot_2_mm(context, builder, sig, args)
+        elif ndims == [2, 1]:
+            return dot_2_mv(context, builder, sig, args)
+        elif ndims == [1, 2]:
+            return dot_2_vm(context, builder, sig, args)
+        elif ndims == [1, 1]:
+            return dot_2_vv(context, builder, sig, args)
+        else:
+            assert 0
 
 @lower_builtin(numpy.vdot, types.Array, types.Array)
 def vdot(context, builder, sig, args):
@@ -306,9 +314,9 @@ def vdot(context, builder, sig, args):
     np.vdot(a, b)
     """
     ensure_blas()
-    sig, args = make_contiguous(context, builder, sig, args)
 
-    return dot_2_vv(context, builder, sig, args, conjugate=True)
+    with make_contiguous(context, builder, sig, args) as (sig, args):
+        return dot_2_vv(context, builder, sig, args, conjugate=True)
 
 
 def dot_3_vm(context, builder, sig, args):
@@ -455,15 +463,15 @@ def dot_3(context, builder, sig, args):
     np.dot(a, b, out)
     """
     ensure_blas()
-    sig, args = make_contiguous(context, builder, sig, args)
 
-    ndims = set(x.ndim for x in sig.args[:2])
-    if ndims == set([2]):
-        return dot_3_mm(context, builder, sig, args)
-    elif ndims == set([1, 2]):
-        return dot_3_vm(context, builder, sig, args)
-    else:
-        assert 0
+    with make_contiguous(context, builder, sig, args) as (sig, args):
+        ndims = set(x.ndim for x in sig.args[:2])
+        if ndims == set([2]):
+            return dot_3_mm(context, builder, sig, args)
+        elif ndims == set([1, 2]):
+            return dot_3_vm(context, builder, sig, args)
+        else:
+            assert 0
 
 
 def call_xxgetrf(context, builder, a_type, a_shapes, a_data, ipiv, info):
@@ -545,7 +553,7 @@ def mat_inv(context, builder, sig, args):
     # Scipy for which one can specify to whether or not to overwrite the
     # input).
     def create_out(a):
-        m, n= a.shape
+        m, n = a.shape
         if m != n:
             raise ValueError("np.linalg.inv can only work on square "
                              "arrays.")
@@ -572,15 +580,14 @@ def create_out(a):
     invalid_arg = builder.icmp_signed('<', info_val, zero)
 
     with builder.if_then(lapack_error, False):
+        context.nrt_decref(builder, ipiv_t, ipiv)
         with builder.if_else(invalid_arg) as (then, otherwise):
             raise_err = context.call_conv.return_user_exc
             with then:
-                context.nrt_decref(builder, ipiv_t, ipiv)
                 raise_err(builder, ValueError,
                           ('One argument passed to getrf is invalid',)
                           )
             with otherwise:
-                context.nrt_decref(builder, ipiv_t, ipiv)
                 raise_err(builder, ValueError,
                           ('Matrix is singular and cannot be inverted',)
                           )


Test patch:
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -102,7 +102,7 @@ def assertRefCount(self, *objects):
         """
         A context manager that asserts the given objects have the
         same reference counts before and after executing the
-        enclosed blocks.
+        enclosed block.
         """
         old_refcounts = [sys.getrefcount(x) for x in objects]
         yield
@@ -112,6 +112,25 @@ def assertRefCount(self, *objects):
                 self.fail("Refcount changed from %d to %d for object: %r"
                           % (old, new, obj))
 
+    @contextlib.contextmanager
+    def assertNoNRTLeak(self):
+        """
+        A context manager that asserts no NRT leak was created during
+        the execution of the enclosed block.
+        """
+        old = rtsys.get_allocation_stats()
+        yield
+        new = rtsys.get_allocation_stats()
+        total_alloc = new.alloc - old.alloc
+        total_free = new.free - old.free
+        total_mi_alloc = new.mi_alloc - old.mi_alloc
+        total_mi_free = new.mi_free - old.mi_free
+        self.assertEqual(total_alloc, total_free,
+                         "number of data allocs != number of data frees")
+        self.assertEqual(total_mi_alloc, total_mi_free,
+                         "number of meminfo allocs != number of meminfo frees")
+
+
     _exact_typesets = [(bool, np.bool_), utils.INT_TYPES, (str,), (np.integer,), (utils.text_type), ]
     _approx_typesets = [(float,), (complex,), (np.inexact)]
     _sequence_typesets = [(tuple, list)]
diff --git a/numba/tests/test_linalg.py b/numba/tests/test_linalg.py
--- a/numba/tests/test_linalg.py
+++ b/numba/tests/test_linalg.py
@@ -1,6 +1,7 @@
 from __future__ import division, print_function
 
 import contextlib
+import gc
 from itertools import product
 import sys
 import warnings
@@ -38,6 +39,10 @@ class TestProduct(TestCase):
 
     dtypes = (np.float64, np.float32, np.complex128, np.complex64)
 
+    def setUp(self):
+        # Collect leftovers from previous test cases before checking for leaks
+        gc.collect()
+
     def sample_vector(self, n, dtype):
         # Be careful to generate only exactly representable float values,
         # to avoid rounding discrepancies between Numpy and Numba
@@ -66,16 +71,20 @@ def check_contiguity_warning(self, pyfunc):
         self.assertEqual(w[0].lineno, pyfunc.__code__.co_firstlineno + 1)
 
     def check_func(self, pyfunc, cfunc, args):
-        expected = pyfunc(*args)
-        got = cfunc(*args)
-        self.assertPreciseEqual(got, expected, ignore_sign_on_zero=True)
+        with self.assertNoNRTLeak():
+            expected = pyfunc(*args)
+            got = cfunc(*args)
+            self.assertPreciseEqual(got, expected, ignore_sign_on_zero=True)
+            del got, expected
 
     def check_func_out(self, pyfunc, cfunc, args, out):
-        expected = np.copy(out)
-        got = np.copy(out)
-        self.assertIs(pyfunc(*args, out=expected), expected)
-        self.assertIs(cfunc(*args, out=got), got)
-        self.assertPreciseEqual(got, expected, ignore_sign_on_zero=True)
+        with self.assertNoNRTLeak():
+            expected = np.copy(out)
+            got = np.copy(out)
+            self.assertIs(pyfunc(*args, out=expected), expected)
+            self.assertIs(cfunc(*args, out=got), got)
+            self.assertPreciseEqual(got, expected, ignore_sign_on_zero=True)
+            del got, expected
 
     def assert_mismatching_sizes(self, cfunc, args, is_out=False):
         with self.assertRaises(ValueError) as raises:
@@ -301,6 +310,10 @@ class TestLinalgInv(TestCase):
 
     dtypes = (np.float64, np.float32, np.complex128, np.complex64)
 
+    def setUp(self):
+        # Collect leftovers from previous test cases before checking for leaks
+        gc.collect()
+
     def sample_vector(self, n, dtype):
         # Be careful to generate only exactly representable float values,
         # to avoid rounding discrepancies between Numpy and Numba
@@ -344,21 +357,23 @@ def test_linalg_inv(self):
         """
         n = 10
         cfunc = jit(nopython=True)(invert_matrix)
+
+        def check(a, **kwargs):
+            with self.assertNoNRTLeak():
+                expected = invert_matrix(a).copy(order='C')
+                got = cfunc(a)
+                # XXX had to use that function otherwise comparison fails
+                # because of +0, -0 discrepancies
+                np.testing.assert_array_almost_equal_nulp(got, expected, **kwargs)
+                del got, expected
+
         for dtype, order in product(self.dtypes, 'CF'):
             a = self.sample_matrix(n, dtype, order)
-            expected = invert_matrix(a).copy(order='C')
-            got = cfunc(a)
-            # XXX add to use that function otherwise comparison fails
-            # because of +0, -0 discrepancies
-            np.testing.assert_array_almost_equal_nulp(got, expected, nulp=3)
+            check(a, nulp=3)
 
         for order in 'CF':
             a = np.array(((2, 1), (2, 3)), dtype=np.float64, order=order)
-            expected = invert_matrix(a).copy(order='C')
-            got = cfunc(a)
-            # XXX add to use that function otherwise comparison fails
-            # because of +0, -0 discrepancies
-            np.testing.assert_array_almost_equal_nulp(got, expected)
+            check(a)
 
         # Non square matrices
         self.assert_non_square(cfunc, (np.ones((2,3)),))