import numpy as np

from numba.roc.vectorizers import HsaGUFuncVectorize
from numba.roc.dispatch import HSAGenerializedUFunc
from numba import guvectorize
import unittest


def ufunc_add_core(a, b, c):
    for i in range(c.size):
        c[i] = a[i] + b[i]


class TestGUFuncBuilding(unittest.TestCase):
    def test_gufunc_building(self):
        ufbldr = HsaGUFuncVectorize(ufunc_add_core, "(x),(x)->(x)")
        ufbldr.add("(float32[:], float32[:], float32[:])")
        ufbldr.add("(intp[:], intp[:], intp[:])")
        ufunc = ufbldr.build_ufunc()
        self.assertIsInstance(ufunc, HSAGenerializedUFunc)

        # Test integer version
        A = np.arange(100, dtype=np.intp)
        B = np.arange(100, dtype=np.intp) + 1
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_equal(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.intp), got.dtype)

        # Test integer version with 2D inputs
        A = A.reshape(50, 2)
        B = B.reshape(50, 2)
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_equal(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.intp), got.dtype)

        # Test integer version with 3D inputs
        A = A.reshape(5, 10, 2)
        B = B.reshape(5, 10, 2)
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_equal(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.intp), got.dtype)

        # Test real version
        A = np.arange(100, dtype=np.float32)
        B = np.arange(100, dtype=np.float32) + 1
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_allclose(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.float32), got.dtype)

        # Test real version with 2D inputs
        A = A.reshape(50, 2)
        B = B.reshape(50, 2)
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_allclose(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.float32), got.dtype)

    def test_gufunc_building_scalar_output(self):
        def sum_row(inp, out):
            tmp = 0.
            for i in range(inp.shape[0]):
                tmp += inp[i]
            out[0] = tmp

        ufbldr = HsaGUFuncVectorize(sum_row, "(n)->()")
        ufbldr.add("void(int32[:], int32[:])")
        ufunc = ufbldr.build_ufunc()

        inp = np.arange(300, dtype=np.int32).reshape(100, 3)
        out = ufunc(inp)

        for i in range(inp.shape[0]):
            np.testing.assert_equal(inp[i].sum(), out[i])

    def test_gufunc_scalar_input_saxpy(self):
        def axpy(a, x, y, out):
            for i in range(out.shape[0]):
                out[i] = a * x[i] + y[i]

        ufbldr = HsaGUFuncVectorize(axpy, '(),(t),(t)->(t)')
        ufbldr.add("void(float32, float32[:], float32[:], float32[:])")
        saxpy = ufbldr.build_ufunc()

        A = np.float32(2)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i])

        X = np.arange(10, dtype=np.float32)
        Y = np.arange(10, dtype=np.float32)
        out = saxpy(A, X, Y)

        for j in range(10):
            exp = A * X[j] + Y[j]
            self.assertTrue(exp == out[j], (exp, out[j]))

        A = np.arange(5, dtype=np.float32)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A[j] * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i], (exp, out[j, i]))


class TestGUFuncDecor(unittest.TestCase):
    def test_gufunc_decorator(self):
        @guvectorize(["void(float32, float32[:], float32[:], float32[:])"],
                     '(),(t),(t)->(t)', target='roc')
        def saxpy(a, x, y, out):
            for i in range(out.shape[0]):
                out[i] = a * x[i] + y[i]

        A = np.float32(2)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i])

        X = np.arange(10, dtype=np.float32)
        Y = np.arange(10, dtype=np.float32)
        out = saxpy(A, X, Y)

        for j in range(10):
            exp = A * X[j] + Y[j]
            self.assertTrue(exp == out[j], (exp, out[j]))

        A = np.arange(5, dtype=np.float32)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A[j] * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i], (exp, out[j, i]))


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

