triton.language.floordiv

triton.language.floordiv(x, y)

Element-wise integer division of x by y, the // operator on tensors.

The result truncates toward zero. Only integer operands are supported. Signed integers use signed division (sdiv); unsigned integers use unsigned division (udiv). Floating-point operands raise TypeError.

参数:

• x (Block or scalar number) -- the dividend

• y (Block or scalar number) -- the divisor

示例

import triton
import triton.language as tl
import torch


def torch_floordiv(x0, x1):
    res = x0 // x1
    return res


@triton.jit
def triton_kernel(out_ptr0, in_ptr0, in_ptr1, N: tl.constexpr):
    idx = tl.arange(0, N)
    x = tl.load(in_ptr0 + idx)
    y = tl.load(in_ptr1 + idx)
    ret = x // y
    tl.store(out_ptr0 + idx, ret)


def test_floordiv():
    param_list = ['int32', (2, 256, 2), 2]
    dtype, shape, ncore = param_list
    x0 = torch.randint(1, 10, size=shape, dtype=eval('torch.' + dtype)).npu()
    x1 = torch.randint(1, 10, size=shape, dtype=eval('torch.' + dtype)).npu()

    torch_res = torch_floordiv(x0, x1)
    triton_res = torch.empty_like(x0)
    triton_kernel[ncore, 1, 1](triton_res, x0, x1, N=x0.numel())

    torch.testing.assert_close(torch_res, triton_res, rtol=0, atol=0, equal_nan=True)


if __name__ == '__main__':
    test_floordiv()

数据类型支持

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