weir/ifds_resident_types/
prepared.rs

1use crate::ifds_gpu::IfdsShape;
2
3/// Reusable IFDS exploded-supergraph CSR prepared once for many seed queries.
4///
5/// The packed ProgramGraph buffers are retained in the shape consumed by
6/// `ifds_gpu_step`, so callers that solve many source/fact queries over one
7/// graph avoid rebuilding and repacking the invariant CSR every time.
8#[derive(Clone, Debug, PartialEq)]
9pub struct PreparedIfdsCsr {
10    pub(crate) shape: IfdsShape,
11    pub(crate) node_count: u32,
12    pub(crate) words: usize,
13    pub(crate) pg_nodes_b: Vec<u8>,
14    pub(crate) row_ptr_b: Vec<u8>,
15    pub(crate) col_idx_b: Vec<u8>,
16    pub(crate) pg_edge_kind_mask_b: Vec<u8>,
17    pub(crate) pg_node_tags_b: Vec<u8>,
18}
19
20impl PreparedIfdsCsr {
21    /// IFDS problem shape encoded by this prepared CSR.
22    #[must_use]
23    pub fn shape(&self) -> IfdsShape {
24        self.shape
25    }
26
27    /// Exploded-supergraph node count.
28    #[must_use]
29    pub fn node_count(&self) -> u32 {
30        self.node_count
31    }
32
33    /// Number of u32 words in one frontier bitset.
34    #[must_use]
35    pub fn frontier_words(&self) -> usize {
36        self.words
37    }
38
39    /// Total resident bytes required by invariant packed graph buffers.
40    #[must_use]
41    pub fn retained_graph_bytes(&self) -> usize {
42        fn padded(bytes: &[u8]) -> &[u8] {
43            const EMPTY_U32: [u8; std::mem::size_of::<u32>()] = [0; std::mem::size_of::<u32>()];
44            if bytes.is_empty() {
45                &EMPTY_U32
46            } else {
47                bytes
48            }
49        }
50
51        let pg_nodes_b = padded(&self.pg_nodes_b);
52        let row_ptr_b = padded(&self.row_ptr_b);
53        let col_idx_b = padded(&self.col_idx_b);
54        let pg_edge_kind_mask_b = padded(&self.pg_edge_kind_mask_b);
55        let pg_node_tags_b = padded(&self.pg_node_tags_b);
56        let node_tag_bytes = if pg_node_tags_b == pg_nodes_b {
57            0
58        } else {
59            pg_node_tags_b.len()
60        };
61        pg_nodes_b.len()
62            + row_ptr_b.len()
63            + col_idx_b.len()
64            + pg_edge_kind_mask_b.len()
65            + node_tag_bytes
66    }
67
68    /// Deterministic hash of the prepared IFDS CSR layout and packed bytes.
69    #[must_use]
70    pub fn stable_layout_hash(&self) -> u64 {
71        super::hash::stable_prepared_ifds_csr_hash(
72            self.shape,
73            self.node_count,
74            self.words,
75            &self.pg_nodes_b,
76            &self.row_ptr_b,
77            &self.col_idx_b,
78            &self.pg_edge_kind_mask_b,
79            &self.pg_node_tags_b,
80        )
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    fn inflated(bytes: &[u8], extra_capacity: usize) -> Vec<u8> {
89        let mut out = Vec::with_capacity(bytes.len() + extra_capacity);
90        out.extend_from_slice(bytes);
91        out
92    }
93
94    fn resident_upload_len(bytes: &[u8]) -> usize {
95        if bytes.is_empty() {
96            std::mem::size_of::<u32>()
97        } else {
98            bytes.len()
99        }
100    }
101
102    fn expected_physical_bytes(prepared: &PreparedIfdsCsr) -> usize {
103        fn padded_owned(bytes: &[u8]) -> Vec<u8> {
104            if bytes.is_empty() {
105                vec![0; std::mem::size_of::<u32>()]
106            } else {
107                bytes.to_vec()
108            }
109        }
110
111        let node_tags_alias_nodes =
112            padded_owned(&prepared.pg_nodes_b) == padded_owned(&prepared.pg_node_tags_b);
113        resident_upload_len(&prepared.pg_nodes_b)
114            + resident_upload_len(&prepared.row_ptr_b)
115            + resident_upload_len(&prepared.col_idx_b)
116            + resident_upload_len(&prepared.pg_edge_kind_mask_b)
117            + if node_tags_alias_nodes {
118                0
119            } else {
120                resident_upload_len(&prepared.pg_node_tags_b)
121            }
122    }
123
124    #[test]
125    fn retained_graph_bytes_counts_physical_resident_uploads_not_vec_capacity() {
126        let prepared = PreparedIfdsCsr {
127            shape: crate::ifds_gpu::IfdsShape {
128                num_procs: 1,
129                blocks_per_proc: 1,
130                facts_per_proc: 1,
131                edge_count: 0,
132            },
133            node_count: 1,
134            words: 1,
135            pg_nodes_b: inflated(&[0, 0, 0, 0], 1024),
136            row_ptr_b: inflated(&[0, 0, 0, 0, 0, 0, 0, 0], 512),
137            col_idx_b: inflated(&[], 256),
138            pg_edge_kind_mask_b: inflated(&[], 128),
139            pg_node_tags_b: inflated(&[0, 0, 0, 0], 4096),
140        };
141
142        assert_eq!(
143            prepared.retained_graph_bytes(),
144            20,
145            "resident retained-byte accounting must count four unique physical buffers: pg_nodes, row_ptr, padded col_idx, padded edge masks"
146        );
147    }
148
149    #[test]
150    fn retained_graph_bytes_matches_generated_alias_and_padding_cases() {
151        for case in 0..2048usize {
152            let node_len = case % 17;
153            let row_len = (case / 3) % 23;
154            let col_len = (case / 5) % 29;
155            let mask_len = (case / 7) % 31;
156            let tag_mode = case % 4;
157            let nodes = (0..node_len)
158                .map(|index| (case.wrapping_add(index * 17) & 0xff) as u8)
159                .collect::<Vec<_>>();
160            let tags = match tag_mode {
161                0 => nodes.clone(),
162                1 => Vec::new(),
163                2 => vec![0; std::mem::size_of::<u32>()],
164                _ => (0..node_len.max(1))
165                    .map(|index| (case.wrapping_mul(3).wrapping_add(index * 11) | 1) as u8)
166                    .collect(),
167            };
168            let prepared = PreparedIfdsCsr {
169                shape: crate::ifds_gpu::IfdsShape {
170                    num_procs: 1,
171                    blocks_per_proc: 1,
172                    facts_per_proc: 1,
173                    edge_count: col_len as u32,
174                },
175                node_count: node_len.max(1) as u32,
176                words: 1,
177                pg_nodes_b: inflated(&nodes, case % 13),
178                row_ptr_b: inflated(&vec![0x11; row_len], case % 11),
179                col_idx_b: inflated(&vec![0x22; col_len], case % 7),
180                pg_edge_kind_mask_b: inflated(&vec![0x33; mask_len], case % 5),
181                pg_node_tags_b: inflated(&tags, case % 19),
182            };
183
184            assert_eq!(
185                prepared.retained_graph_bytes(),
186                expected_physical_bytes(&prepared),
187                "case {case} must account padded physical resident bytes and node/tag aliasing"
188            );
189        }
190    }
191}
weir/ifds_resident_types/prepared.rs

weir/ifds_resident_types/
prepared.rs
