//! 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::device;
use xserv_cuda::GpuBuffer;

pub use ffi::NcclUniqueId as UniqueId;

/// The CUDA "null" (default) stream. The model's kernels and cuBLAS calls run
/// on it, so issuing NCCL on the same stream keeps AllReduce correctly ordered
/// after the producing matmul and before the consuming kernel — no extra sync.
const NULL_STREAM: xserv_cuda::ffi::CudaStream = std::ptr::null_mut();

/// 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,
                    NULL_STREAM,
                )
            },
            "ncclAllReduce",
        );
    }
}

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