View page source

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 APIC graph capture (ScopedCapture with apic=True), int32, float32, int64, and float64 scalar and vector scans are recorded into the byte stream, including positively-strided (non-contiguous) 1D arrays. Negatively-strided CPU 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 inside a CPU APIC ScopedCapture with a negatively-strided array.

Return type:

None