from __future__ import print_function
from numba import unittest_support as unittest
from numba import cuda
import numpy as np


class TestCudaAutojit(unittest.TestCase):
    def test_device_array(self):
        @cuda.autojit
        def foo(x, y):
            i = cuda.grid(1)
            y[i] = x[i]

        x = np.arange(10)
        y = np.empty_like(x)

        dx = cuda.to_device(x)
        dy = cuda.to_device(y)

        foo[10, 1](dx, dy)

        dy.copy_to_host(y)

        self.assertTrue(np.all(x == y))

    def test_device_auto_jit(self):
        @cuda.jit(device=True)
        def mapper(args):
            a, b, c = args
            return a + b + c


        @cuda.jit(device=True)
        def reducer(a, b):
            return a + b


        @cuda.jit
        def driver(A, B):
            i = cuda.grid(1)
            if i < B.size:
                args = A[i], A[i] + B[i], B[i]
                B[i] = reducer(mapper(args), 1)

        A = np.arange(100, dtype=np.float32)
        B = np.arange(100, dtype=np.float32)

        Acopy = A.copy()
        Bcopy = B.copy()

        driver[1, 100](A, B)

        np.testing.assert_allclose(Acopy + Acopy + Bcopy + Bcopy + 1, B)


if __name__ == '__main__':
    unittest.main()

@@ -1190,16 +1190,532 @@
  1, B)%0A%0A
+    def test_device_auto_jit_2(self):%0A        @cuda.jit(device=True)%0A        def inner(arg):%0A            return arg + 1%0A%0A        @cuda.jit%0A        def outer(argin, argout):%0A            argout%5B0%5D = inner(argin%5B0%5D) + inner(2)%0A%0A        a = np.zeros(1)%0A        b = np.zeros(1)%0A%0A        stream = cuda.stream()%0A        d_a = cuda.to_device(a, stream)%0A        d_b = cuda.to_device(b, stream)%0A%0A        outer%5B1, 1, stream%5D(d_a, d_b)%0A%0A        d_b.copy_to_host(b, stream)%0A%0A        self.assertEqual(b%5B0%5D, (a%5B0%5D + 1) + (2 + 1))%0A%0A
 %0Aif __na

Add test