View page source

warp.kernel#

warp.kernel(
f=None,
*,
enable_backward=None,
launch_bounds=None,
cluster_dim=None,
module=None,
module_options=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(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.

  • 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.

  • 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.

  • 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.