[((5288, 5303), 'unittest.main', 'unittest.main', ([], {}), '()\n', (5301, 5303), False, 'import unittest\n'), ((357, 407), 'numba.roc.vectorizers.HsaGUFuncVectorize', 'HsaGUFuncVectorize', (['ufunc_add_core', '"""(x),(x)->(x)"""'], {}), "(ufunc_add_core, '(x),(x)->(x)')\n", (375, 407), False, 'from numba.roc.vectorizers import HsaGUFuncVectorize\n'), ((657, 686), 'numpy.arange', 'np.arange', (['(100)'], {'dtype': 'np.intp'}), '(100, dtype=np.intp)\n', (666, 686), True, 'import numpy as np\n'), ((793, 831), 'numpy.testing.assert_equal', 'np.testing.assert_equal', (['expected', 'got'], {}), '(expected, got)\n', (816, 831), True, 'import numpy as np\n'), ((1104, 1142), 'numpy.testing.assert_equal', 'np.testing.assert_equal', (['expected', 'got'], {}), '(expected, got)\n', (1127, 1142), True, 'import numpy as np\n'), ((1421, 1459), 'numpy.testing.assert_equal', 'np.testing.assert_equal', (['expected', 'got'], {}), '(expected, got)\n', (1444, 1459), True, 'import numpy as np\n'), ((1608, 1640), 'numpy.arange', 'np.arange', (['(100)'], {'dtype': 'np.float32'}), '(100, dtype=np.float32)\n', (1617, 1640), True, 'import numpy as np\n'), ((1750, 1791), 'numpy.testing.assert_allclose', 'np.testing.assert_allclose', (['expected', 'got'], {}), '(expected, got)\n', (1776, 1791), True, 'import numpy as np\n'), ((2064, 2105), 'numpy.testing.assert_allclose', 'np.testing.assert_allclose', (['expected', 'got'], {}), '(expected, got)\n', (2090, 2105), True, 'import numpy as np\n'), ((2434, 2472), 'numba.roc.vectorizers.HsaGUFuncVectorize', 'HsaGUFuncVectorize', (['sum_row', '"""(n)->()"""'], {}), "(sum_row, '(n)->()')\n", (2452, 2472), False, 'from numba.roc.vectorizers import HsaGUFuncVectorize\n'), ((2921, 2964), 'numba.roc.vectorizers.HsaGUFuncVectorize', 'HsaGUFuncVectorize', (['axpy', '"""(),(t),(t)->(t)"""'], {}), "(axpy, '(),(t),(t)->(t)')\n", (2939, 2964), False, 'from numba.roc.vectorizers import HsaGUFuncVectorize\n'), ((3087, 3100), 'numpy.float32', 'np.float32', (['(2)'], {}), '(2)\n', (3097, 3100), True, 'import numpy as np\n'), ((3412, 3443), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (3421, 3443), True, 'import numpy as np\n'), ((3456, 3487), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (3465, 3487), True, 'import numpy as np\n'), ((3651, 3681), 'numpy.arange', 'np.arange', (['(5)'], {'dtype': 'np.float32'}), '(5, dtype=np.float32)\n', (3660, 3681), True, 'import numpy as np\n'), ((4091, 4194), 'numba.guvectorize', 'guvectorize', (["['void(float32, float32[:], float32[:], float32[:])']", '"""(),(t),(t)->(t)"""'], {'target': '"""roc"""'}), "(['void(float32, float32[:], float32[:], float32[:])'],\n    '(),(t),(t)->(t)', target='roc')\n", (4102, 4194), False, 'from numba import guvectorize\n'), ((4341, 4354), 'numpy.float32', 'np.float32', (['(2)'], {}), '(2)\n', (4351, 4354), True, 'import numpy as np\n'), ((4666, 4697), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (4675, 4697), True, 'import numpy as np\n'), ((4710, 4741), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (4719, 4741), True, 'import numpy as np\n'), ((4905, 4935), 'numpy.arange', 'np.arange', (['(5)'], {'dtype': 'np.float32'}), '(5, dtype=np.float32)\n', (4914, 4935), True, 'import numpy as np\n'), ((699, 728), 'numpy.arange', 'np.arange', (['(100)'], {'dtype': 'np.intp'}), '(100, dtype=np.intp)\n', (708, 728), True, 'import numpy as np\n'), ((909, 926), 'numpy.dtype', 'np.dtype', (['np.intp'], {}), '(np.intp)\n', (917, 926), True, 'import numpy as np\n'), ((1220, 1237), 'numpy.dtype', 'np.dtype', (['np.intp'], {}), '(np.intp)\n', (1228, 1237), True, 'import numpy as np\n'), ((1537, 1554), 'numpy.dtype', 'np.dtype', (['np.intp'], {}), '(np.intp)\n', (1545, 1554), True, 'import numpy as np\n'), ((1653, 1685), 'numpy.arange', 'np.arange', (['(100)'], {'dtype': 'np.float32'}), '(100, dtype=np.float32)\n', (1662, 1685), True, 'import numpy as np\n'), ((1869, 1889), 'numpy.dtype', 'np.dtype', (['np.float32'], {}), '(np.float32)\n', (1877, 1889), True, 'import numpy as np\n'), ((2183, 2203), 'numpy.dtype', 'np.dtype', (['np.float32'], {}), '(np.float32)\n', (2191, 2203), True, 'import numpy as np\n'), ((2572, 2602), 'numpy.arange', 'np.arange', (['(300)'], {'dtype': 'np.int32'}), '(300, dtype=np.int32)\n', (2581, 2602), True, 'import numpy as np\n'), ((3113, 3144), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (3122, 3144), True, 'import numpy as np\n'), ((3171, 3202), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (3180, 3202), True, 'import numpy as np\n'), ((3694, 3725), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (3703, 3725), True, 'import numpy as np\n'), ((3752, 3783), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (3761, 3783), True, 'import numpy as np\n'), ((4367, 4398), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (4376, 4398), True, 'import numpy as np\n'), ((4425, 4456), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (4434, 4456), True, 'import numpy as np\n'), ((4948, 4979), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (4957, 4979), True, 'import numpy as np\n'), ((5006, 5037), 'numpy.arange', 'np.arange', (['(10)'], {'dtype': 'np.float32'}), '(10, dtype=np.float32)\n', (5015, 5037), True, 'import numpy as np\n')]