from __future__ import print_function, absolute_import, division

import numpy as np

from numba.cuda.testing import unittest, skip_on_cudasim
from numba import cuda, jit
from numba.errors import TypingError


class TestDeviceFunc(unittest.TestCase):

    def test_use_add2f(self):

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

        def use_add2f(ary):
            i = cuda.grid(1)
            ary[i] = add2f(ary[i], ary[i])

        compiled = cuda.jit("void(float32[:])")(use_add2f)

        nelem = 10
        ary = np.arange(nelem, dtype=np.float32)
        exp = ary + ary
        compiled[1, nelem](ary)

        self.assertTrue(np.all(ary == exp), (ary, exp))

    def test_indirect_add2f(self):

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

        @cuda.jit("float32(float32, float32)", device=True)
        def indirect(a, b):
            return add2f(a, b)

        def indirect_add2f(ary):
            i = cuda.grid(1)
            ary[i] = indirect(ary[i], ary[i])

        compiled = cuda.jit("void(float32[:])")(indirect_add2f)

        nelem = 10
        ary = np.arange(nelem, dtype=np.float32)
        exp = ary + ary
        compiled[1, nelem](ary)

        self.assertTrue(np.all(ary == exp), (ary, exp))

    def _check_cpu_dispatcher(self, add):
        @cuda.jit
        def add_kernel(ary):
            i = cuda.grid(1)
            ary[i] = add(ary[i], 1)

        ary = np.arange(10)
        expect = ary + 1
        add_kernel[1, ary.size](ary)
        np.testing.assert_equal(expect, ary)

    def test_cpu_dispatcher(self):
        # Test correct usage
        @jit
        def add(a, b):
            return a + b

        self._check_cpu_dispatcher(add)

    @skip_on_cudasim('not supported in cudasim')
    def test_cpu_dispatcher_invalid(self):
        # Test invalid usage
        # Explicit signature disables compilation, which also disable
        # compiling on CUDA.
        @jit('(i4, i4)')
        def add(a, b):
            return a + b

        # Check that the right error message is provided.
        with self.assertRaises(TypingError) as raises:
            self._check_cpu_dispatcher(add)
        msg = "Untyped global name 'add': using cpu function on device"
        self.assertIn(msg, str(raises.exception))


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

@@ -59,16 +59,31 @@
 vision%0A%0A
+%0Aimport types%0A%0A
 import n
@@ -2400,16 +2400,469 @@
 tion))%0A%0A
+    def test_cpu_dispatcher_other_module(self):%0A        @jit%0A        def add(a, b):%0A            return a + b%0A%0A        mymod = types.ModuleType(name='mymod')%0A        mymod.add = add%0A        del add%0A%0A        @cuda.jit%0A        def add_kernel(ary):%0A            i = cuda.grid(1)%0A            ary%5Bi%5D = mymod.add(ary%5Bi%5D, 1)%0A%0A        ary = np.arange(10)%0A        expect = ary + 1%0A        add_kernel%5B1, ary.size%5D(ary)%0A        np.testing.assert_equal(expect, ary)%0A%0A
 %0Aif __na

Add test that uses a cpu dispatcher object as an attribute of a module