diff --git a/pystencils_tests/test_jacobi_llvm.py b/pystencils_tests/test_llvm.py similarity index 72% rename from pystencils_tests/test_jacobi_llvm.py rename to pystencils_tests/test_llvm.py index 953432c11a91616dac17d776a1b1e822ea5f8848..54c60ff45dd99c5157977acdb86b052d3dd67214 100644 --- a/pystencils_tests/test_jacobi_llvm.py +++ b/pystencils_tests/test_llvm.py @@ -6,11 +6,11 @@ try: from pystencils.cpu.cpujit import get_llc_command from pystencils import Assignment, Field, show_code import numpy as np + import sympy as sp except ModuleNotFoundError: pytest.importorskip("llvmlite") - def test_jacobi_fixed_field_size(): size = (30, 20) @@ -93,5 +93,50 @@ def test_jacobi_variable_field_size(): np.testing.assert_almost_equal(error, 0.0) +def test_pow_llvm(): + size = (30, 20) + + src_field_llvm = 4 * np.ones(size) + dst_field_llvm = np.zeros(size) + + f = Field.create_from_numpy_array("f", src_field_llvm) + d = Field.create_from_numpy_array("d", dst_field_llvm) + + ur = Assignment(d[0, 0], sp.Pow(f[0, 0], -1.0)) + ast = create_kernel([ur]) + + jit = generate_and_jit(ast) + jit('kernel', dst_field_llvm, src_field_llvm) + assert np.all(0.25 == dst_field_llvm) + + ur = Assignment(d[0, 0], sp.Pow(f[0, 0], 0.5)) + ast = create_kernel([ur]) + + jit = generate_and_jit(ast) + jit('kernel', dst_field_llvm, src_field_llvm) + assert np.all(2.0 == dst_field_llvm) + + ur = Assignment(d[0, 0], sp.Pow(f[0, 0], 2.0)) + ast = create_kernel([ur]) + + jit = generate_and_jit(ast) + jit('kernel', dst_field_llvm, src_field_llvm) + assert np.all(16.0 == dst_field_llvm) + + ur = Assignment(d[0, 0], sp.Pow(f[0, 0], 3.0)) + ast = create_kernel([ur]) + + jit = generate_and_jit(ast) + jit('kernel', dst_field_llvm, src_field_llvm) + assert np.all(64.0 == dst_field_llvm) + + ur = Assignment(d[0, 0], sp.Pow(f[0, 0], 4.0)) + ast = create_kernel([ur]) + + jit = generate_and_jit(ast) + jit('kernel', dst_field_llvm, src_field_llvm) + assert np.all(256.0 == dst_field_llvm) + + if __name__ == "__main__": test_jacobi_fixed_field_size_gpu()