vyre_driver_cuda/
token_fact_graph_cuda_adapter.rs

1//! CUDA adapter for the unified resident token/fact graph.
2
3use std::{cmp::Reverse, collections::BinaryHeap};
4
5use crate::backend::accounting::{
6    checked_add_u64_count as checked_add, checked_mul_u64_count as checked_mul,
7    CudaArithmeticOverflow,
8};
9use vyre_driver::megakernel_execution::MegakernelGraphShape;
10use vyre_self_substrate::device_resident_token_fact_graph::DeviceResidentTokenFactGraph;
11
12/// Number of rank buckets carried for token/fact out-degree skew planning.
13pub const CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS: usize = 16;
14
15/// Power-of-two ranks used by the token/fact out-degree profile.
16pub const CUDA_TOKEN_FACT_DEGREE_PROFILE_RANKS: [u64; CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS] = [
17    1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1_024, 2_048, 4_096, 8_192, 16_384, 32_768,
18];
19
20const CUDA_TOKEN_FACT_DEGREE_PROFILE_MAX_RANK: usize = 32_768;
21
22/// CUDA resident byte envelope for the unified compiler/dataflow graph.
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub struct CudaTokenFactGraphLayout {
25    /// Scheduler-visible graph shape.
26    pub graph_shape: MegakernelGraphShape,
27    /// Maximum outgoing CSR row degree in the resident token/fact graph.
28    pub max_out_degree: u64,
29    /// Prefix sums of top out-degrees at `CUDA_TOKEN_FACT_DEGREE_PROFILE_RANKS`.
30    pub top_out_degree_prefix_sums: [u64; CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS],
31    /// Fixed bytes per resident node record.
32    pub node_record_bytes: u64,
33    /// Fixed bytes per resident edge record.
34    pub edge_record_bytes: u64,
35    /// Bytes for resident node records.
36    pub node_bytes: u64,
37    /// Bytes for resident edge records.
38    pub edge_bytes: u64,
39    /// Bytes for the shared token/fact payload slab.
40    pub payload_bytes: u64,
41    /// Total bytes that must remain device-resident for the layout.
42    pub resident_bytes: u64,
43}
44
45impl CudaTokenFactGraphLayout {
46    /// Build a layout from aggregate byte fields when CSR row offsets are not
47    /// available to the caller. This preserves correctness by treating total
48    /// edge count as the maximum possible row degree.
49    #[must_use]
50    pub const fn from_aggregate_fields(
51        graph_shape: MegakernelGraphShape,
52        node_record_bytes: u64,
53        edge_record_bytes: u64,
54        node_bytes: u64,
55        edge_bytes: u64,
56        payload_bytes: u64,
57        resident_bytes: u64,
58    ) -> Self {
59        Self {
60            graph_shape,
61            max_out_degree: graph_shape.edge_count,
62            top_out_degree_prefix_sums: [graph_shape.edge_count;
63                CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS],
64            node_record_bytes,
65            edge_record_bytes,
66            node_bytes,
67            edge_bytes,
68            payload_bytes,
69            resident_bytes,
70        }
71    }
72}
73
74/// CUDA token/fact adapter errors.
75#[derive(Clone, Debug, Eq, PartialEq)]
76pub enum CudaTokenFactGraphLayoutError {
77    /// Record widths must be explicit, non-zero ABI values.
78    ZeroRecordWidth {
79        /// Field that was zero.
80        field: &'static str,
81    },
82    /// Public CSR fields are inconsistent with each other.
83    InvalidCsrShape {
84        /// Invalid CSR field or relationship.
85        field: &'static str,
86    },
87    /// Byte arithmetic overflowed.
88    ByteCountOverflow {
89        /// Field being computed.
90        field: &'static str,
91    },
92}
93
94impl std::fmt::Display for CudaTokenFactGraphLayoutError {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        match self {
97            Self::ZeroRecordWidth { field } => write!(
98                f,
99                "CUDA token/fact graph adapter received zero {field}. Fix: pass the concrete resident ABI record width."
100            ),
101            Self::InvalidCsrShape { field } => write!(
102                f,
103                "CUDA token/fact graph adapter received invalid CSR {field}. Fix: rebuild the token/fact graph through the canonical resident graph planner."
104            ),
105            Self::ByteCountOverflow { field } => write!(
106                f,
107                "CUDA token/fact graph adapter overflowed while computing {field}. Fix: shard the token/fact graph before resident upload."
108            ),
109        }
110    }
111}
112
113impl std::error::Error for CudaTokenFactGraphLayoutError {}
114
115impl CudaArithmeticOverflow for CudaTokenFactGraphLayoutError {
116    fn arithmetic_overflow(field: &'static str) -> Self {
117        Self::ByteCountOverflow { field }
118    }
119}
120
121/// Convert the unified token/fact graph into CUDA scheduler shape and bytes.
122pub fn adapt_token_fact_graph_to_cuda_layout(
123    graph: &DeviceResidentTokenFactGraph,
124    node_record_bytes: u64,
125    edge_record_bytes: u64,
126) -> Result<CudaTokenFactGraphLayout, CudaTokenFactGraphLayoutError> {
127    if node_record_bytes == 0 {
128        return Err(CudaTokenFactGraphLayoutError::ZeroRecordWidth {
129            field: "node_record_bytes",
130        });
131    }
132    if edge_record_bytes == 0 {
133        return Err(CudaTokenFactGraphLayoutError::ZeroRecordWidth {
134            field: "edge_record_bytes",
135        });
136    }
137    let node_count = u64::try_from(graph.node_ids.len()).map_err(|_| {
138        CudaTokenFactGraphLayoutError::ByteCountOverflow {
139            field: "node count",
140        }
141    })?;
142    let edge_count = u64::try_from(graph.column_indices.len()).map_err(|_| {
143        CudaTokenFactGraphLayoutError::ByteCountOverflow {
144            field: "edge count",
145        }
146    })?;
147    let (max_out_degree, top_out_degree_prefix_sums) = csr_out_degree_profile(graph, edge_count)?;
148    let node_bytes = checked_mul(node_count, node_record_bytes, "node bytes")?;
149    let edge_bytes = checked_mul(edge_count, edge_record_bytes, "edge bytes")?;
150    let resident_without_payload = checked_add(node_bytes, edge_bytes, "node plus edge bytes")?;
151    let resident_bytes = checked_add(
152        resident_without_payload,
153        graph.payload_bytes,
154        "resident bytes",
155    )?;
156
157    Ok(CudaTokenFactGraphLayout {
158        graph_shape: MegakernelGraphShape {
159            node_count,
160            edge_count,
161        },
162        max_out_degree,
163        top_out_degree_prefix_sums,
164        node_record_bytes,
165        edge_record_bytes,
166        node_bytes,
167        edge_bytes,
168        payload_bytes: graph.payload_bytes,
169        resident_bytes,
170    })
171}
172
173fn csr_out_degree_profile(
174    graph: &DeviceResidentTokenFactGraph,
175    edge_count: u64,
176) -> Result<(u64, [u64; CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS]), CudaTokenFactGraphLayoutError> {
177    let expected_row_offsets = graph.node_ids.len().checked_add(1).ok_or(
178        CudaTokenFactGraphLayoutError::ByteCountOverflow {
179            field: "row offset count",
180        },
181    )?;
182    if graph.row_offsets.len() != expected_row_offsets {
183        return Err(CudaTokenFactGraphLayoutError::InvalidCsrShape {
184            field: "row_offsets length",
185        });
186    }
187    let declared_edges = u64::from(*graph.row_offsets.last().ok_or(
188        CudaTokenFactGraphLayoutError::InvalidCsrShape {
189            field: "row_offsets terminator",
190        },
191    )?);
192    if declared_edges != edge_count {
193        return Err(CudaTokenFactGraphLayoutError::InvalidCsrShape {
194            field: "row_offsets edge count",
195        });
196    }
197    let profile_capacity = CUDA_TOKEN_FACT_DEGREE_PROFILE_MAX_RANK.min(graph.node_ids.len());
198    let mut top_degrees = BinaryHeap::with_capacity(profile_capacity);
199    let mut max_out_degree = 0_u64;
200    for row in graph.row_offsets.windows(2) {
201        let start = row[0];
202        let end = row[1];
203        if end < start {
204            return Err(CudaTokenFactGraphLayoutError::InvalidCsrShape {
205                field: "row_offsets ordering",
206            });
207        }
208        let degree = u64::from(end - start);
209        max_out_degree = max_out_degree.max(degree);
210        if top_degrees.len() < profile_capacity {
211            top_degrees.push(Reverse(degree));
212        } else if let Some(mut min_degree) = top_degrees.peek_mut() {
213            if degree > min_degree.0 {
214                *min_degree = Reverse(degree);
215            }
216        }
217    }
218    let mut degrees = top_degrees
219        .into_iter()
220        .map(|Reverse(degree)| degree)
221        .collect::<Vec<_>>();
222    degrees.sort_unstable_by(|lhs, rhs| rhs.cmp(lhs));
223
224    let mut prefix_sum = 0_u64;
225    let mut prefix_sums = [0_u64; CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS];
226    let mut bucket = 0_usize;
227    for (index, degree) in degrees.into_iter().enumerate() {
228        prefix_sum = checked_add(prefix_sum, degree, "top out-degree prefix sum")?;
229        let rank = u64::try_from(index + 1).map_err(|_| {
230            CudaTokenFactGraphLayoutError::ByteCountOverflow {
231                field: "out-degree profile rank",
232            }
233        })?;
234        while bucket < CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS
235            && rank >= CUDA_TOKEN_FACT_DEGREE_PROFILE_RANKS[bucket]
236        {
237            prefix_sums[bucket] = prefix_sum;
238            bucket += 1;
239        }
240    }
241    while bucket < CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS {
242        prefix_sums[bucket] = prefix_sum;
243        bucket += 1;
244    }
245    Ok((max_out_degree, prefix_sums))
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use vyre_driver::megakernel_execution::{
252        plan_megakernel_memory_budget, MegakernelExecutionTopology,
253    };
254    use vyre_self_substrate::device_resident_token_fact_graph::{
255        plan_device_resident_token_fact_graph, TokenFactEdge, TokenFactEdgeKind, TokenFactNode,
256        TokenFactNodeKind,
257    };
258
259    #[test]
260    fn adapter_accounts_for_cuda_resident_token_fact_layout() {
261        let graph = plan_device_resident_token_fact_graph(
262            &[
263                node(1, TokenFactNodeKind::Token, 0, 8),
264                node(2, TokenFactNodeKind::Semantic, 8, 8),
265                node(3, TokenFactNodeKind::Fact, 16, 8),
266            ],
267            &[
268                edge(1, 2, TokenFactEdgeKind::SemanticFact),
269                edge(2, 3, TokenFactEdgeKind::FactDependency),
270            ],
271            24,
272        )
273        .expect("Fix: token/fact graph should pack");
274
275        let cuda = adapt_token_fact_graph_to_cuda_layout(&graph, 32, 16)
276            .expect("Fix: token/fact graph should adapt to CUDA layout");
277
278        assert_eq!(cuda.graph_shape.node_count, 3);
279        assert_eq!(cuda.graph_shape.edge_count, 2);
280        assert_eq!(cuda.max_out_degree, 1);
281        assert_eq!(cuda.top_out_degree_prefix_sums[0], 1);
282        assert_eq!(cuda.top_out_degree_prefix_sums[1], 2);
283        assert_eq!(cuda.top_out_degree_prefix_sums[15], 2);
284        assert_eq!(cuda.node_bytes, 96);
285        assert_eq!(cuda.edge_bytes, 32);
286        assert_eq!(cuda.resident_bytes, 152);
287        let memory = plan_megakernel_memory_budget(
288            MegakernelExecutionTopology::SparseFrontier,
289            cuda.graph_shape,
290            cuda.node_record_bytes,
291            cuda.edge_record_bytes,
292            64,
293            cuda.payload_bytes,
294            16,
295            512,
296        )
297        .expect("Fix: adapted token/fact graph should feed CUDA memory planning");
298        assert_eq!(memory.graph_bytes, 128);
299    }
300
301    #[test]
302    fn adapter_exports_max_out_degree_for_hub_heavy_queue_planning() {
303        let graph = plan_device_resident_token_fact_graph(
304            &[
305                node(1, TokenFactNodeKind::Fact, 0, 4),
306                node(2, TokenFactNodeKind::Fact, 4, 4),
307                node(3, TokenFactNodeKind::Fact, 8, 4),
308                node(4, TokenFactNodeKind::Fact, 12, 4),
309            ],
310            &[
311                edge(1, 2, TokenFactEdgeKind::FactDependency),
312                edge(1, 3, TokenFactEdgeKind::FactDependency),
313                edge(1, 4, TokenFactEdgeKind::FactDependency),
314                edge(2, 3, TokenFactEdgeKind::FactDependency),
315            ],
316            16,
317        )
318        .expect("Fix: hub-heavy token/fact graph should pack");
319
320        let cuda = adapt_token_fact_graph_to_cuda_layout(&graph, 32, 16)
321            .expect("Fix: hub-heavy token/fact graph should adapt to CUDA layout");
322
323        assert_eq!(cuda.graph_shape.edge_count, 4);
324        assert_eq!(cuda.max_out_degree, 3);
325        assert_eq!(cuda.top_out_degree_prefix_sums[0], 3);
326        assert_eq!(cuda.top_out_degree_prefix_sums[1], 4);
327        assert_eq!(cuda.top_out_degree_prefix_sums[2], 4);
328    }
329
330    #[test]
331    fn generated_adapter_profiles_top_out_degree_prefixes() {
332        let mut state = 0x5eec_c0de_f00d_7715_u64;
333        for case_index in 0..4096_u64 {
334            let node_count = 1 + (next_u64(&mut state) % 64) as u32;
335            let nodes = (0..node_count)
336                .map(|index| node(index + 1, TokenFactNodeKind::Fact, u64::from(index) * 4, 4))
337                .collect::<Vec<_>>();
338            let mut edges = Vec::new();
339            if case_index % 4 == 0 {
340                for to in 2..=node_count {
341                    edges.push(edge(1, to, TokenFactEdgeKind::FactDependency));
342                }
343            }
344            let attempts = next_u64(&mut state) % (u64::from(node_count) * 5 + 1);
345            for _ in 0..attempts {
346                let from = 1 + (next_u64(&mut state) % u64::from(node_count)) as u32;
347                let to = 1 + (next_u64(&mut state) % u64::from(node_count)) as u32;
348                let kind = if next_u64(&mut state) & 1 == 0 {
349                    TokenFactEdgeKind::FactDependency
350                } else {
351                    TokenFactEdgeKind::DiagnosticProvenance
352                };
353                edges.push(edge(from, to, kind));
354            }
355            let graph =
356                plan_device_resident_token_fact_graph(&nodes, &edges, u64::from(node_count) * 4)
357                    .expect("Fix: generated token/fact graph should pack");
358            let cuda = adapt_token_fact_graph_to_cuda_layout(&graph, 32, 16)
359                .expect("Fix: generated token/fact graph should adapt");
360            let mut degrees = graph
361                .row_offsets
362                .windows(2)
363                .map(|row| u64::from(row[1] - row[0]))
364                .collect::<Vec<_>>();
365            degrees.sort_unstable_by(|lhs, rhs| rhs.cmp(lhs));
366
367            assert_eq!(
368                cuda.max_out_degree,
369                degrees.first().copied().unwrap_or(0),
370                "case {case_index}"
371            );
372            for (bucket, rank) in CUDA_TOKEN_FACT_DEGREE_PROFILE_RANKS.iter().enumerate() {
373                let expected = degrees
374                    .iter()
375                    .take((*rank as usize).min(degrees.len()))
376                    .copied()
377                    .sum::<u64>();
378                assert_eq!(
379                    cuda.top_out_degree_prefix_sums[bucket], expected,
380                    "case {case_index} bucket {bucket}"
381                );
382            }
383        }
384    }
385
386    #[test]
387    fn adapter_profiles_large_graph_with_bounded_top_rank_storage() {
388        let node_count = 32_770_u32;
389        let nodes = (0..node_count)
390            .map(|index| node(index + 1, TokenFactNodeKind::Fact, u64::from(index) * 4, 4))
391            .collect::<Vec<_>>();
392        let mut edges = Vec::with_capacity(32_858);
393        for to in 2..=51 {
394            edges.push(edge(1, to, TokenFactEdgeKind::FactDependency));
395        }
396        for to in 3..=42 {
397            edges.push(edge(2, to, TokenFactEdgeKind::FactDependency));
398        }
399        for from in 3..=node_count {
400            edges.push(edge(from, 1, TokenFactEdgeKind::FactDependency));
401        }
402        let graph =
403            plan_device_resident_token_fact_graph(&nodes, &edges, u64::from(node_count) * 4)
404                .expect("Fix: large skewed token/fact graph should pack");
405
406        let cuda = adapt_token_fact_graph_to_cuda_layout(&graph, 32, 16)
407            .expect("Fix: large skewed token/fact graph should adapt");
408
409        assert_eq!(cuda.graph_shape.node_count, u64::from(node_count));
410        assert_eq!(cuda.graph_shape.edge_count, 32_858);
411        assert_eq!(cuda.max_out_degree, 50);
412        assert_eq!(cuda.top_out_degree_prefix_sums[0], 50);
413        assert_eq!(cuda.top_out_degree_prefix_sums[1], 90);
414        assert_eq!(cuda.top_out_degree_prefix_sums[2], 92);
415        assert_eq!(cuda.top_out_degree_prefix_sums[15], 32_856);
416    }
417
418    #[test]
419    fn aggregate_layout_constructor_preserves_legacy_safe_edge_bound() {
420        let layout = CudaTokenFactGraphLayout::from_aggregate_fields(
421            MegakernelGraphShape {
422                node_count: 4,
423                edge_count: 9,
424            },
425            32,
426            16,
427            128,
428            144,
429            64,
430            336,
431        );
432
433        assert_eq!(layout.max_out_degree, 9);
434        assert_eq!(
435            layout.top_out_degree_prefix_sums,
436            [9; CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS]
437        );
438        assert_eq!(layout.resident_bytes, 336);
439    }
440
441    #[test]
442    fn adapter_rejects_missing_abi_widths() {
443        let graph = plan_device_resident_token_fact_graph(&[], &[], 0)
444            .expect("Fix: empty graph still has a valid resident layout");
445
446        assert_eq!(
447            adapt_token_fact_graph_to_cuda_layout(&graph, 0, 8)
448                .expect_err("zero node record width should fail"),
449            CudaTokenFactGraphLayoutError::ZeroRecordWidth {
450                field: "node_record_bytes",
451            }
452        );
453        assert_eq!(
454            adapt_token_fact_graph_to_cuda_layout(&graph, 8, 0)
455                .expect_err("zero edge record width should fail"),
456            CudaTokenFactGraphLayoutError::ZeroRecordWidth {
457                field: "edge_record_bytes",
458            }
459        );
460    }
461
462    #[test]
463    fn adapter_rejects_public_graphs_with_invalid_csr_rows() {
464        let mut graph = plan_device_resident_token_fact_graph(
465            &[node(1, TokenFactNodeKind::Fact, 0, 4)],
466            &[],
467            4,
468        )
469        .expect("Fix: token/fact graph should pack before adversarial mutation");
470        graph.row_offsets[1] = 1;
471
472        assert_eq!(
473            adapt_token_fact_graph_to_cuda_layout(&graph, 32, 16)
474                .expect_err("invalid CSR row offsets should fail before CUDA planning"),
475            CudaTokenFactGraphLayoutError::InvalidCsrShape {
476                field: "row_offsets edge count",
477            }
478        );
479    }
480
481    fn node(
482        id: u32,
483        kind: TokenFactNodeKind,
484        payload_offset: u64,
485        payload_bytes: u64,
486    ) -> TokenFactNode {
487        TokenFactNode {
488            id,
489            kind,
490            payload_offset,
491            payload_bytes,
492        }
493    }
494
495    fn edge(from: u32, to: u32, kind: TokenFactEdgeKind) -> TokenFactEdge {
496        TokenFactEdge { from, to, kind }
497    }
498
499    fn next_u64(state: &mut u64) -> u64 {
500        *state = state
501            .wrapping_mul(6_364_136_223_846_793_005)
502            .wrapping_add(1_442_695_040_888_963_407);
503        *state
504    }
505}
vyre_driver_cuda/token_fact_graph_cuda_adapter.rs

vyre_driver_cuda/
token_fact_graph_cuda_adapter.rs
