JIT Compiler

The JIT (Just-In-Time) module is the engine of HeteroSymNN. It handles the translation of symbolic math expressions into optimized machine code.

Symbolic Compiler

class HeteroSymNN.JIT.compiler.SymbolicJITCompiler(configs: list[NodeConfig], calculation_method: Literal['GPU_CUDA', 'CPU_JIT', 'CPU_PYTHON'], device_id: int, mode: Literal['activation', 'loss'] = 'activation')[source]

Bases: object

The SymbolicJITCompiler is the computational heart of HeteroSymNN. It is responsible for transforming high-level symbolic definitions of mathematical functions (activations and losses) into highly optimized, hardware-specific executable kernels at runtime.

This compiler bridges the gap between flexibility and performance by leveraging SymPy for symbolic differentiation and code generation, and then compiling that code into:

  • CUDA Kernels (GPU_CUDA): For massive parallelism on NVIDIA GPUs using CuPy.

  • C++ Shared Libraries (CPU_JIT): For high-performance CPU execution using OpenMP and system compilers (MSVC/GCC).

  • Python Lambdas (CPU_PYTHON): As a fallback for maximum compatibility.

It handles the automatic differentiation of user-defined formulas, manages the compilation cache to avoid redundant work, and provides a unified interface (forward_kernel, backward_kernel) for the rest of the library to execute these functions without worrying about the underlying hardware implementation.

Parameters:

  • configs (list[NodeConfig]) – A list of configurations defining the functions to compile. For ‘activation’ mode, this is a list of (function_name_or_expression, constants_dict) for each node. For ‘loss’ mode, this is a list containing a single tuple with the loss expression and its constants.

  • calculation_method (Literal["GPU_CUDA", "CPU_JIT", "CPU_PYTHON"]) – The target backend for compilation.

  • device_id (int) – The ID of the GPU device to use if compiling for CUDA.

  • mode (Literal["activation", "loss"], optional) – The type of function being compiled. Determines the kernel signature and symbolic variables used (‘num’ for activations, ‘y_pred’/’y_true’ for losses). Defaults to “activation”.

forward_kernel

The compiled executable function for the forward pass.

Type:

Callable

backward_kernel

The compiled executable function for the backward pass (gradient calculation).

Type:

Callable

calculation_method

The current active calculation method.

Type:

str

device_id

The current GPU ID.

Type:

int

Examples

Although this class is primarily used internally, it can be instantiated for testing custom symbolic expressions.

>>> import numpy as np
>>> from HeteroSymNN.JIT.compiler import SymbolicJITCompiler
>>>
>>> configs = [("Max(0, num)", {})]
>>> method = "CPU_PYTHON"
>>> mode = "activation"
>>>
>>> py_compiler = SymbolicJITCompiler(
...     configs=configs,
...     calculation_method=method,
...     mode=mode
... )
>>>
>>> configs_2 =[("d*y_pred-y_true", {"d": 2.0})]
>>> method_2 = "GPU_CUDA"
>>> mode_2 = "loss"
>>> device_id = 0
>>>
>>> cuda_compiler = SymbolicJITCompiler(
...     configs=configs_2,
...     calculation_method=method_2,
...     device_id=device_id,
...     mode=mode_2
... )

Capabilities:

  • Parses SymPy expressions from string configs.

  • Differentiates expressions symbolically for backpropagation.

  • Compiles to C++ (OpenMP) or CUDA kernels depending on the hardware.

set_gpu_id(new_id: int)[source]

Updates the active GPU ID and recompiles CUDA kernels if necessary.

Parameters:

new_id (int) – The new GPU device ID.