vyre_spec/
subgroup_reduce_op.rs

1//! Frozen subgroup (warp) reduction-operation contracts.
2//!
3//! Warp-scoped collective reductions across the active subgroup lanes. Distinct
4//! from [`crate::collective_op::CollectiveOp`], which is distributed-scoped
5//! (NCCL/MPI) and lacks the multiplicative reduction the warp ISA exposes.
6//! This is the complete set the hardware/`naga::SubgroupOperation` surface
7//! supports.
8// TAG RESERVATIONS: Add=0x01, Mul=0x02, Min=0x03, Max=0x04, And=0x05,
9// Or=0x06, Xor=0x07, 0x08..=0x7F reserved.
10
11/// Reduction operator applied across the active subgroup lanes.
12///
13/// Stability: matches must include a fallback arm so the contract can grow
14/// without breaking `SemVer`.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize)]
16#[non_exhaustive]
17pub enum SubgroupReduceOp {
18    /// Sum across the subgroup.
19    Add,
20    /// Product across the subgroup.
21    Mul,
22    /// Minimum across the subgroup.
23    Min,
24    /// Maximum across the subgroup.
25    Max,
26    /// Bitwise AND across the subgroup.
27    And,
28    /// Bitwise OR across the subgroup.
29    Or,
30    /// Bitwise XOR across the subgroup.
31    Xor,
32}
33
34impl SubgroupReduceOp {
35    /// Every builtin reduction operator, in wire-tag order.
36    pub const ALL: [Self; 7] = [
37        Self::Add,
38        Self::Mul,
39        Self::Min,
40        Self::Max,
41        Self::And,
42        Self::Or,
43        Self::Xor,
44    ];
45
46    /// Frozen builtin wire tag for this reduction operator.
47    #[must_use]
48    pub const fn builtin_wire_tag(self) -> u8 {
49        match self {
50            Self::Add => 0x01,
51            Self::Mul => 0x02,
52            Self::Min => 0x03,
53            Self::Max => 0x04,
54            Self::And => 0x05,
55            Self::Or => 0x06,
56            Self::Xor => 0x07,
57        }
58    }
59
60    /// Decode a frozen builtin wire tag.
61    ///
62    /// # Errors
63    ///
64    /// Returns an actionable diagnostic when `tag` is not assigned.
65    pub fn from_wire_tag(tag: u8) -> Result<Self, String> {
66        match tag {
67            0x01 => Ok(Self::Add),
68            0x02 => Ok(Self::Mul),
69            0x03 => Ok(Self::Min),
70            0x04 => Ok(Self::Max),
71            0x05 => Ok(Self::And),
72            0x06 => Ok(Self::Or),
73            0x07 => Ok(Self::Xor),
74            value => Err(format!(
75                "Fix: unknown SubgroupReduceOp tag {value}; use a Program serializer compatible with this vyre version."
76            )),
77        }
78    }
79
80    /// Lower-case spelling used in op-id strings and diagnostics.
81    #[must_use]
82    pub const fn as_str(self) -> &'static str {
83        match self {
84            Self::Add => "add",
85            Self::Mul => "mul",
86            Self::Min => "min",
87            Self::Max => "max",
88            Self::And => "and",
89            Self::Or => "or",
90            Self::Xor => "xor",
91        }
92    }
93
94    /// True when this operator is bitwise (integer-only): `And`/`Or`/`Xor`.
95    ///
96    /// Bitwise reductions reject floating-point operands during type checking.
97    #[must_use]
98    pub const fn is_bitwise(self) -> bool {
99        matches!(self, Self::And | Self::Or | Self::Xor)
100    }
101
102    /// Canonical integer reduction of `lanes` under this operator.
103    ///
104    /// This is the single source of truth for the operator's semantics:
105    /// the CPU reference oracle, constant folding, and any host-side
106    /// evaluation route through it so they cannot drift. Wrapping arithmetic
107    /// matches the GPU ISA (`redux.sync` / `subgroupAdd` wrap on overflow).
108    /// Neutral elements: `Add`=0, `Mul`=1, `Min`=`u32::MAX`, `Max`=0,
109    /// `And`=`u32::MAX`, `Or`=0, `Xor`=0.
110    #[must_use]
111    pub fn reduce_u32(self, lanes: impl IntoIterator<Item = u32>) -> u32 {
112        let lanes = lanes.into_iter();
113        match self {
114            Self::Add => lanes.fold(0u32, u32::wrapping_add),
115            Self::Mul => lanes.fold(1u32, u32::wrapping_mul),
116            Self::Min => lanes.fold(u32::MAX, u32::min),
117            Self::Max => lanes.fold(0u32, u32::max),
118            Self::And => lanes.fold(u32::MAX, |acc, lane| acc & lane),
119            Self::Or => lanes.fold(0u32, |acc, lane| acc | lane),
120            Self::Xor => lanes.fold(0u32, |acc, lane| acc ^ lane),
121        }
122    }
123
124    /// Floating-point identity (neutral) element for this operator, or `None`
125    /// for the bitwise operators (`And`/`Or`/`Xor`), which are integer-only.
126    ///
127    /// Callers fold with [`Self::combine_f32`] starting from this identity so
128    /// they can apply their own per-step canonicalization (e.g. NaN folding).
129    #[must_use]
130    pub fn f32_identity(self) -> Option<f32> {
131        match self {
132            Self::Add => Some(0.0),
133            Self::Mul => Some(1.0),
134            Self::Min => Some(f32::INFINITY),
135            Self::Max => Some(f32::NEG_INFINITY),
136            Self::And | Self::Or | Self::Xor => None,
137        }
138    }
139
140    /// Combine one f32 lane into a running accumulator, or `None` for the
141    /// bitwise operators (`And`/`Or`/`Xor`), which are integer-only.
142    #[must_use]
143    pub fn combine_f32(self, acc: f32, lane: f32) -> Option<f32> {
144        match self {
145            Self::Add => Some(acc + lane),
146            Self::Mul => Some(acc * lane),
147            Self::Min => Some(acc.min(lane)),
148            Self::Max => Some(acc.max(lane)),
149            Self::And | Self::Or | Self::Xor => None,
150        }
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn wire_tag_roundtrips_every_op() {
160        for op in SubgroupReduceOp::ALL {
161            let tag = op.builtin_wire_tag();
162            assert_eq!(
163                SubgroupReduceOp::from_wire_tag(tag).unwrap(),
164                op,
165                "Fix: SubgroupReduceOp wire tag {tag:#04x} must round-trip"
166            );
167        }
168    }
169
170    #[test]
171    fn wire_tags_are_distinct_and_dense() {
172        let tags: Vec<u8> = SubgroupReduceOp::ALL
173            .iter()
174            .map(|op| op.builtin_wire_tag())
175            .collect();
176        assert_eq!(tags, vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
177    }
178
179    #[test]
180    fn unknown_tag_is_rejected_loudly() {
181        let err = SubgroupReduceOp::from_wire_tag(0x42).unwrap_err();
182        assert!(err.contains("unknown SubgroupReduceOp tag 66"), "{err}");
183    }
184
185    #[test]
186    fn reduce_u32_computes_exact_values() {
187        let lanes = [3u32, 1, 4, 1, 5];
188        assert_eq!(SubgroupReduceOp::Add.reduce_u32(lanes), 14);
189        assert_eq!(SubgroupReduceOp::Mul.reduce_u32(lanes), 60);
190        assert_eq!(SubgroupReduceOp::Min.reduce_u32(lanes), 1);
191        assert_eq!(SubgroupReduceOp::Max.reduce_u32(lanes), 5);
192        assert_eq!(SubgroupReduceOp::And.reduce_u32([0b1100u32, 0b1010]), 0b1000);
193        assert_eq!(SubgroupReduceOp::Or.reduce_u32([0b1100u32, 0b1010]), 0b1110);
194        assert_eq!(SubgroupReduceOp::Xor.reduce_u32([0b1100u32, 0b1010]), 0b0110);
195    }
196
197    #[test]
198    fn reduce_u32_add_wraps_like_the_isa() {
199        assert_eq!(SubgroupReduceOp::Add.reduce_u32([u32::MAX, 1]), 0);
200        assert_eq!(SubgroupReduceOp::Mul.reduce_u32([u32::MAX, 2]), u32::MAX - 1);
201    }
202
203    #[test]
204    fn reduce_u32_empty_yields_neutral() {
205        assert_eq!(SubgroupReduceOp::Add.reduce_u32([]), 0);
206        assert_eq!(SubgroupReduceOp::Mul.reduce_u32([]), 1);
207        assert_eq!(SubgroupReduceOp::Min.reduce_u32([]), u32::MAX);
208        assert_eq!(SubgroupReduceOp::Max.reduce_u32([]), 0);
209        assert_eq!(SubgroupReduceOp::And.reduce_u32([]), u32::MAX);
210    }
211
212    #[test]
213    fn f32_reduce_helpers_match_op_and_reject_bitwise() {
214        let lanes = [3.0f32, 1.0, 4.0];
215        let fold = |op: SubgroupReduceOp| {
216            op.f32_identity()
217                .map(|id| lanes.iter().fold(id, |acc, &l| op.combine_f32(acc, l).unwrap()))
218        };
219        assert_eq!(fold(SubgroupReduceOp::Add), Some(8.0));
220        assert_eq!(fold(SubgroupReduceOp::Mul), Some(12.0));
221        assert_eq!(fold(SubgroupReduceOp::Min), Some(1.0));
222        assert_eq!(fold(SubgroupReduceOp::Max), Some(4.0));
223        // Bitwise ops have no f32 identity.
224        assert_eq!(SubgroupReduceOp::And.f32_identity(), None);
225        assert_eq!(SubgroupReduceOp::Or.combine_f32(1.0, 2.0), None);
226    }
227
228    #[test]
229    fn only_bitwise_ops_are_bitwise() {
230        assert!(SubgroupReduceOp::And.is_bitwise());
231        assert!(SubgroupReduceOp::Or.is_bitwise());
232        assert!(SubgroupReduceOp::Xor.is_bitwise());
233        assert!(!SubgroupReduceOp::Add.is_bitwise());
234        assert!(!SubgroupReduceOp::Mul.is_bitwise());
235        assert!(!SubgroupReduceOp::Min.is_bitwise());
236        assert!(!SubgroupReduceOp::Max.is_bitwise());
237    }
238}
vyre_spec/subgroup_reduce_op.rs

vyre_spec/
subgroup_reduce_op.rs
