vyre_driver_reference/
lib.rs

1#![forbid(unsafe_code)]
2
3//! Registry adapter that exposes `vyre-reference` as a `VyreBackend`.
4
5use std::sync::Arc;
6
7use vyre_driver::backend::private;
8use vyre_driver::backend::{
9    core_supported_ops, BackendCapability, BackendError, BackendPrecedence, BackendRegistration,
10};
11use vyre_driver::{DispatchConfig, VyreBackend};
12use vyre_foundation::ir::{BufferAccess, BufferDecl, Program};
13use vyre_reference::value::Value;
14
15/// Stable backend id for the pure-Rust reference interpreter.
16pub const CPU_REF_BACKEND_ID: &str = "cpu-ref";
17
18/// Dispatch backend backed by `vyre_reference::reference_eval`.
19#[derive(Debug, Default, Clone, Copy)]
20pub struct CpuRefBackend;
21
22impl private::Sealed for CpuRefBackend {}
23
24impl VyreBackend for CpuRefBackend {
25    fn id(&self) -> &'static str {
26        CPU_REF_BACKEND_ID
27    }
28
29    fn version(&self) -> &'static str {
30        env!("CARGO_PKG_VERSION")
31    }
32
33    fn dispatch(
34        &self,
35        program: &Program,
36        inputs: &[Vec<u8>],
37        _config: &DispatchConfig,
38    ) -> Result<Vec<Vec<u8>>, BackendError> {
39        let values = reference_values(program, inputs)?;
40        vyre_reference::reference_eval(program, &values)
41            .map(|outputs| outputs.iter().map(Value::to_bytes).collect())
42            .map_err(|error| {
43                BackendError::new(format!(
44                    "cpu-ref reference dispatch failed: {error}. Fix: validate the Program and input buffer ABI before dispatch."
45                ))
46            })
47    }
48
49    fn supported_ops(&self) -> &std::collections::HashSet<vyre_foundation::ir::OpId> {
50        core_supported_ops()
51    }
52
53    fn max_workgroup_size(&self) -> [u32; 3] {
54        [1024, 1, 1]
55    }
56
57    fn max_compute_workgroups_per_dimension(&self) -> u32 {
58        u32::MAX
59    }
60}
61
62fn reference_values(program: &Program, inputs: &[Vec<u8>]) -> Result<Vec<Value>, BackendError> {
63    let backend_allocated_output = |buffer: &BufferDecl| {
64        buffer.is_output()
65            || buffer.access() == BufferAccess::WriteOnly
66            || (buffer.is_pipeline_live_out() && buffer.access() == BufferAccess::ReadWrite)
67    };
68    let logical_input_count = program
69        .buffers()
70        .iter()
71        .filter(|buffer| {
72            buffer.access() != BufferAccess::Workgroup && !backend_allocated_output(buffer)
73        })
74        .count();
75    let legacy_input_count = program
76        .buffers()
77        .iter()
78        .filter(|buffer| buffer.access() != BufferAccess::Workgroup)
79        .count();
80    let legacy_input_mode =
81        inputs.len() == legacy_input_count && inputs.len() != logical_input_count;
82    let mut next_input = 0usize;
83    let mut values = Vec::new();
84    for buffer in program.buffers() {
85        if buffer.access() == BufferAccess::Workgroup {
86            continue;
87        }
88        let bytes = if backend_allocated_output(buffer) {
89            if legacy_input_mode {
90                let _legacy_initializer = inputs.get(next_input).ok_or_else(|| {
91                    BackendError::new(format!(
92                        "cpu-ref missing legacy output initializer for buffer `{}`. Fix: pass one buffer for every non-workgroup declaration or migrate to logical backend inputs.",
93                        buffer.name()
94                    ))
95                })?;
96                next_input += 1;
97            }
98            synthesized_zero_buffer(buffer, "backend-allocated output")?
99        } else if let Some(input) = inputs.get(next_input) {
100            next_input += 1;
101            input.clone()
102        } else {
103            synthesized_zero_buffer(buffer, "missing input")?
104        };
105        values.push(Value::Bytes(Arc::from(bytes)));
106    }
107    if next_input != inputs.len() {
108        return Err(BackendError::new(format!(
109            "cpu-ref received {} extra input buffer(s). Fix: pass inputs in Program::buffers order without trailing buffers.",
110            inputs.len() - next_input
111        )));
112    }
113    Ok(values)
114}
115
116fn synthesized_zero_buffer(
117    buffer: &BufferDecl,
118    role: &'static str,
119) -> Result<Vec<u8>, BackendError> {
120    let element_size = buffer.element().size_bytes().ok_or_else(|| {
121        BackendError::new(format!(
122            "cpu-ref cannot synthesize {role} buffer `{}` because its element type is unsized. Fix: declare fixed-width buffers or pass an explicit input buffer.",
123            buffer.name()
124        ))
125    })?;
126    let byte_len = usize::try_from(buffer.count())
127        .ok()
128        .and_then(|count| count.checked_mul(element_size))
129        .ok_or_else(|| {
130            BackendError::new(format!(
131                "cpu-ref {role} buffer `{}` size overflows usize. Fix: use a representable buffer size.",
132                buffer.name()
133            ))
134        })?;
135    Ok(vec![0u8; byte_len])
136}
137
138fn acquire_cpu_ref() -> Result<Box<dyn VyreBackend>, BackendError> {
139    Ok(Box::new(CpuRefBackend))
140}
141
142inventory::submit! {
143    BackendRegistration {
144        id: CPU_REF_BACKEND_ID,
145        factory: acquire_cpu_ref,
146        supported_ops: core_supported_ops,
147    }
148}
149
150inventory::submit! {
151    BackendCapability {
152        id: CPU_REF_BACKEND_ID,
153        dispatches: true,
154    }
155}
156
157inventory::submit! {
158    BackendPrecedence {
159        id: CPU_REF_BACKEND_ID,
160        rank: 900,
161    }
162}
vyre_driver_reference/lib.rs

vyre_driver_reference/
lib.rs
