warp.kernel#

warp.kernel(
f=None,
*,
name=None,
enable_backward=None,
launch_bounds=None,
cuda_max_registers=None,
enable_cuda_smem_spilling=None,
cluster_dim=None,
module=None,
module_options=None,
entry_point_abi=None,
grid_stride=None,
)[source]#

Decorator to register a Warp kernel from a Python function. The function must be defined with type annotations for all arguments. The function must not return anything.

Example:

@wp.kernel
def my_kernel(a: wp.array[float], b: wp.array[float]):
    tid = wp.tid()
    b[tid] = a[tid] + 1.0


@wp.kernel(enable_backward=False)
def my_kernel_no_backward(a: wp.array2d[float], x: float):
    # the backward pass will not be generated
    i, j = wp.tid()
    a[i, j] = x


@wp.kernel(module="unique")
def my_kernel_unique_module(a: wp.array[float], b: wp.array[float]):
    # the kernel will be registered in new unique module created just for this
    # kernel and its dependent functions and structs
    tid = wp.tid()
    b[tid] = a[tid] + 1.0


@wp.kernel(launch_bounds=(256, 1))
def my_kernel_with_launch_bounds(a: wp.array[float]):
    # CUDA __launch_bounds__ will be set to (256, 1)
    tid = wp.tid()
    a[tid] = a[tid] * 2.0


@wp.kernel(cuda_max_registers=64)
def my_kernel_with_cuda_max_registers(a: wp.array[float]):
    # CUDA __maxnreg__(64) will be set when supported
    tid = wp.tid()
    a[tid] = a[tid] * 2.0


@wp.kernel(enable_cuda_smem_spilling=True, launch_bounds=256)
def my_kernel_with_cuda_smem_spilling(a: wp.array[float]):
    # CUDA 13+ may use shared memory for register spills
    tid = wp.tid()
    a[tid] = a[tid] * 2.0


@wp.kernel(module_options={"fast_math": True}, module="unique")
def my_kernel_fast(a: wp.array[float], b: wp.array[float]):
    # fast_math is a module-level option, so module="unique" is required
    tid = wp.tid()
    b[tid] = a[tid] + 1.0
Parameters:

  • f (Callable | None) – The function to be registered as a kernel.

  • name (str | None) – Sets the kernel key used for registration and native code generation. If None, Warp derives the key from f. A custom name must be a valid C++ identifier. When strip_hash=True, Warp uses the key without a hash suffix as the base of the generated native entry-point names.

  • enable_backward (bool | None) – If False, the backward pass will not be generated.

  • launch_bounds (tuple[int, ...] | int | None) – CUDA __launch_bounds__ attribute for the kernel. Can be an int (maxThreadsPerBlock) or a tuple of 1-2 ints (maxThreadsPerBlock, minBlocksPerMultiprocessor). Only applies to CUDA kernels. Note: The block_dim parameter in warp.launch() must not exceed the maxThreadsPerBlock value specified here.

  • cuda_max_registers (int | None) – CUDA __maxnreg__ attribute specifying the maximum number of registers allocated per thread. Must be a positive int and cannot be combined with launch_bounds. Only applies to CUDA kernels and is ignored when Warp was built with CUDA Toolkit earlier than 12.4 or when warp.config.llvm_cuda is True.

  • enable_cuda_smem_spilling (bool | None) – If True, allow the CUDA Toolkit used to build Warp, when version 13.0 or later, to use shared memory for register spills. Applied independently to the forward and backward kernels and silently ignored when an entry point uses dynamic shared memory, on CPU, with older CUDA Toolkits, in unsupported device-debug compilation, or when warp.config.llvm_cuda is True. Explicit launch_bounds are recommended to avoid over-allocating shared memory and reducing occupancy.

  • cluster_dim (int | None) – CUDA Thread Block Cluster size as a 1D CTA count. Warp emits CUDA __cluster_dims__(cluster_dim, 1, 1) because kernels use a 1D hardware launch grid. Must be a positive int <= 16 (Hopper non-portable cap). Default 1 means no clustering. Only effective on devices with compute capability >= 9.0; silently ignored on older archs and on CPU. See warp.get_cuda_max_cluster_dim().

  • module (Module | Literal['unique'] | str | None) – The warp._src.context.Module to which the kernel belongs. Alternatively, if a string "unique" is provided, the kernel is assigned to a new module named after the kernel name and hash. If None, the module is inferred from the function’s module.

  • module_options (dict[str, Any] | None) – A dict of module-level compilation options (e.g. fast_math, mode, max_unroll, deterministic, deterministic_max_records) that are applied to the kernel’s module. Requires module="unique"; raises ValueError otherwise. For shared modules, use warp.set_module_options() instead. See warp.set_module_options() for the full list of supported options.

  • entry_point_abi (Literal['warp', 'external_constant_params'] | None) – Experimental entry-point ABI. "warp" (the default) is supported by CPU and CUDA and emits the regular warp.launch()-compatible kernel signature. "external_constant_params" is supported by CUDA only; it emits a no-argument external entry point and binds the kernel’s single Warp struct argument to the constant-memory symbol params. It cannot be launched with warp.launch() and requires enable_backward=False. This API may change in future releases.

  • grid_stride (bool | None) – Whether to emit a grid-stride loop. False opts into a lean launch (no grid-stride loop) with lower per-thread overhead and register pressure, but the block count cannot be capped: launching it with max_blocks > 0 raises. None defers to the "default_grid_stride" module option, then warp.config.default_grid_stride.

Returns:

The registered kernel.

On this page