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. cuda_max_registers must be a positive int and cannot be combined with launch_bounds. The cuda_max_registers option applies only 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. Warp applies enable_cuda_smem_spilling independently to the forward and backward kernels and silently ignores the option 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