triton.language.extra.cann.libdevice.pow

Contents

triton.language.extra.cann.libdevice.pow#

triton.language.extra.cann.libdevice.pow(arg0: None, arg1: None, _semantic: None = 'None')#

Computes arg0 raised to the power of arg1.

Parameters:

  • arg0 (tl.tensor) – The base tensor. Supported dtypes: fp32, fp16, bf16.

  • arg1 (tl.tensor) – The exponent tensor. Supported dtypes: fp32, fp16, bf16, int32.

Example

import triton
import triton.language as tl
import triton.language.extra.cann.libdevice as libdevice
import torch


@triton.jit
def triton_kernel(input, input2, output, n_elements, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr):
    offset = tl.program_id(0) * XBLOCK
    base = tl.arange(0, XBLOCK_SUB)
    loops: tl.constexpr = XBLOCK // XBLOCK_SUB
    for loop in range(loops):
        x0 = offset + (loop * XBLOCK_SUB) + base
        mask = x0 < n_elements
        tmp0 = tl.load(input + (x0), mask=mask)
        tmp1 = tl.load(input2 + (x0), mask=mask)
        tmp2 = libdevice.pow(tmp0, tmp1)
        tl.store(output + (x0), tmp2, mask=mask)


def test_pow():
    param_list = [(2, 256, 4), 2, 2048, 1024]
    shape, ncore, xblock, xblock_sub = param_list
    x0 = (torch.rand(size=shape, dtype=torch.float32) * 2 + 0.1).npu()
    x1 = (torch.rand(size=shape, dtype=torch.float32) * 2).npu()

    torch_res = torch.pow(x0, x1)
    triton_res = torch.empty_like(x0)
    triton_kernel[ncore, 1, 1](x0, x1, triton_res, x0.numel(), xblock, xblock_sub)

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


if __name__ == "__main__":
    test_pow()

Special Restrictions

    • x: float32

    • y: float32

  • Return value: tl.tensor, returns x raised to the power y.

  • Return type: float32

  • Compilation modes: SIMT, SIMD

Contents