Object-Oriented CPU JIT API and Prototype Implementation

src/pystencils/sympyextensions/__init__.py

 from .astnodes import ConditionalFieldAccess
-from .typed_sympy import TypedSymbol, tcast
+from .typed_sympy import TypedSymbol, CastFunc, tcast, DynamicType
 from .pointers import mem_acc
 
 from .math import (
@@ -34,6 +34,7 @@ from .math import (
 __all__ = [
     "ConditionalFieldAccess",
     "TypedSymbol",
+    "CastFunc",
     "tcast",
     "mem_acc",
     "remove_higher_order_terms",
@@ -61,5 +62,6 @@ __all__ = [
     "count_operations_in_ast",
     "common_denominator",
     "get_symmetric_part",
-    "SymbolCreator"
+    "SymbolCreator",
+    "DynamicType"
 ]

tests/fixtures.py

@@ -15,7 +15,6 @@ by your tests:
 import pytest
 
 from types import ModuleType
-from dataclasses import replace
 
 import pystencils as ps
 
@@ -32,6 +31,14 @@ AVAILABLE_TARGETS += ps.Target.available_vector_cpu_targets()
 TARGET_IDS = [t.name for t in AVAILABLE_TARGETS]
 
 
+def pytest_addoption(parser: pytest.Parser):
+    parser.addoption(
+        "--experimental-cpu-jit",
+        dest="experimental_cpu_jit",
+        action="store_true"
+    )
+
+
 @pytest.fixture(params=AVAILABLE_TARGETS, ids=TARGET_IDS)
 def target(request) -> ps.Target:
     """Provides all code generation targets available on the current hardware"""
@@ -39,7 +46,7 @@ def target(request) -> ps.Target:
 
 
 @pytest.fixture
-def gen_config(target: ps.Target):
+def gen_config(request: pytest.FixtureRequest, target: ps.Target):
     """Default codegen configuration for the current target.
 
     For GPU targets, set default indexing options.
@@ -52,6 +59,11 @@ def gen_config(target: ps.Target):
         gen_config.cpu.vectorize.enable = True
         gen_config.cpu.vectorize.assume_inner_stride_one = True
 
+    if target.is_cpu() and request.config.getoption("experimental_cpu_jit"):
+        from pystencils.jit.cpu import CpuJit, GccInfo
+
+        gen_config.jit = CpuJit.create(compiler_info=GccInfo(target=target))
+
     return gen_config
 
 

tests/jit/test_cpujit.py

+import pytest
+
+import sympy as sp
+import numpy as np
+from pystencils import create_kernel, Assignment, fields, Field
+from pystencils.jit import CpuJit
+
+
+@pytest.fixture
+def cpu_jit(tmp_path) -> CpuJit:
+    return CpuJit.create(objcache=tmp_path)
+
+
+def test_basic_cpu_kernel(cpu_jit):
+    f, g = fields("f, g: [2D]")
+    asm = Assignment(f.center(), 2.0 * g.center())
+    ker = create_kernel(asm)
+    kfunc = cpu_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)
+
+
+def test_argument_type_error(cpu_jit):
+    f, g = fields("f, g: [2D]")
+    c = sp.Symbol("c")
+    asm = Assignment(f.center(), c * g.center())
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    arr_fp16 = np.zeros((23, 12), dtype="float16")
+    arr_fp32 = np.zeros((23, 12), dtype="float32")
+    arr_fp64 = np.zeros((23, 12), dtype="float64")
+
+    with pytest.raises(TypeError):
+        kfunc(f=arr_fp32, g=arr_fp64, c=2.0)
+
+    with pytest.raises(TypeError):
+        kfunc(f=arr_fp64, g=arr_fp32, c=2.0)
+
+    with pytest.raises(TypeError):
+        kfunc(f=arr_fp16, g=arr_fp16, c=2.0)
+
+    #   Wrong scalar types are OK, though
+    kfunc(f=arr_fp64, g=arr_fp64, c=np.float16(1.0))
+
+
+def test_fixed_shape(cpu_jit):
+    a = np.zeros((12, 23), dtype="float64")
+    b = np.zeros((13, 21), dtype="float64")
+    
+    f = Field.create_from_numpy_array("f", a)
+    g = Field.create_from_numpy_array("g", a)
+
+    asm = Assignment(f.center(), 2.0 * g.center())
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    kfunc(f=a, g=a)
+
+    with pytest.raises(ValueError):
+        kfunc(f=b, g=a)
+
+    with pytest.raises(ValueError):
+        kfunc(f=a, g=b)
+
+
+def test_fixed_index_shape(cpu_jit):
+    f, g = fields("f(3), g(2, 2): [2D]")
+
+    asm = Assignment(f.center(1), g.center(0, 0) + g.center(0, 1) + g.center(1, 0) + g.center(1, 1))
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    f_arr = np.zeros((12, 14, 3))
+    g_arr = np.zeros((12, 14, 2, 2))
+    kfunc(f=f_arr, g=g_arr)
+
+    with pytest.raises(ValueError):
+        f_arr = np.zeros((12, 14, 2))
+        g_arr = np.zeros((12, 14, 2, 2))
+        kfunc(f=f_arr, g=g_arr)
+
+    with pytest.raises(ValueError):
+        f_arr = np.zeros((12, 14, 3))
+        g_arr = np.zeros((12, 14, 4))
+        kfunc(f=f_arr, g=g_arr)
+
+    with pytest.raises(ValueError):
+        f_arr = np.zeros((12, 14, 3))
+        g_arr = np.zeros((12, 14, 1, 3))
+        kfunc(f=f_arr, g=g_arr)