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

shuffle_down

def shuffle_down[dtype: DType, simd_width: SIMDLength, //](val: SIMD[dtype, simd_width], offset: UInt32) -> SIMD[dtype, simd_width]

Copies values from threads with higher lane IDs in the warp.

Performs a shuffle operation where each thread receives a value from a thread with a higher lane ID, offset by the specified amount. Uses the full warp mask by default.

For example, with offset=1:

• Thread 0 gets value from thread 1
• Thread 1 gets value from thread 2
• Thread N gets value from thread N+1
• Last N threads get undefined values

Parameters:

• ​dtype (DType): The data type of the SIMD elements (e.g. float32, int32).
• ​simd_width (SIMDLength): The number of elements in each SIMD vector.

Args:

• ​val (SIMD[dtype, simd_width]): The SIMD value to be shuffled down the warp.
• ​offset (UInt32): The number of lanes to shift values down by. Must be positive.

Returns:

SIMD[dtype, simd_width]: The SIMD value from the thread offset lanes higher in the warp. Returns undefined values for threads where lane_id + offset >= WARP_SIZE.

def shuffle_down[dtype: DType, simd_width: SIMDLength, //](mask: UInt, val: SIMD[dtype, simd_width], offset: UInt32) -> SIMD[dtype, simd_width]

Copies values from threads with higher lane IDs in the warp using a custom mask.

Performs a shuffle operation where each thread receives a value from a thread with a higher lane ID, offset by the specified amount. The mask parameter controls which threads participate in the shuffle.

For example, with offset=1:

• Thread 0 gets value from thread 1
• Thread 1 gets value from thread 2
• Thread N gets value from thread N+1
• Last N threads get undefined values

Parameters:

• ​dtype (DType): The data type of the SIMD elements (e.g. float32, int32).
• ​simd_width (SIMDLength): The number of elements in each SIMD vector.

Args:

• ​mask (UInt): A bitmask controlling which threads participate in the shuffle. Only threads with their corresponding bit set will exchange values.
• ​val (SIMD[dtype, simd_width]): The SIMD value to be shuffled down the warp.
• ​offset (UInt32): The number of lanes to shift values down by. Must be positive.

Returns:

SIMD[dtype, simd_width]: The SIMD value from the thread offset lanes higher in the warp. Returns undefined values for threads where lane_id + offset >= WARP_SIZE or where the corresponding mask bit is not set.