triton.language.maximum

triton.language.maximum(x, y, propagate_nan: constexpr = <MagicMock name='triton._C.libtriton.ir.PROPAGATE_NAN.NONE' id='129320157032000'>, _semantic=None)

Computes the element-wise maximum of x and y.

参数:

• x (Block) -- the first input tensor

• y (Block) -- the second input tensor

• propagate_nan (tl.PropagateNan) -- whether to propagate NaN values.

参见

tl.PropagateNan

Example

import triton
import triton.language as tl
import torch


def torch_maximum(x0, x1):
    res = torch.maximum(x0, x1)
    return res


@triton.jit
def triton_maximum(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB):
        x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)
        xmask = x_index < xnumel
        tmp0 = tl.load(in_ptr0 + x_index, xmask)
        tmp1 = tl.load(in_ptr1 + x_index, xmask)
        tmp2 = tl.maximum(tmp0, tmp1)
        tl.store(out_ptr0 + x_index, tmp2, xmask)


def test_maximum(param_list):
    param_list = ['float32', (2, 4096, 8), 2, 32768, 1024]
    dtype, shape, ncore, xblock, xblock_sub = param_list
    x0 = torch.randn(size=shape, dtype=eval('torch.' + dtype)).npu()
    x1 = torch.randn(size=shape, dtype=eval('torch.' + dtype)).npu()

    torch_res = torch_maximum(x0, x1)
    triton_res = torch.empty_like(x0)
    triton_maximum[ncore, 1, 1](x0, x1, triton_res, x0.numel(), xblock, xblock_sub)

    torch.testing.assert_close(torch_res, triton_res, rtol=1e-04, atol=1e-04, equal_nan=True)


if __name__ == '__main__':
    test_maximum()

Special Restrictions

• DataType: Ascend does not support fp64 (hardware limitation).