triton.language.reduce_or

triton.language.reduce_or(input, axis, keep_dims=False)

Returns the reduce_or of all elements in the input tensor along the provided axis

参数:

• input (Tensor) -- the input values

• axis (int) -- the dimension along which the reduction should be done. If None, reduce all dimensions

• keep_dims (bool) -- if true, keep the reduced dimensions with length 1

This function can also be called as a member function on tensor, as x.reduce_or(...) instead of reduce_or(x, ...).

示例

import torch
import triton
import triton.language as tl


@triton.jit
def reduce_or_kernel(in_ptr, out_ptr, N: tl.constexpr):
    a = tl.load(in_ptr + tl.arange(0, N))
    b = tl.reduce_or(a, axis=0)
    tl.store(out_ptr, b)


def test_reduce_or():
    N = 16
    x = torch.zeros(N, dtype=torch.int32, device='npu')
    x[3] = 1
    x[10] = 1
    out = torch.zeros(1, dtype=torch.int32, device='npu')
    reduce_or_kernel[(1, )](x, out, N=N)
    expected = (x != 0).any().to(torch.int32).reshape(1)
    assert out.item() == expected.item(), f"reduce_or 错误: {out} vs {expected}"
    print("PASS: test_reduce_or_basic")


if __name__ == "__main__":
    test_reduce_or()

数据类型支持

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