diff --git a/docs/source/api/jit.md b/docs/source/api/jit.md
new file mode 100644
index 0000000000000000000000000000000000000000..06ea0cbaf4e5df4c031a102b248b9c1fdc57efad
--- /dev/null
+++ b/docs/source/api/jit.md
@@ -0,0 +1,102 @@
+# JIT Compilation
+
+## Base Infrastructure
+
+```{eval-rst}
+.. module:: pystencils.jit
+
+.. autosummary::
+ :toctree: generated
+ :nosignatures:
+ :template: autosummary/entire_class.rst
+
+ KernelWrapper
+ JitBase
+ NoJit
+
+.. autodata:: no_jit
+```
+
+## Legacy CPU JIT
+
+The legacy CPU JIT Compiler is a leftover from pystencils 1.3
+which at the moment still drives most CPU JIT-compilation within the package,
+until the new JIT compiler is ready to take over.
+
+```{eval-rst}
+.. autosummary::
+ :toctree: generated
+ :nosignatures:
+ :template: autosummary/entire_class.rst
+
+ LegacyCpuJit
+```
+
+## CPU Just-In-Time Compiler
+
+:::{note}
+The new CPU JIT compiler is still considered experimental and not yet adopted by most of pystencils.
+While the APIs described here will (probably) become the default for pystencils 2.0
+and can (and should) already be used for testing,
+the current implementation is still *very slow*.
+For more information, see [issue !120](https://i10git.cs.fau.de/pycodegen/pystencils/-/issues/120).
+:::
+
+To configure and create an instance of the CPU JIT compiler, use the `CpuJit.create` factory method:
+
+:::{card}
+```{eval-rst}
+.. autofunction:: pystencils.jit.CpuJit.create
+ :no-index:
+```
+:::
+
+### Compiler Infos
+
+The CPU JIT compiler invokes a host C++ compiler to compile and link a Python extension
+module containing the generated kernel.
+The properties of the host compiler are defined in a `CompilerInfo` object.
+To select a custom host compiler and customize its options, set up and pass
+a custom compiler info object to `CpuJit.create`.
+
+```{eval-rst}
+.. module:: pystencils.jit.cpu.compiler_info
+
+.. autosummary::
+ :toctree: generated
+ :nosignatures:
+ :template: autosummary/entire_class.rst
+
+ CompilerInfo
+ GccInfo
+ ClangInfo
+```
+
+### Implementation
+
+```{eval-rst}
+.. module:: pystencils.jit.cpu
+
+.. autosummary::
+ :toctree: generated
+ :nosignatures:
+ :template: autosummary/entire_class.rst
+
+ CpuJit
+ cpujit.ExtensionModuleBuilderBase
+```
+
+## CuPy-based GPU JIT
+
+```{eval-rst}
+.. module:: pystencils.jit.gpu_cupy
+
+.. autosummary::
+ :toctree: generated
+ :nosignatures:
+ :template: autosummary/entire_class.rst
+
+ CupyJit
+ CupyKernelWrapper
+ LaunchGrid
+```
diff --git a/docs/source/api/jit.rst b/docs/source/api/jit.rst
deleted file mode 100644
index f2e271db3917825ed7fb6a97c38633a847b8bbfc..0000000000000000000000000000000000000000
--- a/docs/source/api/jit.rst
+++ /dev/null
@@ -1,40 +0,0 @@
-JIT Compilation
-===============
-
-.. module:: pystencils.jit
-
-Base Infrastructure
--------------------
-
-.. autosummary::
- :toctree: generated
- :nosignatures:
- :template: autosummary/entire_class.rst
-
- KernelWrapper
- JitBase
- NoJit
-
-.. autodata:: no_jit
-
-Legacy CPU JIT
---------------
-
-.. autosummary::
- :toctree: generated
- :nosignatures:
- :template: autosummary/entire_class.rst
-
- LegacyCpuJit
-
-CuPy-based GPU JIT
-------------------
-
-.. autosummary::
- :toctree: generated
- :nosignatures:
- :template: autosummary/entire_class.rst
-
- CupyJit
- CupyKernelWrapper
- LaunchGrid
diff --git a/src/pystencils/codegen/config.py b/src/pystencils/codegen/config.py
index 6587f16b7650d3e3e3f64116cf464ad4e50fa7c6..bce0757313b597cb9f4413e84652b9878c39a2d5 100644
--- a/src/pystencils/codegen/config.py
+++ b/src/pystencils/codegen/config.py
@@ -4,7 +4,7 @@ from warnings import warn
from abc import ABC
from collections.abc import Collection
-from typing import TYPE_CHECKING, Sequence, Generic, TypeVar, Callable, Any, cast, TypeGuard
+from typing import TYPE_CHECKING, Sequence, Generic, TypeVar, Callable, Any, cast
from dataclasses import dataclass, InitVar, fields
from .target import Target
@@ -207,7 +207,9 @@ class Category(Generic[Category_T]):
setattr(obj, self._lookup, cat.copy() if cat is not None else None)
-class _AUTO_TYPE: ... # noqa: E701
+class _AUTO_TYPE:
+ def __repr__(self) -> str:
+ return "AUTO" # for pretty-printing in the docs
AUTO = _AUTO_TYPE()
diff --git a/src/pystencils/jit/cpu/compiler_info.py b/src/pystencils/jit/cpu/compiler_info.py
index 6ed568d72d28b4254825bd3e84901f5f3b48c5be..061f37af50e9903707f823b98c390b89230efc33 100644
--- a/src/pystencils/jit/cpu/compiler_info.py
+++ b/src/pystencils/jit/cpu/compiler_info.py
@@ -8,28 +8,43 @@ from ...codegen.target import Target
@dataclass
class CompilerInfo(ABC):
+ """Base class for compiler infos."""
+
openmp: bool = True
+ """Enable/disable OpenMP compilation"""
optlevel: str | None = "fast"
+ """Compiler optimization level"""
cxx_standard: str = "c++11"
+ """C++ language standard to be compiled with"""
target: Target = Target.CurrentCPU
+ """Hardware target to compile for.
+
+ Here, `Target.CurrentCPU` represents the current hardware,
+ which is reflected by ``-march=native`` in GNU-like compilers.
+ """
@abstractmethod
- def cxx(self) -> str: ...
+ def cxx(self) -> str:
+ """Path to the executable of this compiler"""
@abstractmethod
- def cxxflags(self) -> list[str]: ...
+ def cxxflags(self) -> list[str]:
+ """Compiler flags affecting C++ compilation"""
@abstractmethod
- def linker_flags(self) -> list[str]: ...
+ def linker_flags(self) -> list[str]:
+ """Flags affecting linkage of the extension module"""
@abstractmethod
- def include_flags(self, include_dirs: Sequence[str]) -> list[str]: ...
+ def include_flags(self, include_dirs: Sequence[str]) -> list[str]:
+ """Convert a list of include directories into corresponding compiler flags"""
@abstractmethod
- def restrict_qualifier(self) -> str: ...
+ def restrict_qualifier(self) -> str:
+ """*restrict* memory qualifier recognized by this compiler"""
class _GnuLikeCliCompiler(CompilerInfo):
@@ -67,13 +82,18 @@ class _GnuLikeCliCompiler(CompilerInfo):
class GccInfo(_GnuLikeCliCompiler):
+ """Compiler info for the GNU Compiler Collection C++ compiler (``g++``)."""
+
def cxx(self) -> str:
return "g++"
@dataclass
class ClangInfo(_GnuLikeCliCompiler):
+ """Compiler info for the LLVM C++ compiler (``clang``)."""
+
llvm_version: int | None = None
+ """Major version number of the LLVM installation providing the compiler."""
def cxx(self) -> str:
if self.llvm_version is None:
@@ -81,5 +101,5 @@ class ClangInfo(_GnuLikeCliCompiler):
else:
return f"clang-{self.llvm_version}"
- def cxxflags(self) -> list[str]:
- return super().cxxflags() + ["-lstdc++"]
+ def linker_flags(self) -> list[str]:
+ return super().linker_flags() + ["-lstdc++"]
diff --git a/src/pystencils/jit/cpu/cpujit.py b/src/pystencils/jit/cpu/cpujit.py
index 64045d54951639b47e7beb3d2d99def8db6ce304..5fce101ee716de053df23687580b312d3ba76ac4 100644
--- a/src/pystencils/jit/cpu/cpujit.py
+++ b/src/pystencils/jit/cpu/cpujit.py
@@ -16,24 +16,45 @@ from .compiler_info import CompilerInfo, GccInfo
class CpuJit(JitBase):
"""Just-in-time compiler for CPU kernels.
- :Creation:
- In most cases, objects of this class should be instantiated using the `create` factory method.
-
- :Implementation Details:
-
- The `CpuJit` combines two separate components:
- - The *extension module builder* produces the code of the dynamically built extension module containing
- the kernel
- - The *compiler info* describes the system compiler used to compile and link that extension module.
-
- Both can be dynamically exchanged.
+ **Creation**
+
+ To configure and create a CPU JIT compiler instance, use the `create` factory method.
+
+ **Implementation Details**
+
+ The `CpuJit` class acts as an orchestrator between two components:
+
+ - The *extension module builder* produces the code of the dynamically built extension module
+ that contains the kernel and its invocation wrappers;
+ - The *compiler info* describes the host compiler used to compile and link that extension module.
+
+ Args:
+ compiler_info: The compiler info object defining the capabilities
+ and command-line interface of the host compiler
+ ext_module_builder: Extension module builder object used to generate the kernel extension module
+ objcache: Directory to cache the generated code files and compiled modules in.
+ If `None`, a temporary directory will be used, and compilation results will not be cached.
"""
@staticmethod
def create(
compiler_info: CompilerInfo | None = None,
objcache: str | Path | _AUTO_TYPE | None = AUTO,
- ):
+ ) -> CpuJit:
+ """Configure and create a CPU JIT compiler object.
+
+ Args:
+ compiler_info: Compiler info object defining capabilities and interface of the host compiler.
+ If `None`, a default compiler configuration will be determined from the current OS and runtime
+ environment.
+ objcache: Directory used for caching compilation results.
+ If set to `AUTO`, a persistent cache directory in the current user's home will be used.
+ If set to `None`, compilation results will not be cached--this may impact performance.
+
+ Returns:
+ The CPU just-in-time compiler.
+ """
+
if objcache is AUTO:
from appdirs import AppDirs
@@ -82,34 +103,19 @@ class CpuJit(JitBase):
+ self._compiler_info.linker_flags()
)
- @property
- def objcache(self) -> Path | None:
- return self._objcache
-
- @property
- def compiler_info(self) -> CompilerInfo:
- return self._compiler_info
-
- @compiler_info.setter
- def compiler_info(self, info: CompilerInfo):
- self._compiler_info = info
-
- @property
- def strict_scalar_types(self) -> bool:
- """Enable or disable implicit type casts for scalar parameters.
-
- If `True`, values for scalar kernel parameters must always be provided with the correct NumPy type.
+ def compile(self, kernel: Kernel) -> KernelWrapper:
+ """Compile the given kernel to an executable function.
+
+ Args:
+ kernel: The kernel object to be compiled.
+
+ Returns:
+ Wrapper object around the compiled function
"""
- return self._strict_scalar_types
- @strict_scalar_types.setter
- def strict_scalar_types(self, v: bool):
- self._strict_scalar_types = v
-
- def compile(self, kernel: Kernel) -> KernelWrapper:
# Get the Code
module_name = f"{kernel.function_name}_jit"
- cpp_code = self._ext_module_builder(kernel, module_name)
+ cpp_code = self._ext_module_builder.render_module(kernel, module_name)
# Get compiler information
import sysconfig
@@ -181,14 +187,21 @@ class CpuJit(JitBase):
class ExtensionModuleBuilderBase(ABC):
+ """Base class for CPU extension module builders."""
+
@staticmethod
@abstractmethod
- def include_dirs() -> list[str]: ...
+ def include_dirs() -> list[str]:
+ """List of directories that must be on the include path when compiling
+ generated extension modules."""
@abstractmethod
- def __call__(self, kernel: Kernel, module_name: str) -> str: ...
+ def render_module(self, kernel: Kernel, module_name: str) -> str:
+ """Produce the extension module code for the given kernel."""
@abstractmethod
def get_wrapper(
self, kernel: Kernel, extension_module: ModuleType
- ) -> KernelWrapper: ...
+ ) -> KernelWrapper:
+ """Produce the invocation wrapper for the given kernel
+ and its compiled extension module."""
diff --git a/src/pystencils/jit/cpu/cpujit_pybind11.py b/src/pystencils/jit/cpu/cpujit_pybind11.py
index 2f43ef1965278eb3148fbfb448a9cfb001542adf..1742bf30113f418e5a701db15f75f47bb2882a8d 100644
--- a/src/pystencils/jit/cpu/cpujit_pybind11.py
+++ b/src/pystencils/jit/cpu/cpujit_pybind11.py
@@ -43,7 +43,7 @@ class Pybind11KernelModuleBuilder(ExtensionModuleBuilderBase):
self._param_check_lines: list[str]
self._extraction_lines: list[str]
- def __call__(self, kernel: Kernel, module_name: str) -> str:
+ def render_module(self, kernel: Kernel, module_name: str) -> str:
self._actual_field_types = dict()
self._param_binds = []
self._public_params = []