triton.language.ravel

triton.language.ravel(x, can_reorder=False)

Returns a contiguous flattened view of x.

Parameters:

x – the input tensor

Example

import torch
import triton
import triton.language as tl


@triton.jit
def flatten_kernel(x_ptr, output_ptr, M: tl.constexpr, N: tl.constexpr):
    # Flatten (M, N) to 1D
    offsets_m = tl.arange(0, M)[:, None]
    offsets_n = tl.arange(0, N)[None, :]
    x = tl.load(x_ptr + offsets_m * N + offsets_n)
    x_flat = tl.ravel(x)
    tl.store(output_ptr + tl.arange(0, M * N), x_flat)


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

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


if __name__ == "__main__":
    test_ravel()

DataType Support

平台	uint8	int8	uint16	int16	uint32	int32	uint64	int64	fp16	fp32	fp64	bf16	fp8e(e4m3)	fp8e5(e5m2)	bool
Ascend A2/A3	√	√	×	√	×	√	×	√	√	√	×	√	×	×	√
Ascend 950	√	√	√	√	√	√	√	√	√	√	×	√	√	√	√