pmetal_distributed/

ring.rs

use crate::{
    DistributedBackend,
    config::DistributedConfig,
    error::DistributedError,
    transport::{TcpTransport, TransportReceiver, TransportSender},
};
use anyhow::Result;
use async_trait::async_trait;
use tokio::sync::Mutex;
use zerocopy::{FromBytes, IntoBytes};

pub struct RingBackend {
    rank: usize,
    world_size: usize,
    sender: Mutex<TransportSender>,
    receiver: Mutex<TransportReceiver>,
}

impl RingBackend {
    pub async fn new(config: DistributedConfig) -> Result<Self> {
        config.validate()?;
        let (sender, receiver) = TcpTransport::connect(&config).await?;
        Ok(Self {
            rank: config.rank,
            world_size: config.nodes.len(),
            sender: Mutex::new(sender),
            receiver: Mutex::new(receiver),
        })
    }
}

#[async_trait]
impl DistributedBackend for RingBackend {
    fn rank(&self) -> usize {
        self.rank
    }

    fn world_size(&self) -> usize {
        self.world_size
    }

    async fn all_reduce(&self, buffer: &mut [u8]) -> Result<()> {
        // Validate buffer alignment and size for f32 operations
        if !buffer.len().is_multiple_of(4) {
            return Err(DistributedError::Protocol(format!(
                "Buffer length {} is not a multiple of 4 (f32 size)",
                buffer.len()
            ))
            .into());
        }

        if !(buffer.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
            return Err(DistributedError::Protocol(
                "Buffer is not properly aligned for f32 operations".to_string(),
            )
            .into());
        }

        let floats: &mut [f32] = <[f32]>::mut_from_bytes(buffer)
            .map_err(|e| DistributedError::Protocol(format!("Buffer cast failed: {e}")))?;
        let len = floats.len();

        let chunk_size = len / self.world_size;
        let remainder = len % self.world_size;

        let get_chunk_range = |idx: usize| -> (usize, usize) {
            let start = idx * chunk_size + idx.min(remainder);
            let end = start + chunk_size + (if idx < remainder { 1 } else { 0 });
            (start, end)
        };

        // Lock both transport halves
        let mut sender = self.sender.lock().await;
        let mut receiver = self.receiver.lock().await;

        // 1. Scatter-Reduce
        let mut send_chunk_idx = self.rank;
        let mut recv_chunk_idx = (self.rank + self.world_size - 1) % self.world_size;

        // Recv buffer
        let max_chunk_size = chunk_size + 1;
        let mut recv_buf = vec![0u8; max_chunk_size * 4];

        for _ in 0..self.world_size - 1 {
            let (s_start, s_end) = get_chunk_range(send_chunk_idx);
            let (r_start, r_end) = get_chunk_range(recv_chunk_idx);

            // Prepare send data
            // We need to copy to a temp buffer because floats is borrowed by recv logic?
            // Actually, we can just send. send takes &[u8].
            // BUT: We need to use &mut sender and &mut receiver concurrently.
            // Since we locked them, we own them. We can pass them to concurrent futures.
            // But MutexGuard is not Send if we hold it across await? No, it is.
            // Wait, we can't borrow `sender` and `receiver` into `try_join` if they are MutexGuards?
            // Yes we can, they are distinct.

            let send_buf = floats[s_start..s_end].as_bytes().to_vec();

            let recv_bytes_len = (r_end - r_start) * 4;
            let recv_slice = &mut recv_buf[..recv_bytes_len];

            // Concurrent Send and Recv
            let send_fut = sender.send(&send_buf);
            let recv_fut = receiver.recv(recv_slice);

            tokio::try_join!(send_fut, recv_fut)?;

            // Reduce
            let recv_floats =
                <[f32]>::ref_from_bytes(recv_slice).expect("recv buffer aligned for f32");

            for i in 0..recv_floats.len() {
                floats[r_start + i] += recv_floats[i];
            }

            send_chunk_idx = recv_chunk_idx;
            recv_chunk_idx = (recv_chunk_idx + self.world_size - 1) % self.world_size;
        }

        // 2. All-Gather
        send_chunk_idx = (self.rank + 1) % self.world_size;
        recv_chunk_idx = self.rank;

        for _ in 0..self.world_size - 1 {
            let (s_start, s_end) = get_chunk_range(send_chunk_idx);
            let (r_start, r_end) = get_chunk_range(recv_chunk_idx);

            let send_buf = floats[s_start..s_end].as_bytes().to_vec();

            let recv_bytes_len = (r_end - r_start) * 4;
            let recv_slice = &mut recv_buf[..recv_bytes_len];

            let send_fut = sender.send(&send_buf);
            let recv_fut = receiver.recv(recv_slice);

            tokio::try_join!(send_fut, recv_fut)?;

            // Copy (Gather)
            let recv_floats =
                <[f32]>::ref_from_bytes(recv_slice).expect("recv buffer aligned for f32");
            floats[r_start..r_end].copy_from_slice(recv_floats);

            send_chunk_idx = recv_chunk_idx;
            recv_chunk_idx = (recv_chunk_idx + self.world_size - 1) % self.world_size;
        }

        Ok(())
    }

    async fn barrier(&self) -> Result<()> {
        let mut buf = [0u8; 4]; // Minimum 1 float
        self.all_reduce(&mut buf).await
    }
}

#[cfg(kani)]
mod verification {
    use super::*;

    #[kani::proof]
    #[kani::unwind(17)] // Sufficient for world_size up to 16 for testing
    fn verify_get_chunk_range() {
        let len: usize = kani::any();
        let world_size: usize = kani::any();

        // Preconditions
        kani::assume(world_size > 0 && world_size <= 16);
        kani::assume(len >= world_size && len < 1024);

        let chunk_size = len / world_size;
        let remainder = len % world_size;

        let get_chunk_range = |idx: usize| -> (usize, usize) {
            let start = idx * chunk_size + idx.min(remainder);
            let end = start + chunk_size + (if idx < remainder { 1 } else { 0 });
            (start, end)
        };

        let mut total_elements = 0;
        let mut last_end = 0;

        for i in 0..world_size {
            let (start, end) = get_chunk_range(i);

            // Chunks must be valid ranges
            assert!(start <= end);
            // Chunks must be contiguous
            assert!(start == last_end);
            // Chunks must be within bounds
            assert!(end <= len);

            total_elements += end - start;
            last_end = end;
        }

        // Total elements must match original length
        assert!(total_elements == len);
        assert!(last_end == len);
    }
}