test_cpujit.py

import sympy as sp
import numpy as np
from pystencils import create_kernel, Assignment, fields
from pystencils.jit import CpuJit


def test_basic_cpu_kernel(tmp_path):
    jit = CpuJit.create(objcache=tmp_path)

    f, g = fields("f, g: [2D]")
    asm = Assignment(f.center(), 2.0 * g.center())
    ker = create_kernel(asm)
    kfunc = jit.compile(ker)

    rng = np.random.default_rng()
    f_arr = rng.random(size=(34, 26), dtype="float64")
    g_arr = np.zeros_like(f_arr)

    kfunc(f=f_arr, g=g_arr)

    np.testing.assert_almost_equal(g_arr, 2.0 * f_arr)