triton.language.view

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

Returns a tensor with the same elements as input but a different shape. The order of the elements may not be preserved.

参数:

• input (Block) -- The input tensor.

• shape -- The desired shape.

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

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

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

Example

import torch
import triton
import triton.language as tl


@triton.jit
def view_kernel(x_ptr, out_ptr, D1: tl.constexpr, D2: tl.constexpr, D3: tl.constexpr, R1: tl.constexpr,
                R2: tl.constexpr):
    # (D1, D2, D3) viewed as (R1, R2)
    d1 = tl.arange(0, D1)[:, None, None]
    d2 = tl.arange(0, D2)[None, :, None]
    d3 = tl.arange(0, D3)[None, None, :]
    x = tl.load(x_ptr + d1 * D2 * D3 + d2 * D3 + d3)
    y = tl.view(x, [R1, R2])
    flat = tl.reshape(y, (D1 * D2 * D3, ))
    tl.store(out_ptr + tl.arange(0, D1 * D2 * D3), flat)


def test_view():
    D1, D2, D3 = 2, 3, 4
    R1, R2 = 6, 4
    x = torch.zeros([D1, D2, D3], dtype=torch.float32).npu()
    out = torch.empty((D1 * D2 * D3, ), dtype=torch.float32, device="npu")
    view_kernel[(1, )](out, x, D1, D2, D3, R1, R2)

    assert out.shape == (D1 * D2 * D3, ), f"Shape mismatch: expected {(D1*D2*D3,)} but got {out.shape}."
    print("Test passed.")


if __name__ == "__main__":
    test_view()

Special Restrictions

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