For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
prefix_sum
def prefix_sum[dtype: DType, //, *, block_size: Int, exclusive: Bool = False](val: Scalar[dtype]) -> Scalar[dtype]
Performs a prefix sum (scan) operation across all threads in a 1D block.
This function implements a block-level inclusive or exclusive scan, efficiently computing the cumulative sum for each thread based on thread indices.
Parameters:
- dtype (
DType): The data type of the Scalar elements. - block_size (
Int): The total number of threads in the block. - exclusive (
Bool): If True, perform exclusive scan instead of inclusive.
Args:
- val (
Scalar[dtype]): The Scalar value from each thread to include in the scan.
Returns:
Scalar[dtype]: A Scalar value containing the result of the scan operation for each
thread.
def prefix_sum[dtype: DType, //, *, block_dim_x: Int, block_dim_y: Int, block_dim_z: Int = Int(1), exclusive: Bool = False](val: Scalar[dtype]) -> Scalar[dtype]
Performs a prefix sum (scan) operation across all threads in a multi-dimensional block.
This function implements a block-level inclusive or exclusive scan for 2D
and 3D thread blocks. Thread IDs are linearized in row-major order:
x + y * dim_x + z * dim_x * dim_y.
Parameters:
- dtype (
DType): The data type of the Scalar elements. - block_dim_x (
Int): The number of threads along the X dimension. - block_dim_y (
Int): The number of threads along the Y dimension. - block_dim_z (
Int): The number of threads along the Z dimension (default: 1). - exclusive (
Bool): If True, perform exclusive scan instead of inclusive.
Args:
- val (
Scalar[dtype]): The Scalar value from each thread to include in the scan.
Returns:
Scalar[dtype]: A Scalar value containing the result of the scan operation for each
thread.