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 级别的循环展开器，对此 range 所使用的 for 循环进行展开的次数。如果该值小于 2，则表示不进行展开。

• disallow_acc_multi_buffer – 如果为 true，则在适用情况下，阻止循环中的点积运算累加器进行多缓冲。

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

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

• disable_licm –

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

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

属性

类型