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.

Parameters:

asm – the inline assembly code
constraints – the LLVM constraint string (only ‘l’ is supported on Ascend)
args – the input tensors
dtype – the output data type(s)
is_pure – whether the assembly is pure (no side effects)
pack – the number of elements to process at a time

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!")

DataType Support

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

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.

triton.language.inline_asm_elementwise

Contents

triton.language.inline_asm_elementwise#