triton.language.range

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

一个永远向上计数的迭代器。

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

注意：

这是一个特殊的迭代器，用于在 triton.jit 函数的上下文中实现与 Python 的 range 相似的语义。此外，它允许用户向编译器传递额外属性。

参数：

• arg1 – 起始值。

• arg2 – 结束值。

• step – 步长值。

• num_stages –

  将循环流水线化为指定数量的阶段（以便同时有 num_stages 次循环迭代在进行中）。

  注意，这与将 num_stages 作为内核参数传递略有不同。内核参数仅流水线化馈送到 dot 操作的加载，而此属性尝试流水线化此循环中的大多数（但非全部）加载。

• loop_unroll_factor – 告诉 Triton IR 级的循环展开器将使用此 range 的 for 循环展开多少次。该值小于 2 表示不进行展开。

• disallow_acc_multi_buffer – 如果为 true，则在适用情况下，阻止循环中 dot 操作的累加器进行多重缓冲。

• flatten – 自动展平从此循环开始的循环嵌套，以创建一个单一的展平循环。编译器将尝试对展平后的循环进行流水线化，这可以避免阶段停顿。

• warp_specialize –

  启用循环的自动 warp 特殊化。编译器将尝试将循环中的内存、MMA 和向量操作划分为单独的异步分区。这将增加内核所需的总 warp 数量。

  注意，warp 特殊化仅在 Blackwell GPU 上受支持，并且仅适用于简单的 matmul 循环。对任意循环的支持将随时间推移而扩展。

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

方法

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