triton.language.cumsum

triton.language.cumsum = <function cumsum>

Returns the cumsum of all elements in the input tensor along the provided axis

参数:

• input (Tensor) -- the input values

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

• reverse (bool) -- if true, the scan is performed in the reverse direction

• dtype (tl.dtype) -- the desired data type of the returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. If not specified, small integer types (< 32 bits) are upcasted to prevent overflow. Note that tl.bfloat16 inputs are automatically promoted to tl.float32.

This function can also be called as a member function on tensor, as x.cumsum(...) instead of cumsum(x, ...). .. rubric:: Example

import torch
import triton
import triton.language as tl


@triton.jit
def cumsum_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,
):
    """
    Cumulatively sum the input tensor (shape numel_x × numel_r) along the specified dimension,
    writing the result to out_ptr0.
    """
    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.cumsum(x, axis=dim, reverse=reverse)
    tl.store(out_ptr0 + idx, ret)


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

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

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

    if reverse:
        expected = torch.flip(torch.cumsum(torch.flip(x.cpu(), [dim]), dim=dim), [dim])
    else:
        expected = torch.cumsum(x.cpu(), dim=dim)

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


if __name__ == "__main__":
    test_cumsum_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