Skip to main content

vyre_spec/
collective_op.rs

1//! Frozen collective-communication operation contracts.
2//!
3//! These types are IR-level contracts only. Concrete transport bindings such
4//! as NCCL, UCX, SHARP, or MPI live in backend crates.
5// TAG RESERVATIONS: Sum=0x01, Min=0x02, Max=0x03, BitAnd=0x04,
6// BitOr=0x05, BitXor=0x06, 0x07..=0x7F reserved.
7
8/// Reduction operator used by distributed collective nodes.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize)]
10#[non_exhaustive]
11pub enum CollectiveOp {
12    /// Sum reduction.
13    Sum,
14    /// Minimum reduction.
15    Min,
16    /// Maximum reduction.
17    Max,
18    /// Bitwise AND reduction.
19    BitAnd,
20    /// Bitwise OR reduction.
21    BitOr,
22    /// Bitwise XOR reduction.
23    BitXor,
24}
25
26impl CollectiveOp {
27    /// Frozen builtin wire tag for this collective operator.
28    #[must_use]
29    pub const fn builtin_wire_tag(self) -> u8 {
30        match self {
31            Self::Sum => 0x01,
32            Self::Min => 0x02,
33            Self::Max => 0x03,
34            Self::BitAnd => 0x04,
35            Self::BitOr => 0x05,
36            Self::BitXor => 0x06,
37        }
38    }
39
40    /// Decode a frozen builtin wire tag.
41    ///
42    /// # Errors
43    ///
44    /// Returns an actionable diagnostic when `tag` is not assigned.
45    pub fn from_wire_tag(tag: u8) -> Result<Self, String> {
46        match tag {
47            0x01 => Ok(Self::Sum),
48            0x02 => Ok(Self::Min),
49            0x03 => Ok(Self::Max),
50            0x04 => Ok(Self::BitAnd),
51            0x05 => Ok(Self::BitOr),
52            0x06 => Ok(Self::BitXor),
53            value => Err(format!(
54                "Fix: unknown CollectiveOp tag {value}; use a Program serializer compatible with this vyre version."
55            )),
56        }
57    }
58}
59
60/// Opaque communicator/group handle carried by collective nodes.
61///
62/// `0` is the process/world group by convention. Other ids are backend-owned
63/// handles resolved by the runtime communicator registry.
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize)]
65pub struct CommGroup(pub u32);
66
67impl CommGroup {
68    /// Default world communicator group.
69    pub const WORLD: Self = Self(0);
70
71    /// Return the stable group id.
72    #[must_use]
73    pub const fn as_u32(self) -> u32 {
74        self.0
75    }
76}