Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_simple.py 3.65 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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141 
import numpy as np
from numba import roc
from numba.roc.hsadrv.error import HsaKernelLaunchError
import unittest


class TestSimple(unittest.TestCase):

    def test_array_access(self):
        magic_token = 123

        @roc.jit
        def udt(output):
            output[0] = magic_token

        out = np.zeros(1, dtype=np.intp)
        udt[1, 1](out)

        self.assertEqual(out[0], magic_token)

    def test_array_access_2d(self):
        magic_token = 123

        @roc.jit
        def udt(output):
            for i in range(output.shape[0]):
                for j in range(output.shape[1]):
                    output[i, j] = magic_token

        out = np.zeros((10, 10), dtype=np.intp)
        udt[1, 1](out)
        np.testing.assert_equal(out, magic_token)

    def test_array_access_3d(self):
        magic_token = 123

        @roc.jit
        def udt(output):
            for i in range(output.shape[0]):
                for j in range(output.shape[1]):
                    for k in range(output.shape[2]):
                        output[i, j, k] = magic_token

        out = np.zeros((10, 10, 10), dtype=np.intp)
        udt[1, 1](out)
        np.testing.assert_equal(out, magic_token)

    def test_global_id(self):
        @roc.jit
        def udt(output):
            global_id = roc.get_global_id(0)
            output[global_id] = global_id

        # Allocate extra space to track bad indexing
        out = np.zeros(100 + 2, dtype=np.intp)
        udt[10, 10](out[1:-1])

        np.testing.assert_equal(out[1:-1], np.arange(100))

        self.assertEqual(out[0], 0)
        self.assertEqual(out[-1], 0)

    def test_local_id(self):
        @roc.jit
        def udt(output):
            global_id = roc.get_global_id(0)
            local_id = roc.get_local_id(0)
            output[global_id] = local_id

        # Allocate extra space to track bad indexing
        out = np.zeros(100 + 2, dtype=np.intp)
        udt[10, 10](out[1:-1])

        subarr = out[1:-1]

        for parted in np.split(subarr, 10):
            np.testing.assert_equal(parted, np.arange(10))

        self.assertEqual(out[0], 0)
        self.assertEqual(out[-1], 0)

    def test_group_id(self):
        @roc.jit
        def udt(output):
            global_id = roc.get_global_id(0)
            group_id = roc.get_group_id(0)
            output[global_id] = group_id + 1

        # Allocate extra space to track bad indexing
        out = np.zeros(100 + 2, dtype=np.intp)
        udt[10, 10](out[1:-1])

        subarr = out[1:-1]

        for i, parted in enumerate(np.split(subarr, 10), start=1):
            np.testing.assert_equal(parted, i)

        self.assertEqual(out[0], 0)
        self.assertEqual(out[-1], 0)


    def test_workdim(self):
        @roc.jit
        def udt(output):
            global_id = roc.get_global_id(0)
            workdim = roc.get_work_dim()
            output[global_id] = workdim

        out = np.zeros(10, dtype=np.intp)
        udt[1, 10](out)
        np.testing.assert_equal(out, 1)

        @roc.jit
        def udt2(output):
            g0 = roc.get_global_id(0)
            g1 = roc.get_global_id(1)
            output[g0, g1] = roc.get_work_dim()

        out = np.zeros((2, 5), dtype=np.intp)
        udt2[(1, 1), (2, 5)](out)
        np.testing.assert_equal(out, 2)

    def test_empty_kernel(self):
        @roc.jit
        def udt():
            pass

        udt[1, 1]()

    def test_workgroup_oversize(self):
        @roc.jit
        def udt():
            pass

        with self.assertRaises(HsaKernelLaunchError) as raises:
            udt[1, 2**30]()
        self.assertIn("Try reducing group-size", str(raises.exception))


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