triton.language.static_range

class triton.language.static_range(arg1, arg2=None, step=None)

Iterator that counts upward forever.

@triton.jit
def kernel(...):
    for i in tl.static_range(10):
        ...

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 also guides the compiler to unroll the loop aggressively.

参数:

• arg1 -- the start value.

• arg2 -- the end value.

• step -- the step value.

Example

import pytest
import torch
import torch_npu
from torch.testing import assert_close

import triton
import triton.language as tl


@triton.jit
def static_range_kernel(x_ptr, out_ptr, BLOCK_SIZE: tl.constexpr):
    # static_range unrolls the loop at compile time (BLOCK_SIZE must be constexpr)
    for i in tl.static_range(BLOCK_SIZE):
        x = tl.load(x_ptr + i)
        tl.store(out_ptr + i, x * x)


def test_static_range():
    BLOCK_SIZE = 128
    x = torch.randn(BLOCK_SIZE, device="npu", dtype=torch.float32)
    out = torch.empty(BLOCK_SIZE, device="npu", dtype=torch.float32)

    static_range_kernel[(1, )](x, out, BLOCK_SIZE=BLOCK_SIZE)
    torch.npu.synchronize()

    assert_close(out, x * x, rtol=1e-3, atol=1e-3)


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

Special Restrictions

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

• start, end, step must be compile-time constants (tl.constexpr).

__init__(arg1, arg2=None, step=None)

Methods

__init__(arg1[, arg2, step])

Attributes

type