JIT Compilation

JIT Compilation#

Base Infrastructure#

`KernelWrapper`	Wrapper around a compiled and executable pystencils kernel.
`JitBase`	Base class for just-in-time compilation interfaces implemented in pystencils.
`NoJit`	Not a JIT compiler: Used to explicitly disable JIT compilation on an AST.

pystencils.jit.no_jit = <pystencils.jit.jit.NoJit object>#: Disables just-in-time compilation for a kernel.

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.

LegacyCpuJit

Wrapper around pystencils.cpu.cpujit

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.

To configure and create an instance of the CPU JIT compiler, use the CpuJit.create factory method:

pystencils.jit.CpuJit.create(compiler_info=None, objcache=AUTO)

Configure and create a CPU JIT compiler object.

Parameters:

compiler_info (Optional[CompilerInfo]) – 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 (str | Path | _AUTO_TYPE | None) – 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.

Return type:

CpuJit

Returns:

The CPU just-in-time compiler.

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.

`CompilerInfo`	Base class for compiler infos.
`GccInfo`	Compiler info for the GNU Compiler Collection C++ compiler (`g++`).
`ClangInfo`	Compiler info for the LLVM C++ compiler (`clang`).

Implementation#

`CpuJit`	Just-in-time compiler for CPU kernels.
`cpujit.ExtensionModuleBuilderBase`	Base class for CPU extension module builders.

CuPy-based GPU JIT#

`CupyJit`
`CupyKernelWrapper`
`LaunchGrid`