triton.language.rshift

triton.language.rshift(x, y)

Element-wise right shift of x by y bits, the >> operator on tensors.

Only integer operands are supported. Signed integers use arithmetic shift (ashr, preserving the sign bit); unsigned integers use logical shift (lshr). Bit widths of both operands are checked for consistency before the shift.

参数:

• x (Block) -- the value to shift

• y (Block) -- the shift amount

示例

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


@triton.jit
def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr):
    idx_block = tl.arange(0, NUMEL)
    x = tl.load(in_ptr0 + idx_block, mask=idx_block < N)
    tmp = tl.cast(2, tl.int8)
    ret = x >> tmp
    tl.store(out_ptr0 + idx_block, ret, mask=idx_block < N)


def test_rshift():
    dtype, N = ['int32', 32]
    x0 = torch.randint(low=0, high=256, size=(N, ), dtype=eval(f'torch.{dtype}')).npu()
    ans = x0 >> 2
    out = torch.zeros((N, ), dtype=eval(f'torch.{dtype}')).npu()
    triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=N)
    torch.testing.assert_close(ans, out)


if __name__ == "main":
    test_rshift()

数据类型支持

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