Custom Kernel Framework

Custom Kernel

XLACustomKernel(name[, doc])	Creates and manages a custom JAX primitive for XLA custom calls.
KernelEntry(backend, platform, kernel_generator)	A registered kernel implementation for a specific backend and platform.

CPU Kernel via Numba

numba_kernel(kernel, outs, *[, vmap_method, ...])

Create a JAX-callable function from a Numba CPU kernel.

GPU Kernel via Numba CUDA

numba_cuda_kernel(kernel, outs, *[, grid, ...])	Create a JAX-callable function from a single Numba CUDA kernel.
numba_cuda_callable(func, outs, *[, ...])	Create a JAX-callable from a Python function that launches Numba CUDA kernels.

GPU Kernel via CUDA Source

load_cuda_inline(name, cuda_sources[, ...])	Compile inline CUDA source and load the resulting module.
load_cuda_file(filepath[, functions, name])	Compile a single .cu file and load the resulting module.
load_cuda_dir(directory[, functions, name, ...])	Compile all CUDA files in a directory and load the resulting module.

CPU Kernel via C++ Source

load_cpp_inline(name, cpp_sources[, ...])	Compile inline C++ source for CPU (or CUDA) and load the module.
load_cpp_file(filepath[, functions, name])	Compile a single .cpp / .cc file for CPU and load the module.

Runtime

CompiledModule(so_path, function_names)

A compiled module loaded from a shared library.

register_ffi_target(target_name, module, ...)	Register a compiled function as a JAX FFI target.
list_registered_targets()	Return a sorted list of all registered FFI target names.

Cache Utilities

clear_cache([name])	Clear the compilation cache.
set_cache_dir(path)	Set the compilation cache directory.
get_cache_dir()	Return the current compilation cache directory path.

Diagnostics

print_diagnostics()

Print a summary of the brainevent compilation environment.

Exceptions

BrainEventError
CompilationError(message[, compiler_output, ...])	CUDA or C++ compilation failed.
KernelToolchainError	Compilation toolchain missing or incompatible.
KernelRegistrationError	JAX FFI target registration failed.