add new CPU JIT classes to API reference pages

docs/source/api/jit.md

+# 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
+```

docs/source/api/jit.rst

-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

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()

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++"]

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.
+    **Creation**
     
-    :Implementation Details:
+    To configure and create a CPU JIT compiler instance, use the `create` factory method.
 
-    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.
+    **Implementation Details**
 
-    Both can be dynamically exchanged.
+    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
+    def compile(self, kernel: Kernel) -> KernelWrapper:
+        """Compile the given kernel to an executable function.
         
-    @property
-    def strict_scalar_types(self) -> bool:
-        """Enable or disable implicit type casts for scalar parameters.
+        Args:
+            kernel: The kernel object to be compiled.
         
-        If `True`, values for scalar kernel parameters must always be provided with the correct NumPy type.
+        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."""

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 = []