triton.language.extra.cann.libdevice.jn

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

Computes the Bessel function of the first kind of integer order n.

Example

import os

# The libdevice SIMT ops below are A5-only (Ascend 910_95 / 950) and are
# additionally gated by this env switch; set it so the examples run on A5
# hardware without extra configuration.
os.environ.setdefault("TRITON_ENABLE_LIBDEVICE_SIMT", "1")

import pytest
import triton
import triton.language as tl
import triton.language.extra.cann.libdevice as libdevice
import torch
from triton.backends.ascend.utils import triton_enable_libdevice_simt

_SIMT_SKIP_MSG = ("SIMT libdevice ops require an Ascend 950 target "
                  "with TRITON_ENABLE_LIBDEVICE_SIMT=1; skipping.")


@triton.jit
def triton_kernel(input0, input1, 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(input0 + (x0), mask=mask)
        tmp1 = tl.load(input1 + (x0), mask=mask)
        tmp2 = libdevice.jn(tmp0, tmp1)
        tl.store(output + (x0), tmp2, mask=mask)


if __name__ == "__main__":
    if not triton_enable_libdevice_simt():
        print(_SIMT_SKIP_MSG)
    else:
        x0 = (torch.randint(1, 16, (8, ))).to(torch.int32).npu()
        x1 = (torch.rand((8, )) + 0.1).to(torch.float32).npu()
        output = torch.empty(8, dtype=torch.float32, device='npu')
        triton_kernel[(1, )](x0, x1, output, 8, XBLOCK=8, XBLOCK_SUB=8, force_simt_only=True)

Special Restrictions

• • n: int32

• • x: float32

• Return value: tl.tensor, returns the Bessel function of the first kind of order n.

• Return type: float32

• Compilation modes: SIMT