triton.language.cat

triton.language.cat(input, other, can_reorder=False, _semantic=None)

Concatenate the given blocks

参数:

• input (Tensor) -- The first input tensor.

• other (Tensor) -- The second input tensor.

• reorder -- Compiler hint. If true, the compiler is allowed to reorder elements while concatenating inputs. Only use if the order does not matter (e.g., result is only used in reduction ops). Current implementation of cat supports only can_reorder=True.

Example

import torch
import triton
import triton.language as tl


@triton.jit
def cat_1d_kernel(output_ptr, x_ptr, y_ptr, XB: tl.constexpr):
    """
    Concatenate two 1D tensors x and y of length XB into a continuous tensor of length XB*2.
    """
    idx = tl.arange(0, XB)
    X = tl.load(x_ptr + idx)
    Y = tl.load(y_ptr + idx)

    ret = tl.cat(X, Y, can_reorder=True)

    oidx = tl.arange(0, XB * 2)
    tl.store(output_ptr + oidx, ret)


def test_cat_1d():
    XB = 8

    x = torch.randn(XB, dtype=torch.float32).npu()
    y = torch.randn(XB, dtype=torch.float32).npu()
    output = torch.zeros(XB * 2, dtype=torch.float32).npu()

    cat_1d_kernel[(1, )](output, x, y, XB=XB)

    expected = torch.cat([x.cpu(), y.cpu()], dim=0)

    assert torch.allclose(output.cpu(), expected.cpu())


if __name__ == "__main__":
    test_cat_1d()

Special Restrictions

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

• can_reorder: Both Ascend and GPU only support can_reorder=True