nir_rs/io/
wire.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Vocabulary of the NIR HDF5 wire format.
4//!
5//! Everything here is pure Rust and always compiled, with or without the
6//! `hdf5` feature: it is the shared agreement between the reader and the
7//! writer, and it is useful on its own to tooling that inspects `.nir` files
8//! without going through [`crate::io::read`].
9//!
10//! The structural layout these names describe:
11//!
12//! ```text
13//! /version                 scalar string, e.g. "0.2.0"
14//! /node                    group — the root NIRGraph
15//!   type                   scalar string "NIRGraph"
16//!   edges                  (E, 2) string dataset
17//!   nodes/<name>/          one group per node
18//!     type                 scalar string, e.g. "LIF"
19//!     <field>              one dataset per wire field
20//!     metadata/            group; omitted when empty
21//!   metadata/              group; omitted when empty
22//! ```
23
24use crate::error::{NirError, Result};
25use crate::nodes::Padding;
26
27/// Dataset holding the NIR format version at the root of the file.
28pub const KEY_VERSION: &str = "version";
29/// Group holding the root graph.
30pub const KEY_NODE: &str = "node";
31/// Group holding a graph's named nodes.
32pub const KEY_NODES: &str = "nodes";
33/// Dataset holding a graph's `(E, 2)` edge list.
34pub const KEY_EDGES: &str = "edges";
35/// Group holding free-form metadata; omitted from the file when empty.
36pub const KEY_METADATA: &str = "metadata";
37/// Dataset holding a node's wire `type` string.
38pub const KEY_TYPE: &str = "type";
39
40/// Every node `type` string that can appear in a `.nir` file.
41///
42/// These match [`crate::NirNode::type_name`] exactly — a unit test asserts the
43/// two stay in step, so this list cannot drift from the enum.
44///
45/// Upstream's `Identity` node is deliberately absent: it is not in the Python
46/// serializer registry (`nir.ir.__all_ir`) and so never reaches the wire.
47pub const WIRE_TYPES: [&str; 19] = [
48    "Input",
49    "Output",
50    "Affine",
51    "Linear",
52    "Scale",
53    "Conv1d",
54    "Conv2d",
55    "CubaLI",
56    "CubaLIF",
57    "Delay",
58    "Flatten",
59    "I",
60    "IF",
61    "LI",
62    "LIF",
63    "SumPool2d",
64    "AvgPool2d",
65    "Threshold",
66    "NIRGraph",
67];
68
69/// Whether `name` is a NIR wire node type this crate understands.
70#[must_use]
71pub fn is_wire_type(name: &str) -> bool {
72    WIRE_TYPES.contains(&name)
73}
74
75/// Wire spelling of the symbolic padding modes.
76///
77/// Returns [`None`] for [`Padding::Explicit`], which is written as an integer
78/// dataset rather than a string.
79#[must_use]
80pub fn padding_as_wire_str(padding: &Padding) -> Option<&'static str> {
81    match padding {
82        Padding::Same => Some("same"),
83        Padding::Valid => Some("valid"),
84        Padding::Explicit(_) => None,
85    }
86}
87
88/// Parse a symbolic padding mode from its wire spelling.
89///
90/// # Errors
91///
92/// Returns [`NirError::InvalidGraph`] for anything other than `"same"` or
93/// `"valid"` — the only two strings upstream `Conv1d` / `Conv2d` accept.
94pub fn padding_from_wire_str(s: &str) -> Result<Padding> {
95    match s {
96        "same" => Ok(Padding::Same),
97        "valid" => Ok(Padding::Valid),
98        other => Err(NirError::InvalidGraph(format!(
99            "padding must be \"same\", \"valid\", or integer extents, not {other:?}"
100        ))),
101    }
102}
103
104/// Check that `name` can be used as an HDF5 link name.
105///
106/// Applies to every caller-supplied string that becomes a link in the file:
107/// graph node names and metadata keys. HDF5 splits paths on `/`, so a name
108/// containing one would silently nest and change the graph on the next read;
109/// `.` and `..` are reserved path components; link names are C strings and so
110/// cannot carry an embedded NUL; and an empty name has no valid encoding.
111///
112/// Callers should run this **before** creating the destination file. Every
113/// rejected name otherwise fails at link-creation time, by which point an
114/// existing file at that path has already been truncated.
115///
116/// `kind` names what is being checked (`"node name"`, `"metadata key"`) and
117/// appears in the error.
118///
119/// # Errors
120///
121/// Returns [`NirError::InvalidGraph`] describing the offending name.
122pub fn check_link_name(kind: &str, name: &str) -> Result<()> {
123    if name.is_empty() {
124        return Err(NirError::InvalidGraph(format!("{kind} must not be empty")));
125    }
126    if name.contains('/') {
127        return Err(NirError::InvalidGraph(format!(
128            "{kind} {name:?} must not contain '/' (HDF5 path separator)"
129        )));
130    }
131    if name.contains('\0') {
132        return Err(NirError::InvalidGraph(format!(
133            "{kind} {name:?} must not contain a NUL byte (HDF5 link names are C strings)"
134        )));
135    }
136    if name == "." || name == ".." {
137        return Err(NirError::InvalidGraph(format!(
138            "{kind} {name:?} is a reserved HDF5 path component"
139        )));
140    }
141    Ok(())
142}
143
144/// Check that `value` can be stored in an HDF5 string dataset.
145///
146/// Link names and string payloads share the C-string constraint: an embedded
147/// NUL cannot be encoded. Call this **before** creating the destination file,
148/// alongside [`check_link_name`].
149///
150/// # Errors
151///
152/// Returns [`NirError::InvalidGraph`] when `value` contains a NUL byte.
153pub fn check_hdf5_string(kind: &str, value: &str) -> Result<()> {
154    if value.contains('\0') {
155        return Err(NirError::InvalidGraph(format!(
156            "{kind} {value:?} must not contain a NUL byte (HDF5 strings are C strings)"
157        )));
158    }
159    Ok(())
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use crate::graph::NirGraph;
166    use crate::nodes::{
167        Affine, AvgPool2d, Conv1d, Conv2d, CubaLi, CubaLif, Delay, Flatten, I, If, Input, Li, Lif,
168        Linear, NirNode, Output, Scale, SumPool2d, Threshold,
169    };
170    use crate::types::Tensor;
171
172    /// A length-2 `f64` vector, the shape every neuron parameter below uses.
173    fn v() -> Tensor {
174        Tensor::from_f64([2], vec![1.0, 1.0]).unwrap()
175    }
176
177    /// A length-2 `i64` vector, for pooling windows.
178    fn pool() -> Tensor {
179        Tensor::from_i64([2], vec![2, 2]).unwrap()
180    }
181
182    fn port_and_linear_nodes() -> Vec<NirNode> {
183        let weight = || Tensor::from_f32(vec![2, 2], vec![1., 0., 0., 1.]).unwrap();
184        vec![
185            NirNode::Input(Input {
186                shape: vec![2],
187                metadata: Default::default(),
188            }),
189            NirNode::Output(Output {
190                shape: vec![2],
191                metadata: Default::default(),
192            }),
193            NirNode::Affine(Affine {
194                weight: weight(),
195                bias: Tensor::from_f32([2], vec![0., 0.]).unwrap(),
196                metadata: Default::default(),
197            }),
198            NirNode::Linear(Linear {
199                weight: weight(),
200                metadata: Default::default(),
201            }),
202            NirNode::Scale(Scale {
203                scale: v(),
204                metadata: Default::default(),
205            }),
206        ]
207    }
208
209    fn conv_nodes() -> Vec<NirNode> {
210        vec![
211            NirNode::Conv1d(Conv1d {
212                weight: Tensor::from_f32(vec![1, 1, 3], vec![1., 0., -1.]).unwrap(),
213                stride: vec![1],
214                padding: Padding::single(0),
215                dilation: vec![1],
216                groups: 1,
217                bias: Tensor::from_f32([1], vec![0.]).unwrap(),
218                input_shape: Some(10),
219                metadata: Default::default(),
220            }),
221            NirNode::Conv2d(Conv2d {
222                weight: Tensor::from_f32(vec![1, 1, 2, 2], vec![0.; 4]).unwrap(),
223                stride: vec![1, 1],
224                padding: Padding::Same,
225                dilation: vec![1, 1],
226                groups: 1,
227                bias: Tensor::from_f32([1], vec![0.]).unwrap(),
228                input_shape: Some(vec![8, 8]),
229                metadata: Default::default(),
230            }),
231        ]
232    }
233
234    fn cuba_nodes() -> Vec<NirNode> {
235        vec![
236            NirNode::CubaLi(CubaLi {
237                tau_syn: v(),
238                tau_mem: v(),
239                r: v(),
240                v_leak: v(),
241                w_in: None,
242                metadata: Default::default(),
243            }),
244            NirNode::CubaLif(CubaLif {
245                tau_syn: v(),
246                tau_mem: v(),
247                r: v(),
248                v_leak: v(),
249                v_threshold: v(),
250                v_reset: None,
251                w_in: None,
252                metadata: Default::default(),
253            }),
254        ]
255    }
256
257    fn neuron_nodes() -> Vec<NirNode> {
258        vec![
259            NirNode::I(I {
260                r: v(),
261                metadata: Default::default(),
262            }),
263            NirNode::If(If {
264                r: v(),
265                v_threshold: v(),
266                v_reset: None,
267                metadata: Default::default(),
268            }),
269            NirNode::Li(Li {
270                tau: v(),
271                r: v(),
272                v_leak: v(),
273                metadata: Default::default(),
274            }),
275            NirNode::Lif(Lif {
276                tau: v(),
277                r: v(),
278                v_leak: v(),
279                v_threshold: v(),
280                v_reset: None,
281                metadata: Default::default(),
282            }),
283        ]
284    }
285
286    fn pool_nodes() -> Vec<NirNode> {
287        let no_pad = || Tensor::from_i64([2], vec![0, 0]).unwrap();
288        vec![
289            NirNode::SumPool2d(SumPool2d {
290                kernel_size: pool(),
291                stride: pool(),
292                padding: no_pad(),
293                metadata: Default::default(),
294            }),
295            NirNode::AvgPool2d(AvgPool2d {
296                kernel_size: pool(),
297                stride: pool(),
298                padding: no_pad(),
299                metadata: Default::default(),
300            }),
301        ]
302    }
303
304    /// Stable index of a `NirNode` variant, in [`WIRE_TYPES`] order.
305    ///
306    /// Exhaustive, so a new variant is a compile error here — that is the
307    /// notification, and it is worth being precise about its limits. Nothing
308    /// in the test reads this match's arm count, so adding an arm with the
309    /// next index and updating neither `VARIANT_COUNT`, `WIRE_TYPES` nor
310    /// `one_of_each` still leaves the test green: the three lists agree with
311    /// each other at the old length, and no sample ever exercises the new
312    /// index. Closing that needs the variants, their wire strings and their
313    /// sample values generated from one definition — a macro owning `NirNode`
314    /// itself, since stable Rust cannot enumerate an enum's variants.
315    ///
316    /// What the index does buy over a bare coverage match: `WIRE_TYPES` and
317    /// `one_of_each` are checked position-by-position rather than as two
318    /// sequences that happen to compare equal, so a sample in the wrong slot,
319    /// a duplicate, or a `type_name` that disagrees with `WIRE_TYPES` at that
320    /// index all fail.
321    fn variant_index(node: &NirNode) -> usize {
322        match node {
323            NirNode::Input(_) => 0,
324            NirNode::Output(_) => 1,
325            NirNode::Affine(_) => 2,
326            NirNode::Linear(_) => 3,
327            NirNode::Scale(_) => 4,
328            NirNode::Conv1d(_) => 5,
329            NirNode::Conv2d(_) => 6,
330            NirNode::CubaLi(_) => 7,
331            NirNode::CubaLif(_) => 8,
332            NirNode::Delay(_) => 9,
333            NirNode::Flatten(_) => 10,
334            NirNode::I(_) => 11,
335            NirNode::If(_) => 12,
336            NirNode::Li(_) => 13,
337            NirNode::Lif(_) => 14,
338            NirNode::SumPool2d(_) => 15,
339            NirNode::AvgPool2d(_) => 16,
340            NirNode::Threshold(_) => 17,
341            NirNode::Graph(_) => 18,
342        }
343    }
344
345    /// One value per `NirNode` variant, in [`WIRE_TYPES`] order.
346    ///
347    /// Assembled from the per-family helpers above so the constructors live in
348    /// one place; [`variant_index`] supplies the compile-time coverage.
349    fn one_of_each() -> Vec<NirNode> {
350        let mut nodes = port_and_linear_nodes();
351        nodes.extend(conv_nodes());
352        nodes.extend(cuba_nodes());
353        nodes.push(NirNode::Delay(Delay {
354            delay: v(),
355            metadata: Default::default(),
356        }));
357        nodes.push(NirNode::Flatten(Flatten {
358            start_dim: 1,
359            end_dim: -1,
360            input_type: None,
361            metadata: Default::default(),
362        }));
363        nodes.extend(neuron_nodes());
364        nodes.extend(pool_nodes());
365        nodes.push(NirNode::Threshold(Threshold {
366            threshold: v(),
367            metadata: Default::default(),
368        }));
369        nodes.push(NirNode::Graph(Box::new(NirGraph::new())));
370        nodes
371    }
372
373    #[test]
374    fn wire_types_matches_every_node_variant() {
375        // Independent of `one_of_each`: the index match is exhaustive, so a new
376        // variant forces both this constant and the sample list to grow.
377        const VARIANT_COUNT: usize = 19;
378        assert_eq!(WIRE_TYPES.len(), VARIANT_COUNT);
379
380        let nodes = one_of_each();
381        assert_eq!(nodes.len(), VARIANT_COUNT);
382
383        let mut seen = [false; VARIANT_COUNT];
384        for node in &nodes {
385            let i = variant_index(node);
386            assert!(
387                i < VARIANT_COUNT,
388                "variant_index {i} is outside VARIANT_COUNT"
389            );
390            assert!(!seen[i], "duplicate sample for variant index {i}");
391            seen[i] = true;
392            assert_eq!(
393                node.type_name(),
394                WIRE_TYPES[i],
395                "sample at index {i} must match WIRE_TYPES"
396            );
397        }
398        assert!(
399            seen.iter().all(|&s| s),
400            "one_of_each must cover every variant index"
401        );
402    }
403
404    #[test]
405    fn is_wire_type_rejects_marketing_aliases() {
406        for good in WIRE_TYPES {
407            assert!(is_wire_type(good), "{good} should be a wire type");
408        }
409        for bad in ["CurrLIF", "Convolution", "Integrator", "SumPooling", ""] {
410            assert!(!is_wire_type(bad), "{bad} must not be a wire type");
411        }
412    }
413
414    #[test]
415    fn padding_wire_strings_round_trip() {
416        assert_eq!(padding_as_wire_str(&Padding::Same), Some("same"));
417        assert_eq!(padding_as_wire_str(&Padding::Valid), Some("valid"));
418        assert_eq!(padding_as_wire_str(&Padding::pair(1, 1)), None);
419        assert_eq!(padding_from_wire_str("same").unwrap(), Padding::Same);
420        assert_eq!(padding_from_wire_str("valid").unwrap(), Padding::Valid);
421    }
422
423    #[test]
424    fn padding_from_unknown_string_is_rejected() {
425        let err = padding_from_wire_str("SAME").unwrap_err();
426        assert!(matches!(err, NirError::InvalidGraph(_)));
427        assert!(err.to_string().contains("\"SAME\""));
428    }
429
430    #[test]
431    fn node_names_with_dots_are_allowed() {
432        // Real upstream fixtures use names like "lif1.lif".
433        assert!(check_link_name("node name", "lif1.lif").is_ok());
434        assert!(check_link_name("node name", "0").is_ok());
435    }
436
437    #[test]
438    fn illegal_link_names_are_rejected() {
439        for bad in ["", "a/b", ".", "..", "nul\0inside"] {
440            let err = check_link_name("node name", bad).unwrap_err();
441            assert!(
442                matches!(err, NirError::InvalidGraph(_)),
443                "{bad:?} should be rejected"
444            );
445        }
446    }
447
448    #[test]
449    fn the_kind_label_appears_in_the_error() {
450        let err = check_link_name("metadata key", "a/b").unwrap_err();
451        assert!(err.to_string().contains("metadata key"), "got {err}");
452    }
453
454    #[test]
455    fn nul_bytes_in_string_values_are_rejected() {
456        let err = check_hdf5_string("version", "0.2\0.0").unwrap_err();
457        assert!(matches!(err, NirError::InvalidGraph(_)));
458        assert!(err.to_string().contains("version"), "got {err}");
459        assert!(err.to_string().contains("NUL"), "got {err}");
460    }
461}
nir_rs/io/wire.rs

nir_rs/io/
wire.rs
