triton.language.permute

triton.language.permute(input, *dims, _semantic=None)

Permutes the dimensions of a tensor.

参数:

• input (Block) -- The input tensor.

• dims -- The desired ordering of dimensions. For example, (2, 1, 0) reverses the order dims in a 3D tensor.

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

# These are equivalent
permute(x, (2, 1, 0))
permute(x, 2, 1, 0)

trans() is equivalent to this function, except when dims is empty, it tries to swap the last two axes.

This function can also be called as a member function on tensor, as x.permute(...) instead of permute(x, ...).

Example

import torch
import triton
import triton.language as tl


@triton.jit
def permute_3d_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr):
    # (M, N, K) -> (K, M, N)
    offsets_m = tl.arange(0, M)[:, None, None]
    offsets_n = tl.arange(0, N)[None, :, None]
    offsets_k = tl.arange(0, K)[None, None, :]
    x = tl.load(x_ptr + offsets_m * N * K + offsets_n * K + offsets_k)
    y = tl.permute(x, [2, 0, 1])
    flat = tl.reshape(y, (M * N * K, ))
    tl.store(out_ptr + tl.arange(0, M * N * K), flat)


def test_permute():
    M, N, K = 2, 3, 4
    x = torch.zeros([M, N, K], dtype=torch.float32).npu()
    out = torch.empty((M * N * K, ), dtype=torch.float32, device="npu")

    permute_3d_kernel[(1, )](out, x, M=M, N=N, K=K)

    assert out.shape == (M * N * K, ), "Shape mismatch"
    print("Test passed.")


if __name__ == "__main__":
    test_permute()

Special Restrictions

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