triton.language.ge

triton.language.ge(x, y)

Element-wise greater-or-equal comparison, the >= operator on tensors.

• float — ordered greater-or-equal (fcmp OGE)

• signed int — signed greater-or-equal (icmp SGE)

• unsigned int — unsigned greater-or-equal (icmp UGE)

The result is an int1 (boolean) tensor with the same shape.

参数:

• x (Block) -- the left operand

• y (Block) -- the right operand

示例

import torch

import triton
import triton.language as tl


@triton.jit
def triton_ge_3d(in_ptr0, in_ptr1, out_ptr0, L: tl.constexpr, M: tl.constexpr, N: tl.constexpr):
    lblk_idx = tl.arange(0, L)
    mblk_idx = tl.arange(0, M)
    nblk_idx = tl.arange(0, N)
    idx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :]
    x0 = tl.load(in_ptr0 + idx)
    x1 = tl.load(in_ptr1 + idx)
    ret = x0 >= x1
    odx = lblk_idx[:, None, None] * N * M + mblk_idx[None, :, None] * N + nblk_idx[None, None, :]
    tl.store(out_ptr0 + odx, ret)


def test_ge():
    shape = (2, 4, 8)
    torch.manual_seed(0)
    x_cpu = torch.randn(shape, dtype=torch.float32)
    y_cpu = torch.randn(shape, dtype=torch.float32)

    out = torch.empty(shape, dtype=torch.bool, device="npu")
    triton_ge_3d[(1, 1, 1)](x_cpu.npu(), y_cpu.npu(), out, L=shape[0], M=shape[1], N=shape[2], debug=True)

    assert torch.equal(out.cpu(), torch.ge(x_cpu, y_cpu))


if __name__ == "__main__":
    test_ge()

数据类型支持

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