triton.language.range

class triton.language.range(arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, disallow_acc_multi_buffer=False, flatten=False, warp_specialize=False, disable_licm=False)

Iterator that counts upward forever.

@triton.jit
def kernel(...):
    for i in tl.range(10, num_stages=3):
        ...

Note:

This is a special iterator used to implement similar semantics to Python's range in the context of triton.jit functions. In addition, it allows user to pass extra attributes to the compiler.

参数:

• arg1 -- the start value.

• arg2 -- the end value.

• step -- the step value.

• num_stages --

  pipeline the loop into this many stages (so there are num_stages iterations of the loop in flight at once).

  Note this is subtly different than passing num_stages as a kernel argument. The kernel argument only pipelines loads that feed into dot operations, while this attribute tries to pipeline most (though not all) loads in this loop.

• loop_unroll_factor -- Tells the Triton IR level loop unroller how many times to unroll a for loop that this range is used with. Less than 2 for this value implies no unrolling.

• disallow_acc_multi_buffer -- If true, prevent the accumulator of the dot operation in the loop to be multi-buffered, if applicable.

• flatten -- automatically flatten the loop nest starting at this loop to create a single flattened loop. The compiler will try to pipeline the flattened loop which can avoid stage stalling.

• warp_specialize -- Enable automatic warp specialization on the loop. The compiler will attempt to partition memory, MMA, and vector operations in the loop into separate async partitions. This will increase the total number of warps required by the kernel.

• disable_licm --

  Tells the compiler it shouldn't hoist loop invariant code outside the loop. This is often useful to avoid creating long liveranges within a loop.

  Note that warp specialization is only supported on Blackwell GPUs and only works on simple matmul loops. Support for arbitrary loops will be expanded over time.

Example

import torch
import torch_npu
from torch.testing import assert_close

import triton
import triton.language as tl


@triton.jit
def range_kernel(x_ptr, out_ptr, N: tl.constexpr):
    acc = 0.0
    # Iterate from 0 to N (exclusive), accumulating each element
    for i in tl.range(0, N):
        acc += tl.load(x_ptr + i)
    tl.store(out_ptr, acc)


def test_range():
    N = 128
    x = torch.randn(N, device="npu", dtype=torch.float32)
    out = torch.empty(1, device="npu", dtype=torch.float32)

    range_kernel[(1, )](x, out, N=N)
    torch.npu.synchronize()

    assert_close(out, x.sum().reshape(1), rtol=1e-3, atol=1e-3)


if __name__ == "__main__":
    test_range()
    print("test_range PASSED!")

Special Restrictions

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

• disallow_acc_multi_buffer, flatten, disable_licm: related functionality is incomplete on Ascend.

• warp_specialize: only supported on Blackwell GPU; has no effect on Ascend.

__init__(arg1, arg2=None, step=None, num_stages=None, loop_unroll_factor=None, disallow_acc_multi_buffer=False, flatten=False, warp_specialize=False, disable_licm=False)

Methods

__init__(arg1[, arg2, step, num_stages, ...])

Attributes

type