test_jacobi_llvm.py 3.36 KB
Newer Older
1

2
import pytest
3
4
5
6
7
8
9
10
try:
    from pystencils.llvm.llvmjit import generate_and_jit
    from pystencils.llvm import create_kernel, make_python_function
    from pystencils.cpu.cpujit import get_llc_command
    from pystencils import Assignment, Field, show_code
    import numpy as np
except ModuleNotFoundError:
    pytest.importorskip("llvmlite")
Martin Bauer's avatar
Martin Bauer committed
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


def test_jacobi_fixed_field_size():
    size = (30, 20)

    src_field_llvm = np.random.rand(*size)
    src_field_py = np.copy(src_field_llvm)
    dst_field_llvm = np.zeros(size)
    dst_field_py = np.zeros(size)

    f = Field.create_from_numpy_array("f", src_field_llvm)
    d = Field.create_from_numpy_array("d", dst_field_llvm)

    jacobi = Assignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
    ast = create_kernel([jacobi])

    for x in range(1, size[0] - 1):
        for y in range(1, size[1] - 1):
            dst_field_py[x, y] = 0.25 * (src_field_py[x - 1, y] + src_field_py[x + 1, y] +
                                         src_field_py[x, y - 1] + src_field_py[x, y + 1])

    jit = generate_and_jit(ast)
    jit('kernel', dst_field_llvm, src_field_llvm)
    error = np.sum(np.abs(dst_field_py - dst_field_llvm))
    np.testing.assert_almost_equal(error, 0.0)


39
@pytest.mark.skipif(not get_llc_command(), reason="Tests requires llc in $PATH")
40
41
42
def test_jacobi_fixed_field_size_gpu():
    size = (30, 20)

43
44
45
    import pycuda.autoinit  # noqa
    from pycuda.gpuarray import to_gpu

46
47
48
49
50
    src_field_llvm = np.random.rand(*size)
    src_field_py = np.copy(src_field_llvm)
    dst_field_llvm = np.zeros(size)
    dst_field_py = np.zeros(size)

51
52
53
54
55
    f = Field.create_from_numpy_array("f", src_field_py)
    d = Field.create_from_numpy_array("d", dst_field_py)

    src_field_llvm = to_gpu(src_field_llvm)
    dst_field_llvm = to_gpu(dst_field_llvm)
56
57
58

    jacobi = Assignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
    ast = create_kernel([jacobi], target='gpu')
59
    show_code(ast)
60
61
62
63
64
65
66
67

    for x in range(1, size[0] - 1):
        for y in range(1, size[1] - 1):
            dst_field_py[x, y] = 0.25 * (src_field_py[x - 1, y] + src_field_py[x + 1, y] +
                                         src_field_py[x, y - 1] + src_field_py[x, y + 1])

    jit = generate_and_jit(ast)
    jit('kernel', dst_field_llvm, src_field_llvm)
68
    error = np.sum(np.abs(dst_field_py - dst_field_llvm.get()))
69
70
71
    np.testing.assert_almost_equal(error, 0.0)


72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def test_jacobi_variable_field_size():
    size = (3, 3, 3)
    f = Field.create_generic("f", 3)
    d = Field.create_generic("d", 3)
    jacobi = Assignment(d[0, 0, 0], (f[1, 0, 0] + f[-1, 0, 0] + f[0, 1, 0] + f[0, -1, 0]) / 4)
    ast = create_kernel([jacobi])

    src_field_llvm = np.random.rand(*size)
    src_field_py = np.copy(src_field_llvm)
    dst_field_llvm = np.zeros(size)
    dst_field_py = np.zeros(size)

    for x in range(1, size[0] - 1):
        for y in range(1, size[1] - 1):
            for z in range(1, size[2] - 1):
                dst_field_py[x, y, z] = 0.25 * (src_field_py[x - 1, y, z] + src_field_py[x + 1, y, z] +
                                                src_field_py[x, y - 1, z] + src_field_py[x, y + 1, z])

    kernel = make_python_function(ast, {'f': src_field_llvm, 'd': dst_field_llvm})
    kernel()
    error = np.sum(np.abs(dst_field_py - dst_field_llvm))
    np.testing.assert_almost_equal(error, 0.0)
94
95
96
97


if __name__ == "__main__":
    test_jacobi_fixed_field_size_gpu()