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):
        ...

Note:

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

Parameters:

• arg1 – 起始值。

• arg2 – 结束值。

• step – 步长值。

• num_stages –

  将循环流水线划分为这么多阶段（因此同时会有num_stages次循环迭代在进行中）。

  请注意这与将num_stages作为内核参数传递有微妙差异。内核参数仅对输入到dot操作的加载进行流水线处理，而此属性尝试对该循环中的大多数（尽管不是全部）加载进行流水线处理。

• loop_unroll_factor – 告诉Triton IR级别的循环展开器在使用此范围时对for循环进行多少次展开。该值小于2意味着不进行展开。

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

• flatten - 自动从当前循环开始展平循环嵌套，创建单个扁平化循环。编译器将尝试流水线化这个扁平化循环，从而避免阶段停滞。

• warp_specialize –

  在循环上启用自动warp专业化。 编译器会尝试将循环中的内存操作、MMA和向量操作分区到不同的异步分区中。这将增加内核所需的warp总数。

  请注意，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)

方法

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