Skip to main content
Log in

Mojo function

lane_group_sum

lane_group_sum[val_type: DType, simd_width: Int, //, num_lanes: Int, stride: Int = 1](val: SIMD[val_type, simd_width]) -> SIMD[val_type, simd_width]

Computes the sum of values across a group of lanes using warp-level operations.

This function performs a parallel reduction across a group of lanes to compute their sum. The reduction is done using warp shuffle operations for efficient communication between lanes. The result is stored in all participating lanes.

Parameters:

  • val_type (DType): The data type of the SIMD elements (e.g. float32, int32).
  • simd_width (Int): The number of elements in the SIMD vector.
  • num_lanes (Int): The number of threads participating in the reduction.
  • stride (Int): The stride between lanes participating in the reduction.

Args:

  • val (SIMD[val_type, simd_width]): The SIMD value to reduce. Each lane contributes its value to the sum.

Returns:

A SIMD value where all participating lanes contain the sum found across the lane group. Non-participating lanes (lane_id >= num_lanes) retain their original values.