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 specialization。编译器将尝试将循环中的内存、MMA 和向量操作划分为单独的异步分区。这将增加内核所需的总 warp 数量。

• disable_licm –

  告诉编译器不应将循环不变代码提升到循环之外。这通常有助于避免在循环内创建长的 live range。

  请注意，warp specialization 仅在 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, disable_licm=False)

方法

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

属性

类型