triton.language.neg

triton.language.neg(x)

Element-wise negation, the unary - operator on tensors.

Equivalent to 0 - x. Pointer-typed inputs raise ValueError.

参数:

x (Block) -- the input

示例

import pytest
import torch
import torch_npu
import triton
import triton.language as tl


@triton.jit
def triton_neg(in_ptr0, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr):
    offset = tl.program_id(0) * XBLOCK
    base1 = tl.arange(0, XBLOCK_SUB)
    loops1 = XBLOCK // XBLOCK_SUB
    for loop1 in range(loops1):
        x0 = offset + (loop1 * XBLOCK_SUB) + base1
        tmp0 = tl.load(in_ptr0 + (x0), None)
        tmp1 = -tmp0
        tl.store(out_ptr0 + (x0), tmp1, None)


def test_neg():
    dtype, shape, ncore, xblock, xblock_sub = ['float32', (8, 8), 8, 8, 8]
    x0 = torch.randint(low=0, high=256, size=shape, dtype=eval(f'torch.{dtype}')).npu()
    ans = -x0
    y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu()
    triton_neg[ncore, 1, 1](x0, y_cal, xblock, xblock_sub)
    torch.testing.assert_close(ans, y_cal)


if __name__ == "main":
    test_neg()

数据类型支持

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