Skip to main content

vyre_runtime/megakernel/advanced/
zero_copy_io.rs

1//! GPU-initiated IO fragment via `AsyncLoad`.
2//!
3//! Emits the IR node that asks the runtime scheduler to map source and
4//! destination capability tables onto the concrete ingest path. Linux runtimes
5//! wire that request to registered mapped reads or the native GPUDirect NVMe
6//! driver; this module owns only the device-side request fragment.
7
8use vyre_foundation::ir::{Expr, Node};
9
10use crate::resident_work_queue::io::{
11    IO_DESTINATION_CAPABILITY_TABLE, IO_QUEUE_DMA_TAG, IO_SOURCE_CAPABILITY_TABLE,
12};
13
14/// Binding names for a GPU-initiated direct file pull.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct DirectIoBindings {
17    /// Source handle table or device namespace.
18    pub source: &'static str,
19    /// Destination GPU cache buffer.
20    pub destination: &'static str,
21    /// Variable naming the first byte to pull.
22    pub file_start: &'static str,
23    /// Variable naming one-past-last byte to pull.
24    pub file_end: &'static str,
25    /// Trace tag attached to the async transfer.
26    pub tag: &'static str,
27}
28
29impl Default for DirectIoBindings {
30    fn default() -> Self {
31        Self {
32            source: IO_SOURCE_CAPABILITY_TABLE,
33            destination: IO_DESTINATION_CAPABILITY_TABLE,
34            file_start: "file_start",
35            file_end: "file_end",
36            tag: IO_QUEUE_DMA_TAG,
37        }
38    }
39}
40
41/// Emit a GPU-initiated direct file pull using default megakernel names.
42///
43/// This translates to the backend's async-load path; callers must pair it
44/// with the runtime IO queue that maps `source` to device/offset handles.
45#[must_use]
46pub fn pull_file_async_direct() -> Node {
47    pull_file_async_direct_with(&DirectIoBindings::default())
48}
49
50/// Emit a GPU-initiated direct file pull using custom binding names.
51#[must_use]
52pub fn pull_file_async_direct_with(bindings: &DirectIoBindings) -> Node {
53    Node::async_load_gpu_driven(
54        bindings.source,
55        bindings.destination,
56        Expr::var(bindings.file_start),
57        Expr::sub(Expr::var(bindings.file_end), Expr::var(bindings.file_start)),
58        bindings.tag,
59    )
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn direct_io_uses_extended_async_load_fields() {
68        let node = pull_file_async_direct();
69        let Node::AsyncLoad {
70            source,
71            destination,
72            offset,
73            size,
74            tag,
75        } = node
76        else {
77            panic!("direct IO must emit AsyncLoad");
78        };
79        assert_eq!(source.as_str(), IO_SOURCE_CAPABILITY_TABLE);
80        assert_eq!(destination.as_str(), IO_DESTINATION_CAPABILITY_TABLE);
81        assert_eq!(tag.as_str(), IO_QUEUE_DMA_TAG);
82        assert!(matches!(*offset, Expr::Var(_)));
83        assert!(matches!(*size, Expr::BinOp { .. }));
84    }
85}