test_benchmark.py

import subprocess
import numpy as np
import sympy as sp
import pystencils as ps
from pystencils_benchmark import generate_benchmark, Compiler
from pathlib import Path


def test_generate():
    path = Path.cwd() / 'test'
    path.mkdir(parents=True, exist_ok=True)
    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())
    def vadd():
        a[0] @= b[0] + c[0]
    kernel_vadd = ps.create_kernel(**vadd)

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

    for compiler in [Compiler.GCC, Compiler.GCCdebug, Compiler.Clang]:
        generate_benchmark([kernel_vadd, kernel_daxpy], path, compiler=compiler)
        subprocess.run(['make', '-C', f'{path}'], check=True)
        subprocess.run([f'{path}/benchmark-{compiler.name}', '10'], check=True)