//! Tensor-parallel primitives for xserv.
//!
//! Process model: one OS thread per TP rank, each bound to one GPU. NCCL is
//! used for the collective (AllReduce); a hand-rolled P2P AllReduce may replace
//! it later as a learning exercise (see docs/17-tensor-parallelism.md).

pub mod ffi;

use std::ffi::c_void;

use ffi::{NcclComm, NcclUniqueId};
use xserv_cuda::GpuBuffer;
use xserv_cuda::device;

pub use ffi::NcclUniqueId as UniqueId;

/// NCCL is issued on the thread's current launch stream (legacy null stream
/// by default, the capture stream during CUDA graph capture). The model's
/// kernels run on the same stream, so AllReduce stays correctly ordered after
/// the producing matmul and before the consuming kernel — no extra sync.
fn launch_stream() -> xserv_cuda::ffi::CudaStream {
    xserv_cuda::stream::current_stream_raw()
}

/// Generate a unique id on one rank (typically rank 0) and broadcast the bytes
/// to all ranks out-of-band (e.g. via a shared variable across threads).
pub fn get_unique_id() -> NcclUniqueId {
    let mut id = NcclUniqueId::default();
    ffi::check(unsafe { ffi::ncclGetUniqueId(&mut id) }, "ncclGetUniqueId");
    id
}

/// Per-rank tensor-parallel context: NCCL communicator + a dedicated stream.
pub struct TpContext {
    pub rank: usize,
    pub world: usize,
    pub device: u32,
    comm: NcclComm,
}

// The NCCL communicator is owned by exactly one rank thread.
unsafe impl Send for TpContext {}

impl TpContext {
    /// Initialize this rank. Must be called from the thread that will own this
    /// rank's GPU work; binds the thread to `device` first. All ranks must call
    /// this concurrently with the same `id` and `world`.
    pub fn init(rank: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
        device::set_device(device).expect("set_device");
        let mut comm: NcclComm = std::ptr::null_mut();
        // Wrap the concurrent inits in a group so they rendezvous without deadlock.
        ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
        ffi::check(
            unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, rank as i32) },
            "ncclCommInitRank",
        );
        ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
        Self {
            rank,
            world,
            device,
            comm,
        }
    }

    /// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
    pub fn all_reduce_sum_bf16(&self, buf: &mut GpuBuffer, count: usize) {
        self.all_reduce_sum_bf16_ptr(buf.as_mut_ptr() as *mut c_void, count);
    }

    /// In-place AllReduce(sum) directly on a device pointer (`count` BF16 elems),
    /// issued on the null stream so it is ordered with the model's kernels.
    /// Asynchronous: a later sync (e.g. the D2H logits copy) completes it.
    ///
    /// # Safety
    /// `ptr` must point to at least `count` BF16 elements of valid device memory
    /// on this rank's device. The reduction is in-place (send == recv).
    pub fn all_reduce_sum_bf16_ptr(&self, ptr: *mut c_void, count: usize) {
        if self.world == 1 {
            return; // nothing to reduce
        }
        ffi::check(
            unsafe {
                ffi::ncclAllReduce(
                    ptr as *const c_void,
                    ptr,
                    count,
                    ffi::NCCL_BF16,
                    ffi::NCCL_SUM,
                    self.comm,
                    launch_stream(),
                )
            },
            "ncclAllReduce",
        );
    }
}

impl Drop for TpContext {
    fn drop(&mut self) {
        if !self.comm.is_null() {
            unsafe { ffi::ncclCommDestroy(self.comm) };
        }
    }
}

/// Per-stage pipeline-parallel context: a NCCL communicator spanning all `P`
/// stages plus point-to-point send/recv of the hidden state to the neighbour
/// stages. Init is identical to `TpContext` (one comm across `world` ranks);
/// only the collective differs — PP hands off `[tokens, hidden]` between
/// consecutive stages instead of AllReducing within a layer.
pub struct PpContext {
    pub stage: usize,
    pub world: usize,
    pub device: u32,
    comm: NcclComm,
}

// The NCCL communicator is owned by exactly one stage thread.
unsafe impl Send for PpContext {}

impl PpContext {
    /// Initialize this stage. Must be called from the thread that owns this
    /// stage's GPU; binds the thread to `device` first. All stages call this
    /// concurrently with the same `id` and `world`.
    pub fn init(stage: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
        device::set_device(device).expect("set_device");
        let mut comm: NcclComm = std::ptr::null_mut();
        ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
        ffi::check(
            unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, stage as i32) },
            "ncclCommInitRank",
        );
        ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
        Self {
            stage,
            world,
            device,
            comm,
        }
    }

    /// Send `count` BF16 elements at `ptr` to `peer`, on the null stream so it is
    /// ordered after the producing matmul. Asynchronous — a later `synchronize`
    /// (the caller must do one before reusing/freeing the buffer) completes it.
    ///
    /// # Safety
    /// `ptr` must point to at least `count` BF16 elements of valid device memory.
    pub fn send_bf16_ptr(&self, ptr: *const c_void, count: usize, peer: usize) {
        ffi::check(
            unsafe {
                ffi::ncclSend(
                    ptr,
                    count,
                    ffi::NCCL_BF16,
                    peer as i32,
                    self.comm,
                    launch_stream(),
                )
            },
            "ncclSend",
        );
    }

    /// Receive `count` BF16 elements from `peer` into `ptr`, on the null stream.
    ///
    /// # Safety
    /// `ptr` must point to at least `count` BF16 elements of valid device memory.
    pub fn recv_bf16_ptr(&self, ptr: *mut c_void, count: usize, peer: usize) {
        ffi::check(
            unsafe {
                ffi::ncclRecv(
                    ptr,
                    count,
                    ffi::NCCL_BF16,
                    peer as i32,
                    self.comm,
                    launch_stream(),
                )
            },
            "ncclRecv",
        );
    }
}

impl Drop for PpContext {
    fn drop(&mut self) {
        if !self.comm.is_null() {
            unsafe { ffi::ncclCommDestroy(self.comm) };
        }
    }
}