vyre_primitives/text/
byte_histogram.rs

1//! Byte histogram primitive over source bytes.
2
3use std::sync::Arc;
4
5use vyre_foundation::ir::model::expr::{GeneratorRef, Ident};
6use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
7
8/// Canonical op id for the 256-bin byte histogram primitive.
9pub const BYTE_HISTOGRAM_256_OP_ID: &str = "vyre-primitives::text::byte_histogram_256";
10
11/// Build the reusable histogram body.
12#[must_use]
13pub fn byte_histogram_256_body(input: &str, histogram: &str, count: u32) -> Vec<Node> {
14    let rounds = Expr::div(Expr::add(Expr::u32(count), Expr::u32(255)), Expr::u32(256));
15    let load_byte = |index: Expr| {
16        Expr::bitand(
17            Expr::cast(DataType::U32, Expr::load(input, index)),
18            Expr::u32(0xFF),
19        )
20    };
21
22    vec![
23        Node::let_bind("lane", Expr::InvocationId { axis: 0 }),
24        // Gate the per-lane zero-fill against the bin count. The intended dispatch is
25        // one 256-lane workgroup (lane 0..255 = the 256 bins), but whole-workgroup GPU
26        // dispatch rounds up, so a >256-lane dispatch would otherwise OOB-write the
27        // histogram (memory corruption on CUDA). Transparent for the intended dispatch;
28        // caught by the grid-overfire registry gate.
29        Node::if_then(
30            Expr::lt(Expr::var("lane"), Expr::buf_len(histogram)),
31            vec![Node::store(histogram, Expr::var("lane"), Expr::u32(0))],
32        ),
33        Node::Barrier {
34            ordering: vyre_foundation::MemoryOrdering::SeqCst,
35        },
36        Node::loop_for(
37            "round",
38            Expr::u32(0),
39            rounds,
40            vec![
41                Node::let_bind(
42                    "idx",
43                    Expr::add(
44                        Expr::mul(Expr::var("round"), Expr::u32(256)),
45                        Expr::var("lane"),
46                    ),
47                ),
48                Node::if_then(
49                    Expr::lt(Expr::var("idx"), Expr::u32(count)),
50                    vec![
51                        Node::let_bind("byte", load_byte(Expr::var("idx"))),
52                        Node::let_bind(
53                            "_prev_hist",
54                            Expr::atomic_add(histogram, Expr::var("byte"), Expr::u32(1)),
55                        ),
56                    ],
57                ),
58            ],
59        ),
60        Node::Barrier {
61            ordering: vyre_foundation::MemoryOrdering::SeqCst,
62        },
63    ]
64}
65
66/// Wrap the histogram body as a child of `parent_op_id`.
67#[must_use]
68pub fn byte_histogram_256_child(
69    parent_op_id: &str,
70    input: &str,
71    histogram: &str,
72    count: u32,
73) -> Node {
74    byte_histogram_256_child_with_source_type(parent_op_id, input, histogram, count)
75}
76
77/// Wrap the packed-`u8` histogram body as a child of `parent_op_id`.
78#[must_use]
79pub fn byte_histogram_256_u8_child(
80    parent_op_id: &str,
81    input: &str,
82    histogram: &str,
83    count: u32,
84) -> Node {
85    byte_histogram_256_child_with_source_type(parent_op_id, input, histogram, count)
86}
87
88fn byte_histogram_256_child_with_source_type(
89    parent_op_id: &str,
90    input: &str,
91    histogram: &str,
92    count: u32,
93) -> Node {
94    Node::Region {
95        generator: Ident::from(BYTE_HISTOGRAM_256_OP_ID),
96        source_region: Some(GeneratorRef {
97            name: parent_op_id.to_string(),
98        }),
99        body: Arc::new(byte_histogram_256_body(input, histogram, count)),
100    }
101}
102
103/// Standalone histogram program for primitive-level conformance.
104///
105/// This compatibility entry point expects one `DataType::U32` element per
106/// source byte and reads the low byte of each word. Use
107/// [`byte_histogram_256_u8`] when the source is packed as one byte per
108/// element.
109#[must_use]
110pub fn byte_histogram_256(input: &str, histogram: &str, count: u32) -> Program {
111    byte_histogram_256_with_source_type(input, histogram, count, DataType::U32)
112}
113
114/// Standalone histogram program over a packed `DataType::U8` source buffer.
115///
116/// It emits the same 256-bin histogram as [`byte_histogram_256`] while cutting
117/// source input bandwidth from four bytes per logical byte to one.
118#[must_use]
119pub fn byte_histogram_256_u8(input: &str, histogram: &str, count: u32) -> Program {
120    byte_histogram_256_with_source_type(input, histogram, count, DataType::U8)
121}
122
123fn byte_histogram_256_with_source_type(
124    input: &str,
125    histogram: &str,
126    count: u32,
127    source_type: DataType,
128) -> Program {
129    let input_decl = if source_type == DataType::U8 && count == 0 {
130        BufferDecl::storage(input, 0, BufferAccess::ReadOnly, source_type)
131    } else {
132        BufferDecl::storage(input, 0, BufferAccess::ReadOnly, source_type).with_count(count.max(1))
133    };
134
135    Program::wrapped(
136        vec![
137            input_decl,
138            BufferDecl::output(histogram, 1, DataType::U32)
139                .with_count(256)
140                .with_output_byte_range(0..256 * 4),
141        ],
142        [256, 1, 1],
143        vec![Node::Region {
144            generator: Ident::from(BYTE_HISTOGRAM_256_OP_ID),
145            source_region: None,
146            body: Arc::new(byte_histogram_256_body(input, histogram, count)),
147        }],
148    )
149}
150
151/// Reference oracle for [`byte_histogram_256`].
152#[must_use]
153#[cfg(any(test, feature = "cpu-parity"))]
154pub fn reference_byte_histogram(bytes: &[u8]) -> [u32; 256] {
155    let mut histogram = [0u32; 256];
156    for &byte in bytes {
157        histogram[usize::from(byte)] += 1;
158    }
159    histogram
160}
161
162#[cfg(feature = "inventory-registry")]
163inventory::submit! {
164    vyre_foundation::operation::OperationRegistration::primitive(
165        BYTE_HISTOGRAM_256_OP_ID,
166        || byte_histogram_256("bytes", "histogram", 5),
167        Some(|| {
168            vec![vec![
169                crate::wire::pack_bytes_as_u32_slice(&[b'a', b'b', b'a', 0xC3, 0xA9]),
170                vec![0; 256 * 4],
171            ]]
172        }),
173        Some(|| {
174            let mut histogram = [0u32; 256];
175            histogram[usize::from(b'a')] = 2;
176            histogram[usize::from(b'b')] = 1;
177            histogram[0xC3] = 1;
178            histogram[0xA9] = 1;
179            vec![vec![crate::wire::pack_u32_slice(&histogram)]]
180        }),
181    )
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn reference_counts_each_byte() {
190        let histogram = reference_byte_histogram(&[b'a', b'b', b'a', 0xC3, 0xA9]);
191        assert_eq!(histogram[usize::from(b'a')], 2);
192        assert_eq!(histogram[usize::from(b'b')], 1);
193        assert_eq!(histogram[0xC3], 1);
194        assert_eq!(histogram[0xA9], 1);
195    }
196
197    #[test]
198    fn packed_u8_program_declares_one_source_byte_per_element() {
199        let program = byte_histogram_256_u8("bytes", "histogram", 513);
200        let source = program
201            .buffers()
202            .iter()
203            .find(|buffer| buffer.name() == "bytes")
204            .expect("Fix: packed-u8 byte histogram source buffer must be declared");
205        let histogram = program
206            .buffers()
207            .iter()
208            .find(|buffer| buffer.name() == "histogram")
209            .expect("Fix: byte histogram output buffer must be declared");
210
211        assert_eq!(source.element(), DataType::U8);
212        assert_eq!(source.count(), 513);
213        assert_eq!(histogram.element(), DataType::U32);
214        assert_eq!(histogram.count(), 256);
215        assert_eq!(program.workgroup_size(), [256, 1, 1]);
216    }
217
218    #[test]
219    fn masks_high_bit_source_and_records_no_interpreter_oob() {
220        use vyre_reference::value::Value;
221        // A U32 source element > 255 must fold into the 256-bin histogram via the
222        // `& 0xFF` bin mask, never atomic-add past the last bin. Assert ZERO
223        // interpreter OOB accesses (the mask keeps the atomic in bounds, not the
224        // interpreter's silent drop) and that the low byte's bin is incremented.
225        // Removing the mask would OOB the atomic scatter (memory corruption on CUDA)
226        // and this test would see report.total() > 0.
227        let program = byte_histogram_256("bytes", "histogram", 1);
228        let (outputs, report) = vyre_reference::reference_eval_oob_report(
229            &program,
230            &[
231                Value::from(crate::wire::pack_u32_slice(&[0x0141])), // > 255, low byte 0x41
232                Value::from(vec![0u8; 256 * 4]),
233            ],
234        )
235        .expect("Fix: byte_histogram_256 must reference-evaluate a high-bit source element");
236        assert_eq!(
237            report.total(),
238            0,
239            "Fix: masked bin index must stay in bounds without relying on interpreter OOB masking"
240        );
241        let histogram = crate::wire::decode_u32_le_bytes_all(&outputs[0].to_bytes());
242        assert_eq!(
243            histogram[0x41], 1,
244            "Fix: 0x0141 must count into bin 0x41 via the `& 0xFF` mask"
245        );
246        assert_eq!(
247            histogram.iter().sum::<u32>(),
248            1,
249            "exactly one byte counted, no stray OOB write elsewhere"
250        );
251    }
252}
vyre_primitives/text/byte_histogram.rs

vyre_primitives/text/
byte_histogram.rs
