triton.language.map_elementwise¶

triton.language.map_elementwise(scalar_fn: Callable[[...], Tuple[tensor, ...]], *args: tensor, pack=1, _semantic=None, _generator=None)¶

Map a scalar function over a tensor.

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

This may be useful in allowing control flow over single elements in a tensor, for example a multi-branch function with if/elif/else chains or for loops, providing more flexible element-wise computation than tl.where which only supports binary selection.

@triton.jit
def compare_scalar(x, y):
    if x < y:
        return -1
    elif x == y:
        return 0
    else:
        return 1

@triton.jit
def compare(x, y):
    return tl.map_elementwise(compare_scalar, x, y)

参数:

scalar_fn -- the function to map over.
pack -- the number of elements to be processed by one function call. On Ascend, this parameter has no semantic effect as the implementation is always vectorized.

返回:

one tensor or a tuple of tensors, depending on the mapped function.

示例

import torch
import torch_npu
import triton
import triton.language as tl


@triton.jit
def _compare(x, y):
    if x < y:
        return -1
    elif x == y:
        return 0
    else:
        return 1


@triton.jit
def kernel(X, Y, Z, BLOCK: tl.constexpr):
    x = tl.load(X + tl.arange(0, BLOCK))
    y = tl.load(Y + tl.arange(0, BLOCK))
    z = tl.map_elementwise(_compare, x, y)
    tl.store(Z + tl.arange(0, BLOCK), z)


def test_map_elementwise():
    shape = (128, )
    x = torch.randint(-100, 100, shape, dtype=torch.int32, device='npu')
    y = torch.randint(-100, 100, shape, dtype=torch.int32, device='npu')
    z = torch.zeros(shape, dtype=torch.int32, device='npu')
    kernel[(1, )](x, y, z, BLOCK=shape[0])
    expected = (x > y).int() - (y > x).int()
    assert torch.equal(z.cpu(), expected.cpu())


if __name__ == "__main__":
    test_map_elementwise()

特殊说明

while loops are not supported inside the scalar function.
pack has no semantic effect on NPU backends as the implementation is always vectorized.