triton.language.device_assert

triton.language.device_assert(cond, msg='', mask=None, _semantic=None)

Assert the condition at runtime from the device. Requires that the environment variable TRITON_DEBUG is set to a value besides 0 in order for this to have any effect.

Using the Python assert statement is the same as calling this function, except that the second argument must be provided and must be a string, e.g. assert pid == 0, "pid != 0". The environment variable must be set for this assert statement to have any effect.

tl.device_assert(pid == 0)
assert pid == 0, f"pid != 0"

参数:

• cond -- the condition to assert. This is required to be a boolean tensor.

• msg -- the message to print if the assertion fails. This is required to be a string literal.

Example

import torch
import torch_npu

import triton
import triton.language as tl


@triton.jit
def device_assert_kernel(x_ptr, BLOCK_SIZE: tl.constexpr):
    offsets = tl.arange(0, BLOCK_SIZE)
    x = tl.load(x_ptr + offsets)
    # device_assert checks a condition at runtime on device.
    # Requires both TRITON_DEBUG=1 and TRITON_DEVICE_PRINT=1 to take effect.
    tl.device_assert(x == x, "x contains NaN")


def test_device_assert():
    BLOCK_SIZE = 8
    x = torch.randn(BLOCK_SIZE, device="npu", dtype=torch.float32)
    device_assert_kernel[(1, )](x, BLOCK_SIZE=BLOCK_SIZE)
    torch.npu.synchronize()


if __name__ == "__main__":
    test_device_assert()
    print("test_device_assert PASSED!")