Skip to content
Snippets Groups Projects
Select Git revision
  • master default protected
1 result
Blame Permalink
test_benchmark.py 2.13 KiB
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 






import subprocess
import numpy as np
import sympy as sp
import tempfile

import pytest
import pystencils as ps
from pathlib import Path

from pystencils_benchmark import Compiler
import pystencils_benchmark as pb


compilers = (Compiler.GCC, Compiler.GCCdebug, Compiler.Clang)
config_kwargs = ({},
                 {'cpu_vectorize_info': {'instruction_set': 'best'}},
                 {'cpu_vectorize_info': {'instruction_set': 'best',
                                         'assume_aligned': True}})


@pytest.mark.parametrize('compiler', compilers)
@pytest.mark.parametrize('config_kwarg', config_kwargs)
def test_generate(compiler, config_kwarg):
    a, b, c = ps.fields(a=np.ones(4000000), b=np.ones(4000000), c=np.ones(4000000))
    alpha = sp.symbols('alpha')

    @ps.kernel_config(ps.CreateKernelConfig(**config_kwarg))
    def vadd():
        a[0] @= b[0] + c[0]
    kernel_vadd = ps.create_kernel(**vadd)

    @ps.kernel_config(ps.CreateKernelConfig(**config_kwarg))
    def daxpy():
        b[0] @= alpha * a[0] + b[0]
    kernel_daxpy = ps.create_kernel(**daxpy)

    with tempfile.TemporaryDirectory(dir=Path.cwd()) as temp_dir:
        temp_dir = Path(temp_dir)
        pb.cpu.generate_benchmark([kernel_vadd, kernel_daxpy], temp_dir, compiler=compiler)
        subprocess.run(['make', '-C', f'{temp_dir}'], check=True)
        subprocess.run([f'{temp_dir}/benchmark-{compiler.name}', '10'], check=True)


gpu_kwargs = ({}, {'launch_bounds': (256,)}, {'launch_bounds': (256, 2)})


@pytest.mark.parametrize('kwargs', gpu_kwargs)
def test_generate_gpu(kwargs):
    compiler = Compiler.NVCC
    a, b, c = ps.fields(a=np.ones(4000000), b=np.ones(4000000), c=np.ones(4000000))

    @ps.kernel_config(ps.CreateKernelConfig(target=ps.Target.GPU))
    def vadd():
        a[0] @= b[0] + c[0]
    kernel_vadd = ps.create_kernel(**vadd)

    with tempfile.TemporaryDirectory(dir=Path.cwd()) as temp_dir:
        temp_dir = Path(temp_dir)
        pb.gpu.generate_benchmark(kernel_vadd, temp_dir, compiler=compiler, **kwargs)
        # subprocess.run(['make', '-C', f'{temp_dir}'], check=True)
        # subprocess.run([f'{temp_dir}/benchmark-{compiler.name}', '10'], check=True)