triton.language.inline_asm_elementwise

triton.language.inline_asm_elementwise(asm: str, constraints: str, args: Sequence, dtype: dtype | Sequence[dtype], is_pure: bool, pack: int, _semantic=None)

Execute inline assembly over a tensor. Essentially, this is map where the function is inline assembly.

The input tensors args are implicitly broadcasted to the same shape.

dtype can be a tuple of types, in which case the output is a tuple of tensors.

Each invocation of the inline asm processes pack elements at a time. Exactly which set of inputs a block receives is unspecified. Input elements of size less than 4 bytes are packed into 4-byte registers.

This op does not support empty dtype -- the inline asm must return at least one tensor, even if you don't need it. You can work around this by returning a dummy tensor of arbitrary type; it shouldn't cost you anything if you don't use it.

Example using PTX assembly:

@triton.jit
def kernel(A, B, C, D, BLOCK: tl.constexpr):
    a = tl.load(A + tl.arange(0, BLOCK)) # uint8 tensor
    b = tl.load(B + tl.arange(0, BLOCK)) # float32 tensor

    # For each (a,b) in zip(a,b), perform the following:
    # - Let ai be `a` converted to int32.
    # - Let af be `a` converted to float.
    # - Let m be the max of ai and b.
    # - Return ai and mi.
    # Do the above 4 elements at a time.
    (c, d) = tl.inline_asm_elementwise(
        asm="""
        {
            // Unpack `a` into `ai`.
            .reg .b8 tmp<4>;
            mov.b32 {tmp0, tmp1, tmp2, tmp3}, $8;
            cvt.u32.u8 $0, tmp0;
            cvt.u32.u8 $1, tmp1;
            cvt.u32.u8 $2, tmp2;
            cvt.u32.u8 $3, tmp3;
        }
        // Convert `ai` to float.
        cvt.rn.f32.s32 $4, $0;
        cvt.rn.f32.s32 $5, $1;
        cvt.rn.f32.s32 $6, $2;
        cvt.rn.f32.s32 $7, $3;
        // Take max of `ai` and `b`.
        max.f32 $4, $4, $9;
        max.f32 $5, $5, $10;
        max.f32 $6, $6, $11;
        max.f32 $7, $7, $12;
        """,
        constraints=(
            # 8 output registers, namely
            #   $0=ai0, $1=ai1, $2=ai2, $3=ai3,
            #   $4=m0,  $5=m1,  $6=m2,  $7=m3.
            "=r,=r,=r,=r,=r,=r,=r,=r,"
            # 5 input registers, namely
            #   $8=ai,
            #   $9=b0, $10=b1, $11=b2, $12=b3.
            # The four elements from `a` are all packed into one register.
            "r,r,r,r,r"),
        args=[a, b],
        dtype=(tl.int32, tl.float32),
        is_pure=True,
        pack=4,
    )
    tl.store(C + tl.arange(0, BLOCK), c)
    tl.store(D + tl.arange(0, BLOCK), d)

参数:

• asm -- assembly to run. Must match target's assembly format.

• constraints -- asm constraints in LLVM format

• args -- the input tensors, whose values are passed to the asm block

• dtype -- the element type(s) of the returned tensor(s)

• is_pure -- if true, the compiler assumes the asm block has no side-effects

• pack -- the number of elements to be processed by one instance of inline assembly

返回:

one tensor or a tuple of tensors of the given dtypes

Example

import torch
import torch_npu
from torch.testing import assert_close

import triton
import triton.language as tl


@triton.jit
def inline_asm_add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    out = tl.inline_asm_elementwise(
        asm="ADD.s64 $0, $1, $2",
        constraints="=l,l,l",
        args=[x, y],
        dtype=tl.int64,
        is_pure=True,
        pack=1,
    )
    tl.store(out_ptr + offsets, out, mask=mask)


def test_inline_asm_elementwise():
    N = 128
    BLOCK_SIZE = 128
    x = torch.randint(0, 1000, (N, ), device="npu", dtype=torch.int64)
    y = torch.randint(0, 1000, (N, ), device="npu", dtype=torch.int64)
    out = torch.empty(N, device="npu", dtype=torch.int64)

    grid = (triton.cdiv(N, BLOCK_SIZE), )
    inline_asm_add_kernel[grid](x, y, out, N, BLOCK_SIZE=BLOCK_SIZE)
    torch.npu.synchronize()

    assert_close(out, x + y)


if __name__ == "__main__":
    test_inline_asm_elementwise()
    print("test_inline_asm_elementwise PASSED!")

Special Restrictions

• Inline assembly registers only support int64 (s64) and float32 (f32).

• Only the 'l' LLVM constraint is supported.

• Only 1-D input tensors are supported; higher-dimensional tensors must be flattened.