triton.language.extra.cann.extension.gather_out_to_ub

triton.language.extra.cann.extension.gather_out_to_ub(src: tensor, index: tensor, index_boundary: int, dim: int, src_stride: tuple, end_offset: tuple, start_offset: tuple, other=None, _semantic=None)

Gather from a source tensor in Global Memory (GM) to Unified Buffer (UB) along a specified dimension with out-of-bound handling.

Gather operation for different tensor ranks:

1D index gather:
out[i] = src[start_offset[0] + index[i]]
2D index gather (0 <= dim < 2):

2.1 dim = 0
out[i][j] = src[start_offset[0] + index[i][j]][start_offset[1] + j]

2.2 dim = 1
out[i][j] = src[start_offset[0] + i][start_offset[1] + index[i][j]]
3D index gather (0 <= dim < 3):

3.1 dim = 0
out[i][j][k] = src[start_offset[0] + index[i][j][k]][start_offset[1] + j][start_offset[2] + k]

3.2 dim = 1
out[i][j][k] = src[start_offset[0] + i][start_offset[1] + index[i][j][k]][start_offset[2] + k]

3.3 dim = 2
out[i][j][k] = src[start_offset[0] + i][start_offset[1] + j][start_offset[2] + index[i][j][k]]

参数:

src (pointer type) -- The source tensor pointer (in GM).
index (tensor) -- A tensor to gather (in UB).
index_boundary (int) -- The upper boundary for index values.
dim (int) -- The dimension to gather along.
src_stride (tuple of int) -- The stride of each dimension of src tensor.
end_offset (tuple of int) -- The end offsets of each dimension for index tensor.
start_offset (tuple of int) -- The start offsets of each dimension for index tensor.
other (scalar, optional) -- The default value when index is out of boundary (in UB).

返回:

Tensor with the same shape as index.shape (in UB).

返回类型:

tensor

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

Example

import torch
import triton
import triton.language as tl
from triton.language.extra.cann.extension import gather_out_to_ub
from triton.tools.get_ascend_devices import is_compile_on_910_95


@triton.jit
def kernel(src_ptr, index_ptr, out_ptr):
    """
    Test gather_out_to_ub with a 2D index tile:
    1. Load a [2,2] index tile from GM into UB
    2. Gather values from src (in GM) along dim=0 using the index tile
    3. Store the gathered result (in UB) to output
    """
    # index tile shape: [2,2]
    y0_local = tl.arange(0, 2)[:, None]  # [0,1] rows
    x1_local = tl.arange(0, 2)[None, :]  # [0,1] cols
    mask = (y0_local < 2) & (x1_local < 2)

    # Load index tile to UB
    index = tl.load(index_ptr + y0_local * 2 + x1_local, mask)

    # Call gather_out_to_ub: gather values from src along dim=0
    gathered = gather_out_to_ub(src=src_ptr, index=index, index_boundary=4, dim=0, src_stride=(2, 1), end_offset=(2, 2),
                                start_offset=(0, 0))

    tl.store(out_ptr + y0_local * 2 + x1_local, gathered, mask)


def test_gather_out_to_ub():
    # src(4,2) = [[1.,2.],[3.,4.],[5.,6.],[7.,8.]]
    src = torch.tensor([[1., 2.], [3., 4.], [5., 6.], [7., 8.]], device='npu')
    # index(2,2) = [[0,1],[2,3]]
    index = torch.tensor([[0, 1], [2, 3]], device='npu')
    out = torch.empty((2, 2), device='npu', dtype=torch.float32)

    kernel[(1, )](src, index, out)

    # Reference: gather along dim=0 with src(4,2) and index(2,2)
    # out[i,j] = src[index[i,j], j]
    # out[0,0] = src[index[0,0]=0, 0] = 1.0
    # out[0,1] = src[index[0,1]=1, 1] = 4.0
    # out[1,0] = src[index[1,0]=2, 0] = 5.0
    # out[1,1] = src[index[1,1]=3, 1] = 8.0
    expected = torch.tensor([[1., 4.], [5., 8.]], device='cpu')
    assert torch.allclose(out.cpu(), expected), \
        f"Mismatch: expected {expected}, got {out.cpu()}"


if __name__ == "__main__":
    if not is_compile_on_910_95:
        print("gather_out_to_ub is only supported on Ascend 950, skipping test.")
    else:
        test_gather_out_to_ub()

Special Restrictions

Only supported on Ascend 950.
src.dtype: only supports float16, bfloat16, float32.
src and index must have the same rank.
index: must be an integer tensor, with rank in [1, 5].
dim: must satisfy 0 <= dim < rank(index).
other: must be a scalar value.
For every dimension i not equal to dim, index.size[i] <= src.size[i].
The output shape is the same as index.shape.