Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_intrinsics.py 3.2 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 roc
from numba.core.errors import TypingError
import operator as oper
import unittest

_WAVESIZE = roc.get_context().agent.wavefront_size

@roc.jit(device=True)
def shuffle_up(val, width):
    tid = roc.get_local_id(0)
    roc.wavebarrier()
    idx = (tid + width) % _WAVESIZE
    res = roc.ds_permute(idx, val)
    return res

@roc.jit(device=True)
def shuffle_down(val, width):
    tid = roc.get_local_id(0)
    roc.wavebarrier()
    idx = (tid - width) % _WAVESIZE
    res = roc.ds_permute(idx, val)
    return res

@roc.jit(device=True)
def broadcast(val, from_lane):
    tid = roc.get_local_id(0)
    roc.wavebarrier()
    res = roc.ds_bpermute(from_lane, val)
    return res

def gen_kernel(shuffunc):
    @roc.jit
    def kernel(inp, outp, amount):
        tid = roc.get_local_id(0)
        val = inp[tid]
        outp[tid] = shuffunc(val, amount)
    return kernel


class TestDsPermute(unittest.TestCase):

    def test_ds_permute(self):

        inp = np.arange(_WAVESIZE).astype(np.int32)
        outp = np.zeros_like(inp)

        for shuffler, op in [(shuffle_down, oper.neg), (shuffle_up, oper.pos)]:
            kernel = gen_kernel(shuffler)
            for shuf in range(-_WAVESIZE, _WAVESIZE):
                kernel[1, _WAVESIZE](inp, outp, shuf)
                np.testing.assert_allclose(outp, np.roll(inp, op(shuf)))

    def test_ds_permute_random_floats(self):

        inp = np.linspace(0, 1, _WAVESIZE).astype(np.float32)
        outp = np.zeros_like(inp)

        for shuffler, op in [(shuffle_down, oper.neg), (shuffle_up, oper.pos)]:
            kernel = gen_kernel(shuffler)
            for shuf in range(-_WAVESIZE, _WAVESIZE):
                kernel[1, _WAVESIZE](inp, outp, shuf)
                np.testing.assert_allclose(outp, np.roll(inp, op(shuf)))

    def test_ds_permute_type_safety(self):
        """ Checks that float64's are not being downcast to float32"""
        kernel = gen_kernel(shuffle_down)
        inp = np.linspace(0, 1, _WAVESIZE).astype(np.float64)
        outp = np.zeros_like(inp)
        with self.assertRaises(TypingError) as e:
            kernel[1, _WAVESIZE](inp, outp, 1)
        errmsg = e.exception.msg
        self.assertIn('Invalid use of Function', errmsg)
        self.assertIn('with argument(s) of type(s): (float64, int64)', errmsg)

    def test_ds_bpermute(self):

        @roc.jit
        def kernel(inp, outp, lane):
            tid = roc.get_local_id(0)
            val = inp[tid]
            outp[tid] = broadcast(val, lane)

        inp = np.arange(_WAVESIZE).astype(np.int32)
        outp = np.zeros_like(inp)
        for lane in range(0, _WAVESIZE):
            kernel[1, _WAVESIZE](inp, outp, lane)
            np.testing.assert_allclose(outp, lane)

    def test_ds_bpermute_random_floats(self):

        @roc.jit
        def kernel(inp, outp, lane):
            tid = roc.get_local_id(0)
            val = inp[tid]
            outp[tid] = broadcast(val, lane)

        inp = np.linspace(0, 1, _WAVESIZE).astype(np.float32)
        outp = np.zeros_like(inp)

        for lane in range(0, _WAVESIZE):
            kernel[1, _WAVESIZE](inp, outp, lane)
            np.testing.assert_allclose(outp, inp[lane])


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