triton.language.broadcast_to

triton.language.broadcast_to(input, *shape, _semantic=None)

Tries to broadcast the given tensor to a new shape.

参数:

  • input (Block) -- The input tensor.

  • shape -- The desired shape.

shape can be passed as a tuple or as individual parameters:

# These are equivalent
broadcast_to(x, (32, 32))
broadcast_to(x, 32, 32)

This function can also be called as a member function on :py:class:`tensor`,
as :code:`x.broadcast_to(...)` instead of
:code:`broadcast_to(x, ...)`.

Example

import torch
import triton
import triton.language as tl


@triton.jit
def matrix_add_bias_kernel(x_ptr, bias_ptr, output_ptr, M: tl.constexpr, N: tl.constexpr, BLOCK_M: tl.constexpr,
                           BLOCK_N: tl.constexpr):
    # broadcast bias to (BLOCK_M, BLOCK_N)
    offsets_m = tl.arange(0, BLOCK_M)[:, None]
    offsets_n = tl.arange(0, BLOCK_N)[None, :]
    x = tl.load(x_ptr + offsets_m * N + offsets_n)
    bias = tl.load(bias_ptr + offsets_n)
    bias_broadcast = bias.broadcast_to([BLOCK_M, BLOCK_N])
    output = x + bias_broadcast
    tl.store(output_ptr + offsets_m * N + offsets_n, output)


def test_broadcast_to():
    M, N = 8, 16
    BLOCK_M, BLOCK_N = M, N

    torch.manual_seed(0)
    x = torch.randn((M, N), dtype=torch.float32).npu()
    bias = torch.randn((1, N), dtype=torch.float32).npu()

    output = torch.empty((M, N), dtype=torch.float32, device="npu")

    matrix_add_bias_kernel[(1, )](x, bias, output, M=M, N=N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)

    ref = x.cpu() + bias.cpu()

    torch.testing.assert_close(output.cpu(), ref)
    print("Test passed.")


if __name__ == "__main__":
    test_broadcast_to()

Special Restrictions

  • DataType: Ascend A2/A3 does not support fp64, fp8e4, fp8e5, uint16, uint32, uint64 (hardware limitation).

  • The rank of the input tensor's shape must match the rank of the target shape.