Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_ufuncbuilding.py 3.29 KB
Newer Older
dugupeiwen's avatar
update roc  
dugupeiwen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108 
import numpy as np

from numba import vectorize
from numba.roc.vectorizers import HsaVectorize
from numba.roc.dispatch import HsaUFuncDispatcher
import unittest


def ufunc_add_core(a, b):
    return a + b


class TestUFuncBuilding(unittest.TestCase):
    def test_ufunc_building(self):
        ufbldr = HsaVectorize(ufunc_add_core)
        ufbldr.add("float32(float32, float32)")
        ufbldr.add("intp(intp, intp)")
        ufunc = ufbldr.build_ufunc()
        self.assertIsInstance(ufunc, HsaUFuncDispatcher)

        # 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 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)


class TestVectorizeDecor(unittest.TestCase):
    def test_vectorize_decor(self):
        @vectorize(["float32(float32, float32, float32)",
                    "intp(intp, intp, intp)"],
                   target='roc')
        def axpy(a, x, y):
            return a * x + y


        self.assertIsInstance(axpy, HsaUFuncDispatcher)
        # Test integer version
        A = np.arange(100, dtype=np.intp)
        X = np.arange(100, dtype=np.intp) + 1
        Y = np.arange(100, dtype=np.intp) + 2
        expected = A * X + Y
        got = axpy(A, X, Y)

        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)
        X = np.arange(100, dtype=np.float32) + 1
        Y = np.arange(100, dtype=np.float32) + 2
        expected = A * X + Y
        got = axpy(A, X, Y)

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


class TestVectorizeScalar(unittest.TestCase):
    def test_scalar_input(self):
        @vectorize(["float32(float32, float32, float32)",
                    "intp(intp, intp, intp)"],
                   target='roc')
        def axpy(a, x, y):
            return a * x + y

        self.assertIsInstance(axpy, HsaUFuncDispatcher)
        # Test integer version
        A = 2
        X = np.arange(100, dtype=np.intp) + 1
        Y = np.arange(100, dtype=np.intp) + 2
        expected = A * X + Y
        got = axpy(A, X, Y)

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

        # Test real version
        A = 2.3
        X = np.arange(100, dtype=np.float32) + 1
        Y = np.arange(100, dtype=np.float32) + 2
        expected = A * X + Y
        got = axpy(A, X, Y)

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


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