add experimental-jit mode to test suite

conftest.py

@@ -3,6 +3,7 @@ import runpy
 import sys
 import tempfile
 import warnings
+import pathlib
 
 import nbformat
 import pytest
@@ -185,24 +186,10 @@ class IPyNbFile(pytest.File):
         pass
 
 
-if pytest_version >= 70000:
-    #   Since pytest 7.0, usage of `py.path.local` is deprecated and `pathlib.Path` should be used instead
-    import pathlib
-
-    def pytest_collect_file(file_path: pathlib.Path, parent):
-        glob_exprs = ["*demo*.ipynb", "*tutorial*.ipynb", "test_*.ipynb"]
-        if any(file_path.match(g) for g in glob_exprs):
-            return IPyNbFile.from_parent(path=file_path, parent=parent)
-
-else:
-
-    def pytest_collect_file(path, parent):
-        glob_exprs = ["*demo*.ipynb", "*tutorial*.ipynb", "test_*.ipynb"]
-        if any(path.fnmatch(g) for g in glob_exprs):
-            if pytest_version >= 50403:
-                return IPyNbFile.from_parent(fspath=path, parent=parent)
-            else:
-                return IPyNbFile(path, parent)
+def pytest_collect_file(file_path: pathlib.Path, parent):
+    glob_exprs = ["*demo*.ipynb", "*tutorial*.ipynb", "test_*.ipynb"]
+    if any(file_path.match(g) for g in glob_exprs):
+        return IPyNbFile.from_parent(path=file_path, parent=parent)
 
 
 #   Fixtures

src/pystencils/jit/cpu/cpujit.py

@@ -114,7 +114,7 @@ class CpuJit(JitBase):
         """
 
         #   Get the Code
-        module_name = f"{kernel.function_name}_jit"
+        module_name = f"{kernel.name}_jit"
         cpp_code = self._ext_module_builder.render_module(kernel, module_name)
 
         #   Get compiler information
@@ -124,11 +124,12 @@ class CpuJit(JitBase):
         lib_suffix = f"{so_abi}.so"
 
         #   Compute Code Hash
-        code_utf8 = cpp_code.encode("utf-8")
+        code_utf8: bytes = cpp_code.encode("utf-8")
+        compiler_utf8: bytes = (" ".join([self._cxx] + self._cxx_fixed_flags)).encode("utf-8")
         import hashlib
 
-        code_hash = hashlib.sha256(code_utf8)
-        module_stem = f"module_{code_hash.hexdigest()}"
+        module_hash = hashlib.sha256(code_utf8 + compiler_utf8)
+        module_stem = f"module_{module_hash.hexdigest()}"
 
         def compile_and_load(module_dir: Path):
             cpp_file = module_dir / f"{module_stem}.cpp"

src/pystencils/jit/cpu/cpujit_pybind11.py

@@ -31,12 +31,9 @@ class Pybind11KernelModuleBuilder(ExtensionModuleBuilderBase):
     def __init__(
         self,
         compiler_info: CompilerInfo,
-        strict_scalar_types: bool = False,
     ):
         self._compiler_info = compiler_info
 
-        self._strict_scalar_types = strict_scalar_types
-
         self._actual_field_types: dict[Field, PsType]
         self._param_binds: list[str]
         self._public_params: list[str]
@@ -63,7 +60,7 @@ class Pybind11KernelModuleBuilder(ExtensionModuleBuilderBase):
             includes="\n".join(includes),
             restrict_qualifier=self._compiler_info.restrict_qualifier(),
             module_name=module_name,
-            kernel_name=kernel.function_name,
+            kernel_name=kernel.name,
             param_binds=", ".join(self._param_binds),
             public_params=", ".join(self._public_params),
             param_check_lines=indent("\n".join(self._param_check_lines), prefix="    "),
@@ -118,8 +115,6 @@ class Pybind11KernelModuleBuilder(ExtensionModuleBuilderBase):
 
     def _add_scalar_param(self, sc_param: Parameter):
         param_bind = f'py::arg("{sc_param.name}")'
-        if self._strict_scalar_types:
-            param_bind += ".noconvert()"
         self._param_binds.append(param_bind)
 
         kernel_param = f"{sc_param.dtype.c_string()} {sc_param.name}"

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