triton.language.or

triton.language.or(x, y)

Element-wise bitwise OR of x and y, the | operator on tensors.

Only integer operands are supported.

参数:

• x (Block) -- the first input

• y (Block) -- the second input

示例

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


@triton.jit
def fn_npu_(output_ptr, x_ptr, y_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.constexpr, XNUMEL: tl.constexpr,
            YNUMEL: tl.constexpr, ZNUMEL: tl.constexpr):
    xoffs = tl.program_id(0) * XB
    yoffs = tl.program_id(1) * YB
    zoffs = tl.program_id(2) * ZB

    xidx = tl.arange(0, XB) + xoffs
    yidx = tl.arange(0, YB) + yoffs
    zidx = tl.arange(0, ZB) + zoffs

    idx = xidx[:, None, None] * YNUMEL * ZNUMEL + yidx[None, :, None] * ZNUMEL + zidx[None, None, :]

    X = tl.load(x_ptr + idx)
    Y = tl.load(y_ptr + idx)
    ret = X | Y
    tl.store(output_ptr + idx, ret)


def test_or():
    M, N = (8, 16)
    dtype = 'int32'
    x = torch.randint(low=0, high=256, size=(M, N), dtype=eval(f'torch.{dtype}')).npu()
    y = torch.randint(low=0, high=256, size=(M, N), dtype=eval(f'torch.{dtype}')).npu()
    ans = x | y
    output = torch.zeros_like(ans)
    fn_npu_[1, 1, N](output, x, y, 1, M, 1, 1, M, N)
    torch.testing.assert_close(ans, output)


if __name__ == "main":
    test_or()

数据类型支持

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