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, disable_licm=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 层的循环展开器，对此范围所使用的 for 循环进行多少次展开。此值小于 2 意味着不进行展开。

• disallow_acc_multi_buffer – 如果为真，则阻止循环中 dot 操作的累加器进行多重缓冲（如果适用）。

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

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

• disable_licm –

  告诉编译器不应将循环不变量代码提升到循环外部。这通常有助于避免在循环内创建过长的生存期（liveranges）。

  请注意，Warp 特化仅在 Blackwell GPU 上受支持，并且仅适用于简单的矩阵乘法循环。对任意循环的支持将随着时间的推移逐步扩展。

__init__(self, 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)

方法

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

属性

类型