triton.language.extra.cann.extension.index_select_simd

triton.language.extra.cann.extension.index_select_simd(src, dim, index, src_shape, src_offset, read_shape, _semantic=None) → tensor

Parallel index_select operation from Global Memory to Unified Buffer (SIMD version).

Selects data from multiple indices along a specified dimension and loads them as tiles from GM directly to UB with zero-copy semantics.

参数:

src (tensor (pointer type)) -- Source tensor pointer (in GM)
dim (int or constexpr) -- The dimension along which to select indices
index (tensor) -- 1D tensor of indices to select (in UB)
src_shape (List[Union[int, tensor]]) -- Complete shape of the source tensor (can be int or tensor)
src_offset (List[Union[int, tensor]]) -- Starting offset for reading (can be int or tensor)
read_shape (List[Union[int, tensor]]) -- Size to read (tile shape, can be int or tensor)

Constraints:

read_shape[dim] must be -1
src_offset[dim] can be -1 (will be ignored)
Boundary handling: src_offset + read_shape > src_shape automatically truncates to src_shape boundary
Does not check if index contains out-of-bounds values

Example:

@triton.jit
def kernel(src_ptr, output_ptr, indices_ptr, M, N, D, ...):
    # Load indices (e.g., [5, 10, 15, 20])
    indices = tl.load(indices_ptr + tl.arange(0, 4))

    # Example 1: Static shapes (constants)
    # Index select from dimension 1
    # src: [8, 100, 256], index_select at dim=1
    # Read: [4, ?, 128] starting from [4, ?, 128]
    result = extension.index_select_simd(
        src_ptr,
        dim=1,
        index=indices,
        src_shape=[8, 100, 256],
        src_offset=[4, -1, 128],
        read_shape=[4, -1, 128]
    )
    # result shape: [4, 4, 128]

    # Example 2: Dynamic shapes (variables)
    result2 = extension.index_select_simd(
        src_ptr,
        dim=1,
        index=indices,
        src_shape=[M, N, D],
        src_offset=[4, -1, 128],
        read_shape=[4, -1, 128]
    )

    tl.store(output_ptr + ..., result)

返回:

Result tensor in UB with shape where dim is replaced by the length of index

返回类型:

tensor

This function can also be called as a member function on tensor, as x.index_select_simd(...) instead of index_select_simd(x, ...).

Special Restrictions

DataType: Ascend A2/A3 does not support fp64, fp8e4, fp8e5, uint16, uint32, uint64 (hardware limitation).
index: The data type of the index must be int32 or int64.
dim: The dimension cannot be the trailing axis (the last dimension), i.e., dim < len(src_shape) - 1.
The index_select operation is not supported along the trailing axis (the last dimension).
Out-of-bounds indices are not checked; users must ensure index validity on their own.
The index must be a 1D tensor.