nexar 0.1.2

Distributed runtime with QUIC transport, stream-multiplexed messaging, and built-in collectives
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use crate::client::NexarClient;
use crate::collective::helpers::{CollectiveTag, collective_recv, collective_send};
use crate::error::{NexarError, Result};
use crate::reduce::reduce_slice;
use crate::types::{DataType, Rank, ReduceOp};

/// Tree reduce: reduce data from all ranks to a single root rank.
///
/// Uses a binomial tree algorithm with O(log N) steps. Handles non-power-of-2
/// world sizes by first reducing excess ranks into the lower power-of-2 set.
///
/// # Safety
/// `ptr` must be valid for at least `count * dtype.size_in_bytes()` bytes.
pub(crate) async unsafe fn tree_reduce(
    client: &NexarClient,
    ptr: u64,
    count: usize,
    dtype: DataType,
    op: ReduceOp,
    root: Rank,
    tag: CollectiveTag,
) -> Result<()> {
    let world = client.world_size() as usize;
    let rank = client.rank() as usize;
    let root = root as usize;

    if world <= 1 {
        return Ok(());
    }

    let elem_size = dtype.size_in_bytes();
    let total_bytes = count * elem_size;

    let data = unsafe { client.adapter().stage_for_send(ptr, total_bytes)? };
    let mut buf = data;

    let vrank = (rank + world - root) % world;

    let p2 = if world.is_power_of_two() {
        world
    } else {
        world.next_power_of_two() >> 1
    };
    let excess = world - p2;

    let mut participating = true;
    if vrank < excess {
        let partner_vrank = vrank + p2;
        let partner_real = (partner_vrank + root) % world;
        let received = collective_recv(client, partner_real as u32, "reduce", tag).await?;
        if received.len() != total_bytes {
            return Err(NexarError::BufferSizeMismatch {
                expected: total_bytes,
                actual: received.len(),
            });
        }
        reduce_slice(&mut buf, &received, count, dtype, op)?;
    } else if vrank >= p2 {
        let partner_vrank = vrank - p2;
        let partner_real = (partner_vrank + root) % world;
        collective_send(client, partner_real as u32, &buf, "reduce", tag).await?;
        participating = false;
    }

    if participating {
        let adjusted_vrank = vrank;
        let log2 = p2.trailing_zeros() as usize;

        for round in 0..log2 {
            let mask = 1 << round;
            if adjusted_vrank & mask != 0 {
                let partner_vrank = adjusted_vrank ^ mask;
                let partner_real = (partner_vrank + root) % world;
                collective_send(client, partner_real as u32, &buf, "reduce", tag).await?;
                break;
            } else {
                let partner_vrank = adjusted_vrank ^ mask;
                if partner_vrank < p2 {
                    let partner_real = (partner_vrank + root) % world;
                    let received =
                        collective_recv(client, partner_real as u32, "reduce", tag).await?;
                    if received.len() != total_bytes {
                        return Err(NexarError::BufferSizeMismatch {
                            expected: total_bytes,
                            actual: received.len(),
                        });
                    }
                    reduce_slice(&mut buf, &received, count, dtype, op)?;
                }
            }
        }
    }

    if rank == root {
        unsafe { client.adapter().receive_to_device(&buf, ptr)? };
    }

    Ok(())
}