triton.language.interleave

triton.language.interleave = <function interleave>

Interleaves the values of two tensors along their last dimension. The two tensors must have the same shape. Equivalent to tl.join(a, b).reshape(a.shape[:-1] + [2 * a.shape[-1]])

参数:

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

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

Example

import torch
import triton
import triton.language as tl


@triton.jit
def interleave_2d_kernel(x_ptr, y_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr):
    # (M, N) interleave (M, N) -> (M, 2 * N)
    offsets_m = tl.arange(0, M)[:, None]
    offsets_n = tl.arange(0, N)[None, :]
    x = tl.load(x_ptr + offsets_m * N + offsets_n)
    y = tl.load(y_ptr + offsets_m * N + offsets_n)
    z = tl.interleave(x, y)
    offsets_n2 = tl.arange(0, 2 * N)[None, :]
    tl.store(out_ptr + offsets_m * (2 * N) + offsets_n2, z)


def test_interleave():
    M, N = 2, 3
    x = torch.zeros([M, N], dtype=torch.float32).npu()
    y = torch.zeros([M, N], dtype=torch.float32).npu()

    out = torch.empty((M, 2 * N), dtype=torch.float32, device="npu")
    interleave_2d_kernel[(1, )](out, x, y, M=M, N=N)

    assert out.shape == (M, 2 * N), f"Shape mismatch: expected {(M, 2*N)}, but got {out.shape}."
    print("Test passed.")


if __name__ == "__main__":
    test_interleave()

Special Restrictions

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