warp.utils.array_scan#

warp.utils.array_scan(in_array, out_array, inclusive=True)[source]#

Perform a scan (prefix sum) operation on an array.

This function computes the inclusive or exclusive scan of the input array and stores the result in the output array. The scan operation computes a running sum of elements in the array. Vector types are scanned component-wise.

During CPU graph capture, or CUDA graph capture with apic=True, int32, float32, int64, and float64 scalar and vector scans are recorded into the API Capture (APIC) operation stream, including positively-strided (non-contiguous) 1D arrays. Non-empty, negatively-strided arrays raise NotImplementedError; run those scans outside the captured region.

Parameters:

  • in_array (array) – Input array to scan. Must be a scalar or vector type with scalar type int32, int64, float32, or float64.

  • out_array (array) – Output array to store scan results. Must match input array type and size.

  • inclusive (bool) – If True, performs an inclusive scan (includes current element in sum). If False, performs an exclusive scan (excludes current element).

Raises:

  • RuntimeError – If array storage devices don’t match, if storage size is insufficient, or if data types are unsupported.

  • NotImplementedError – If called with a non-empty, negatively-strided array during CPU graph capture or CUDA graph capture with apic=True.

Return type:

None

On this page