from __future__ import print_function, absolute_import, division
import numpy as np
from numba.cuda.testing import unittest
from numba import cuda


@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 use_add2f(ary):
    i = cuda.grid(1)
    ary[i] = add2f(ary[i], ary[i])


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


class TestDeviceFunc(unittest.TestCase):
    def test_use_add2f(self):
        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):
        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))


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