triton.language.div

triton.language.div(x, y)

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

The implementation is equivalent to fdiv(), but without the floating-point-only restriction: integer operands are automatically converted to floating point before the division.

• int / int — both operands are cast to float32

• int / float or float / int — the integer operand is cast to the float type

• float / float — both operands are unified to the higher-precision float type

参数:

• 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_div(x0, x1):
    res = x0 / x1
    return res


@triton.jit
def triton_div(in_ptr0, in_ptr1, out_ptr0, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr):
    offset = tl.program_id(0) * XBLOCK
    base1 = tl.arange(0, XBLOCK_SUB)
    loops1: tl.constexpr = (XBLOCK + XBLOCK_SUB - 1) // XBLOCK_SUB
    for loop1 in range(loops1):
        x0 = offset + (loop1 * XBLOCK_SUB) + base1
        tmp0 = tl.load(in_ptr0 + (x0), None)
        tmp1 = tl.load(in_ptr1 + (x0), None)
        tmp2 = tmp0 / tmp1
        tl.store(out_ptr0 + (x0), tmp2, None)


def test_div():
    param_list = ['float32', (2, 4096, 8), 2, 32768, 1024]
    dtype, shape, ncore, xblock, xblock_sub = param_list
    x0 = torch.randn(size=shape, dtype=eval('torch.' + dtype)).npu()
    x1 = torch.randn(size=shape, dtype=eval('torch.' + dtype)).npu()

    torch_res = torch_div(x0, x1)
    triton_res = torch.empty_like(x0)
    triton_div[ncore, 1, 1](x0, x1, triton_res, xblock, xblock_sub)

    torch.testing.assert_close(torch_res, triton_res, rtol=1e-04, atol=1e-04, equal_nan=True)


if __name__ == '__main__':
    test_div()

数据类型支持

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