warp.func_native#

warp.func_native(snippet, adj_snippet=None, replay_snippet=None)[source]#

@warp.func_native is an escape hatch for operations that are easier to express in native code than in Warp’s kernel language, such as CUDA intrinsics, shared-memory synchronization patterns, or small C++ helper expressions. The decorated Python function is a typed stub: its argument names become the variable names available inside the snippet, and its body should be ....

Parameters:

snippet (str) – Native C++/CUDA code inserted into the generated forward function body.
adj_snippet (str | None) – Optional native code inserted into the generated adjoint function body. Use adj_-prefixed argument names, and adj_ret for the adjoint of a return value.
replay_snippet (str | None) – Optional native code used when Warp replays the forward function while executing the generated backward pass. Use this when replaying snippet would repeat an unsafe side effect, such as incrementing an atomic counter.

Returns:

A decorator that registers the typed stub as a Warp function.

Note

Compute thread indices in the calling kernel or function and pass them explicitly; snippets cannot call wp.tid(). Pure C++ snippets can be used by CPU kernels, while CUDA-specific constructs require CUDA kernels. If the snippet returns a value, the Python stub must declare the return type. Struct return values are not supported.

Example

Insert a native element-wise operation:

snippet = "out[tid] = x[tid] + 1.0f;"


@warp.func_native(snippet)
def increment(
    x: warp.array[warp.float32],
    out: warp.array[warp.float32],
    tid: int,
): ...


@warp.kernel
def kernel(x: warp.array[warp.float32], out: warp.array[warp.float32]):
    tid = warp.tid()
    increment(x, out, tid)

Use CUDA shared memory:

This pattern assumes the kernel is launched with block_dim=128.

snippet = '''
    int local_tid = threadIdx.x;
    __shared__ int values[128];
    values[local_tid] = arr[tid];
    __syncthreads();
    out[tid] = values[127 - local_tid];
    '''


@warp.func_native(snippet)
def reverse_block(arr: warp.array[int], out: warp.array[int], tid: int): ...

Provide an adjoint snippet for differentiability:

snippet = "out[tid] = 2.0f * x[tid] + y[tid];"
adj_snippet = '''
    adj_x[tid] += 2.0f * adj_out[tid];
    adj_y[tid] += adj_out[tid];
    '''


@warp.func_native(snippet=snippet, adj_snippet=adj_snippet)
def axpy(
    x: warp.array[warp.float32],
    y: warp.array[warp.float32],
    out: warp.array[warp.float32],
    tid: int,
): ...