triton.language.associative_scan

triton.language.associative_scan(input, axis, combine_fn, reverse=False, _semantic=None, _generator=None)

Applies the combine_fn to each elements with a carry in input tensors along the provided axis and update the carry

参数:

• input (Tensor) -- the input tensor, or tuple of tensors

• axis (int) -- the dimension along which the reduction should be done

• combine_fn (Callable) -- a function to combine two groups of scalar tensors (must be marked with @triton.jit)

• reverse -- whether to apply the associative scan in the reverse direction along axis

Example

import torch
import triton
import triton.language as tl


def combine_fn_test(a, b):
    return a + b


@triton.jit
def associative_scan_2d_kernel(
    out_ptr0,
    in_ptr0,
    dim: tl.constexpr,
    reverse: tl.constexpr,
    numel_x: tl.constexpr,
    numel_r: tl.constexpr,
    XBLOCK: tl.constexpr,
    RBLOCK: tl.constexpr,
):
    tl.static_assert(numel_x == XBLOCK, "numel_x must be equal to XBLOCK in this kernel")
    tl.static_assert(numel_r == RBLOCK, "numel_r must be equal to RBLOCK in this kernel")
    idx_x = tl.arange(0, XBLOCK)
    idx_r = tl.arange(0, RBLOCK)
    idx = idx_x[:, None] * numel_r + idx_r[None, :]
    x = tl.load(in_ptr0 + idx)
    ret = tl.associative_scan(x, axis=dim, reverse=reverse, combine_fn=combine_fn_test)
    tl.store(out_ptr0 + idx, ret)


def test_associative_scan_2d():
    X = 4
    R = 8
    dim = 0
    reverse = False

    x = torch.randn((X, R), dtype=torch.float32).npu()
    output = torch.zeros_like(x)

    associative_scan_2d_kernel[(1, )](output, x, dim=dim, reverse=reverse, numel_x=X, numel_r=R, XBLOCK=X, RBLOCK=R)

    expected = torch.cumsum(x.cpu(), dim=dim)

    assert torch.allclose(output.cpu(), expected.cpu())


if __name__ == "__main__":
    test_associative_scan_2d()

Special Restrictions

• DataType: Ascend A2/A3 does not support uint16, uint32, uint64 (hardware limitation).

• reverse=True requires alignment when loading data with tl.load, and mask cannot be used to filter redundant data indices