triton.language.cumprod#

triton.language.cumprod(input, axis=0, reverse=False)#

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

Parameters:

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

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

Example

import torch
import triton
import triton.language as tl


@triton.jit
def cumprod_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 product 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.cumprod(x, axis=dim, reverse=reverse)
    tl.store(out_ptr0 + idx, ret)


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

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

    cumprod_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.cumprod(torch.flip(x.cpu(), [dim]), dim=dim), [dim])
    else:
        expected = torch.cumprod(x.cpu(), dim=dim)

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


if __name__ == "__main__":
    test_cumprod_2d()

DataType Support

平台	uint8	int8	uint16	int16	uint32	int32	uint64	int64	fp16	fp32	fp64	bf16	fp8e(e4m3)	fp8e5(e5m2)	bool
Ascend A2/A3	√	√	×	√	×	√	×	√	√	√	×	√	×	×	√
Ascend 950	√	√	√	√	√	√	√	√	√	√	×	√	×	×	√

Special Restrictions

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

triton.language.cumprod

Contents

triton.language.cumprod#