triton.language.condition

class triton.language.condition(arg1, disable_licm=False)

While loop condition wrapper.

@triton.jit
def kernel(...):
    while tl.condition(c, disable_licm)
        ...

Note:

This is a special wrapper used to annotate while loops in the context of triton.jit functions. It allows user to pass extra attributes to the compiler.

参数:

disable_licm -- Tells the compiler it shouldn't hoist loop invariant code outside the loop. This is often useful to avoid creating long liveranges within a loop.

示例

import torch
import torch_npu
from torch.testing import assert_close

import triton
import triton.language as tl


@triton.jit
def condition_kernel(x_ptr, out_ptr, N: tl.constexpr):
    acc = 0.0
    i = 0
    # Annotate the while-loop condition with tl.condition so that the
    # compiler keeps loop-invariant code inside the loop (disable_licm).
    while tl.condition(i < N, disable_licm=True):
        acc += tl.load(x_ptr + i)
        i += 1
    tl.store(out_ptr, acc)


def test_condition():
    N = 128
    x = torch.randn(N, device="npu", dtype=torch.float32)
    out = torch.empty(1, device="npu", dtype=torch.float32)

    condition_kernel[(1, )](x, out, N=N)
    torch.npu.synchronize()

    assert_close(out, x.sum().reshape(1), rtol=1e-3, atol=1e-3)
    print("test_condition PASSED!")


if __name__ == "__main__":
    test_condition()

数据类型支持

平台	uint8	int8	uint16	int16	uint32	int32	uint64	int64	fp16	fp32	fp64	bf16	fp8e(e4m3)	fp8e5(e5m2)	bool
Ascend A2/A3	×	×	×	×	×	×	×	×	×	×	×	×	×	×	√
Ascend 950	×	×	×	×	×	×	×	×	×	×	×	×	×	×	√

__init__(arg1, disable_licm=False)

Methods

__init__(name)

Initialize self.

Attributes

type