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_idx

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

Copies a value from a source lane to other lanes in a warp.

Broadcasts a value from a source thread in a warp to all participating threads
without using shared memory. This is a convenience wrapper that uses the full
warp mask by default.

Example:

    from max.gpu.primitives.warp import shuffle_idx

    val = SIMD[.float32, 16](1.0)

    # Broadcast value from lane 0 to all lanes
    result = shuffle_idx(val, 0)

    # Get value from lane 5
    result = shuffle_idx(val, 5)

Parameters:

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

Args:

val (SIMD[dtype, simd_width]): The SIMD value to be broadcast from the source lane.
offset (UInt32): The source lane ID to copy the value from.

Returns:

SIMD[dtype, simd_width]: A SIMD vector where all lanes contain the value from the source lane specified by offset.

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

Copies a value from a source lane to other lanes in a warp with explicit mask control.

Broadcasts a value from a source thread in a warp to participating threads specified by
the mask. This provides fine-grained control over which threads participate in the shuffle
operation.

Example:

    from max.gpu.primitives.warp import shuffle_idx

    # Only broadcast to first 16 lanes
    var mask: UInt = 0xFFFF  # 16 ones
    var val = SIMD[.float32, 32](1.0)
    var result = shuffle_idx(mask, val, 5)

Parameters:

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

Args:

mask (UInt): A bit mask specifying which lanes participate in the shuffle (1 bit per lane).
val (SIMD[dtype, simd_width]): The SIMD value to be broadcast from the source lane.
offset (UInt32): The source lane ID to copy the value from.

Returns:

SIMD[dtype, simd_width]: A SIMD vector where participating lanes (set in mask) contain the value from the source lane specified by offset. Non-participating lanes retain their original values.