xserv/crates/xserv-kernels/src/transpose.rs

use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};

unsafe extern "C" {
    fn launch_reshape_heads_bf16(
        inp: *const c_void,
        out: *mut c_void,
        seq_len: i32,
        num_heads: i32,
        head_dim: i32,
        stream: *mut c_void,
    );
    fn launch_merge_heads_bf16(
        inp: *const c_void,
        out: *mut c_void,
        seq_len: i32,
        num_heads: i32,
        head_dim: i32,
        stream: *mut c_void,
    );
    fn launch_transpose_hsd_to_shd_bf16(
        inp: *const c_void,
        out: *mut c_void,
        seq_len: i32,
        num_heads: i32,
        head_dim: i32,
        stream: *mut c_void,
    );
    fn launch_transpose_shd_to_hsd_bf16(
        inp: *const c_void,
        out: *mut c_void,
        seq_len: i32,
        num_heads: i32,
        head_dim: i32,
        stream: *mut c_void,
    );
    fn launch_repeat_kv_bf16(
        inp: *const c_void,
        out: *mut c_void,
        kv_heads: i32,
        n_rep: i32,
        seq_len: i32,
        head_dim: i32,
        stream: *mut c_void,
    );
    fn launch_strided_copy_bf16(
        inp: *const c_void,
        out: *mut c_void,
        numel: i32,
        ndim: i32,
        shape0: i32,
        shape1: i32,
        shape2: i32,
        shape3: i32,
        in_stride0: i32,
        in_stride1: i32,
        in_stride2: i32,
        in_stride3: i32,
        in_offset: i32,
        stream: *mut c_void,
    );
    fn launch_strided_copy_f32(
        inp: *const c_void,
        out: *mut c_void,
        numel: i32,
        ndim: i32,
        shape0: i32,
        shape1: i32,
        shape2: i32,
        shape3: i32,
        in_stride0: i32,
        in_stride1: i32,
        in_stride2: i32,
        in_stride3: i32,
        in_offset: i32,
        stream: *mut c_void,
    );
}

/// [S, H*D] → [1, H, S, D] on GPU (BF16)
pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
    unsafe {
        launch_reshape_heads_bf16(
            x.data_ptr() as _,
            out.data_ptr() as *mut c_void,
            seq_len as i32,
            num_heads as i32,
            head_dim as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    out
}

/// [1, H, S, D] → [S, H*D] on GPU (BF16)
pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let hidden = num_heads * head_dim;
    let out = Tensor::empty(&[seq_len, hidden], DType::BF16, x.device());
    unsafe {
        launch_merge_heads_bf16(
            x.data_ptr() as _,
            out.data_ptr() as *mut c_void,
            seq_len as i32,
            num_heads as i32,
            head_dim as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    out
}

/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
pub fn transpose_for_rope_gpu(
    x: &Tensor,
    seq_len: usize,
    num_heads: usize,
    head_dim: usize,
) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let out = Tensor::empty(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
    unsafe {
        launch_transpose_hsd_to_shd_bf16(
            x.data_ptr() as _,
            out.data_ptr() as *mut c_void,
            seq_len as i32,
            num_heads as i32,
            head_dim as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    out
}

/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
pub fn transpose_from_rope_gpu(
    x: &Tensor,
    seq_len: usize,
    num_heads: usize,
    head_dim: usize,
) -> Tensor {
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
    unsafe {
        launch_transpose_shd_to_hsd_bf16(
            x.data_ptr() as _,
            out.data_ptr() as *mut c_void,
            seq_len as i32,
            num_heads as i32,
            head_dim as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    out
}

/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
    if n_rep == 1 {
        return x.clone();
    }
    assert_eq!(x.dtype(), DType::BF16);
    assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
    let kv_heads = x.shape()[1];
    let seq_len = x.shape()[2];
    let head_dim = x.shape()[3];
    let new_heads = kv_heads * n_rep;
    let out = Tensor::empty(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
    unsafe {
        launch_repeat_kv_bf16(
            x.data_ptr() as _,
            out.data_ptr() as *mut c_void,
            kv_heads as i32,
            n_rep as i32,
            seq_len as i32,
            head_dim as i32,
            xserv_cuda::current_stream_raw(),
        );
    }
    out
}

/// Make a non-contiguous GPU tensor contiguous via a strided copy kernel.
/// Supports BF16 and F32, up to 4D tensors (padded to 4D internally).
pub fn strided_to_contiguous_gpu(x: &Tensor) -> Tensor {
    assert!(matches!(x.device(), Device::Cuda(_)), "expected GPU tensor");
    assert!(!x.is_contiguous(), "tensor is already contiguous");
    assert!(x.ndim() <= 4, "strided_to_contiguous_gpu supports up to 4D");

    let ndim = x.ndim();
    let numel = x.numel();

    // Pad shape and strides to 4D (prepend 1s for shape, 0s for strides)
    let mut shape4 = [1i32; 4];
    let mut strides4 = [0i32; 4];
    let pad = 4 - ndim;
    for i in 0..ndim {
        shape4[pad + i] = x.shape()[i] as i32;
        strides4[pad + i] = x.strides()[i] as i32;
    }

    let out = Tensor::empty(x.shape(), x.dtype(), x.device());

    // Use storage base pointer + element offset, because strides are relative to
    // element 0 of the storage, not the data_ptr() (which already adds byte offset).
    let storage_ptr = x.storage().gpu_buffer().as_ptr();
    let in_offset = x.offset() as i32;

    unsafe {
        match x.dtype() {
            DType::BF16 => launch_strided_copy_bf16(
                storage_ptr as _,
                out.data_ptr() as *mut c_void,
                numel as i32,
                ndim as i32,
                shape4[0],
                shape4[1],
                shape4[2],
                shape4[3],
                strides4[0],
                strides4[1],
                strides4[2],
                strides4[3],
                in_offset,
                xserv_cuda::current_stream_raw(),
            ),
            DType::F32 => launch_strided_copy_f32(
                storage_ptr as _,
                out.data_ptr() as *mut c_void,
                numel as i32,
                ndim as i32,
                shape4[0],
                shape4[1],
                shape4[2],
                shape4[3],
                strides4[0],
                strides4[1],
                strides4[2],
                strides4[3],
                in_offset,
                xserv_cuda::current_stream_raw(),
            ),
            _ => panic!(
                "strided_to_contiguous_gpu: unsupported dtype {:?}",
                x.dtype()
            ),
        }
    }
    out
}