vyre_primitives/bitset/
zero.rs

1//! `bitset_zero` - per-word device clear (`target[w] = 0`).
2//!
3//! Resident graph pipelines use this to clear scratch/output bitsets on device
4//! instead of uploading zero-filled host buffers every iteration.
5
6use std::sync::Arc;
7
8use vyre_foundation::ir::model::expr::Ident;
9use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
10
11/// Canonical op id.
12pub const OP_ID: &str = "vyre-primitives::bitset::zero";
13
14/// Build a Program: `target[w] = 0` for `w` in `0..words`.
15#[must_use]
16pub fn bitset_zero(target: &str, words: u32) -> Program {
17    let w = Expr::InvocationId { axis: 0 };
18    Program::wrapped(
19        vec![
20            BufferDecl::storage(target, 0, BufferAccess::ReadWrite, DataType::U32)
21                .with_count(words),
22        ],
23        [256, 1, 1],
24        vec![Node::Region {
25            generator: Ident::from(OP_ID),
26            source_region: None,
27            body: Arc::new(vec![Node::if_then(
28                Expr::lt(w.clone(), Expr::u32(words)),
29                vec![Node::store(target, w, Expr::u32(0))],
30            )]),
31        }],
32    )
33}
34
35/// CPU reference. Clears every target word.
36#[cfg(any(test, feature = "cpu-parity"))]
37pub fn cpu_ref(target: &mut [u32]) {
38    target.fill(0);
39}
40
41#[cfg(feature = "inventory-registry")]
42inventory::submit! {
43    vyre_foundation::operation::OperationRegistration::primitive(
44        OP_ID,
45        || bitset_zero("target", 3),
46        Some(|| {
47            let to_bytes = |w: &[u32]| crate::wire::pack_u32_slice(w);
48            vec![vec![to_bytes(&[1, 0xDEAD_BEEF, u32::MAX])]]
49        }),
50        Some(|| {
51            let to_bytes = |w: &[u32]| crate::wire::pack_u32_slice(w);
52            vec![vec![to_bytes(&[0, 0, 0])]]
53        }),
54    )
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn cpu_ref_clears_all_words() {
63        let mut words = vec![1u32, 0xDEAD_BEEF, u32::MAX];
64        cpu_ref(&mut words);
65        assert_eq!(words, vec![0, 0, 0]);
66    }
67
68    #[test]
69    fn emitted_program_has_one_rw_target_buffer() {
70        let program = bitset_zero("target", 17);
71        assert_eq!(program.workgroup_size, [256, 1, 1]);
72        assert_eq!(program.buffers.len(), 1);
73        assert_eq!(program.buffers[0].count, 17);
74    }
75}
vyre_primitives/bitset/zero.rs

vyre_primitives/bitset/
zero.rs
