triton.language.device_print#

triton.language.device_print(prefix, *args, hex=False, _semantic=None)#

Print the values at runtime from the device. String formatting does not work for runtime values, so you should provide the values you want to print as arguments. The first value must be a string, all following values must be scalars or tensors.

Calling the Python builtin print is the same as calling this function, and the requirements for the arguments will match this function (not the normal requirements for print).

tl.device_print("pid", pid)
print("pid", pid)

On Ascend, device_print is available when the environment variable TRITON_DEVICE_PRINT=1 is set. The output is streamed through a device buffer of limited size.

Parameters:

prefix – a prefix to print before the values. This is required to be a string literal.
args – the values to print. They can be any tensor or scalar.
hex – print all values as hex instead of decimal

Example

import os
import torch
import torch_npu

import triton
import triton.language as tl

os.environ["TRITON_DEVICE_PRINT"] = "1"


@triton.jit
def device_print_kernel(x_ptr, BLOCK_SIZE: tl.constexpr):
    offsets = tl.arange(0, BLOCK_SIZE)
    x = tl.load(x_ptr + offsets)
    # device_print runs at runtime; the first argument must be a string prefix.
    # Set TRITON_DEVICE_PRINT=1 to see the output.
    tl.device_print("x =", x)


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


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

DataType Support

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

Special Restrictions

prefix: the first argument must be a string prefix; omitting it causes a compilation error.
Set environment variable TRITON_DEVICE_PRINT=1 to enable.

triton.language.device_print

Contents

triton.language.device_print#