triton.language.arange

triton.language.arange(start, end, _semantic=None)

Returns contiguous values within the half-open interval [start, end). end - start must be less than or equal to TRITON_MAX_TENSOR_NUMEL = 1048576

参数:

• start (int32) -- Start of the interval. Must be a power of two.

• end (int32) -- End of the interval. Must be a power of two greater than start.

Example

import torch
import triton
import triton.language as tl


@triton.jit
def arange_kernel(output_ptr, BLOCK: tl.constexpr, START: tl.constexpr, END: tl.constexpr):
    off = tl.arange(0, BLOCK)
    val = tl.arange(START, END)
    tl.store(output_ptr + off, val)


def test_arange():
    start = 0
    end = 128
    BLOCK = end - start
    output = torch.zeros(BLOCK, dtype=torch.int32).npu()
    arange_kernel[(1, )](output, BLOCK=BLOCK, START=start, END=end)
    expected = torch.arange(start, end, dtype=torch.int32)
    assert torch.allclose(output.cpu(), expected.cpu())


if __name__ == "__main__":
    test_arange()

Special Restrictions

• Arguments must be tl.constexpr of uint or int type (int64 not supported).

• start and end must be greater than or equal to 0.

• end - start < 1048576 (TRITON_MAX_TENSOR_NUMEL), applicable to both NV and Ascend.

• relaxed_requirement (Ascend extension): CUDA requires range=(end-start) to be a power of 2, Ascend has no such requirement.