Skip to main content

weir/
ifds_vyre_resident.rs

1//! Adapter from the common Vyre resident backend contract to Weir IFDS.
2//!
3//! IFDS owns its resident dispatch trait so Weir does not depend on concrete
4//! driver crates. This module bridges that trait to any `VyreBackend` that
5//! implements resident allocation, batched upload, resident dispatch, ranged
6//! download, and free through the common Vyre contract.
7
8use crate::ifds_gpu::IfdsResidentDispatch;
9
10/// Weir IFDS resident dispatcher backed by a generic Vyre backend.
11#[derive(Clone, Copy, Debug)]
12pub struct VyreResidentIfds<'a, B: vyre::VyreBackend + ?Sized> {
13    backend: &'a B,
14}
15
16impl<'a, B: vyre::VyreBackend + ?Sized> VyreResidentIfds<'a, B> {
17    /// Create an IFDS resident dispatcher over `backend`.
18    #[must_use]
19    pub const fn new(backend: &'a B) -> Self {
20        Self { backend }
21    }
22
23    /// Underlying Vyre backend.
24    #[must_use]
25    pub const fn backend(&self) -> &'a B {
26        self.backend
27    }
28}
29
30impl<B: vyre::VyreBackend + ?Sized> IfdsResidentDispatch for VyreResidentIfds<'_, B> {
31    type Resource = vyre::backend::Resource;
32
33    fn resident_backend_id(&self) -> &'static str {
34        self.backend.id()
35    }
36
37    fn resident_backend_version(&self) -> &'static str {
38        self.backend.version()
39    }
40
41    fn allocate_resident(&self, byte_len: usize) -> Result<Self::Resource, String> {
42        vyre::VyreBackend::allocate_resident(self.backend, byte_len)
43            .map_err(|error| error.to_string())
44    }
45
46    fn upload_resident(&self, resource: &Self::Resource, bytes: &[u8]) -> Result<(), String> {
47        vyre::VyreBackend::upload_resident(self.backend, resource, bytes)
48            .map_err(|error| error.to_string())
49    }
50
51    fn upload_resident_many(&self, uploads: &[(&Self::Resource, &[u8])]) -> Result<(), String> {
52        vyre::VyreBackend::upload_resident_many(self.backend, uploads)
53            .map_err(|error| error.to_string())
54    }
55
56    fn upload_resident_at_many(
57        &self,
58        uploads: &[(&Self::Resource, usize, &[u8])],
59    ) -> Result<(), String> {
60        vyre::VyreBackend::upload_resident_at_many(self.backend, uploads)
61            .map_err(|error| error.to_string())
62    }
63
64    fn download_resident(&self, resource: &Self::Resource) -> Result<Vec<u8>, String> {
65        vyre::VyreBackend::download_resident(self.backend, resource)
66            .map_err(|error| error.to_string())
67    }
68
69    fn download_resident_into(
70        &self,
71        resource: &Self::Resource,
72        output: &mut Vec<u8>,
73    ) -> Result<(), String> {
74        vyre::VyreBackend::download_resident_into(self.backend, resource, output)
75            .map_err(|error| error.to_string())
76    }
77
78    fn download_resident_range(
79        &self,
80        resource: &Self::Resource,
81        byte_offset: usize,
82        byte_len: usize,
83    ) -> Result<Vec<u8>, String> {
84        vyre::VyreBackend::download_resident_range(self.backend, resource, byte_offset, byte_len)
85            .map_err(|error| error.to_string())
86    }
87
88    fn download_resident_range_into(
89        &self,
90        resource: &Self::Resource,
91        byte_offset: usize,
92        byte_len: usize,
93        output: &mut Vec<u8>,
94    ) -> Result<(), String> {
95        vyre::VyreBackend::download_resident_range_into(
96            self.backend,
97            resource,
98            byte_offset,
99            byte_len,
100            output,
101        )
102        .map_err(|error| error.to_string())
103    }
104
105    fn download_resident_ranges_into(
106        &self,
107        ranges: &[(&Self::Resource, usize, usize)],
108        outputs: &mut [&mut Vec<u8>],
109    ) -> Result<(), String> {
110        vyre::VyreBackend::download_resident_ranges_into(self.backend, ranges, outputs)
111            .map_err(|error| error.to_string())
112    }
113
114    fn free_resident(&self, resource: Self::Resource) -> Result<(), String> {
115        vyre::VyreBackend::free_resident(self.backend, resource).map_err(|error| error.to_string())
116    }
117
118    fn dispatch_resident(
119        &self,
120        program: &vyre::Program,
121        resources: &[Self::Resource],
122        grid_override: Option<[u32; 3]>,
123    ) -> Result<(), String> {
124        let mut config = vyre::DispatchConfig::default();
125        config.grid_override = grid_override;
126        vyre::VyreBackend::dispatch_resident_timed(self.backend, program, resources, &config)
127            .map(|_| ())
128            .map_err(|error| error.to_string())
129    }
130
131    fn dispatch_resident_sequence_read_ranges_into(
132        &self,
133        steps: &[(&vyre::Program, &[Self::Resource], Option<[u32; 3]>)],
134        read_ranges: &[(&Self::Resource, usize, usize)],
135        outputs: &mut [&mut Vec<u8>],
136    ) -> Result<(), String> {
137        let resident_steps = resident_steps_from(steps)?;
138        let resident_read_ranges = resident_read_ranges_from(read_ranges)?;
139        vyre::VyreBackend::dispatch_resident_sequence_read_ranges_into(
140            self.backend,
141            &resident_steps,
142            &resident_read_ranges,
143            outputs,
144        )
145        .map_err(|error| error.to_string())
146    }
147
148    fn dispatch_resident_repeated_sequence_read_ranges_into(
149        &self,
150        prefix_steps: &[(&vyre::Program, &[Self::Resource], Option<[u32; 3]>)],
151        repeated_steps: &[(&vyre::Program, &[Self::Resource], Option<[u32; 3]>)],
152        repeat_count: u32,
153        read_ranges: &[(&Self::Resource, usize, usize)],
154        outputs: &mut [&mut Vec<u8>],
155    ) -> Result<(), String> {
156        let resident_prefix_steps = resident_steps_from(prefix_steps)?;
157        let resident_repeated_steps = resident_steps_from(repeated_steps)?;
158        let resident_read_ranges = resident_read_ranges_from(read_ranges)?;
159        vyre::VyreBackend::dispatch_resident_repeated_sequence_read_ranges_into(
160            self.backend,
161            &resident_prefix_steps,
162            &resident_repeated_steps,
163            repeat_count,
164            &resident_read_ranges,
165            outputs,
166        )
167        .map_err(|error| error.to_string())
168    }
169}
170
171#[allow(clippy::type_complexity)]
172fn resident_steps_from<'a>(
173    steps: &'a [(
174        &'a vyre::Program,
175        &'a [vyre::backend::Resource],
176        Option<[u32; 3]>,
177    )],
178) -> Result<Vec<vyre::backend::ResidentDispatchStep<'a>>, String> {
179    let mut resident_steps =
180        crate::staging_reserve::reserved_vec(steps.len(), "Vyre resident IFDS dispatch step")?;
181    for &(program, resources, grid_override) in steps {
182        resident_steps.push(vyre::backend::ResidentDispatchStep {
183            program,
184            resources,
185            grid_override,
186            workgroup_override: None,
187        });
188    }
189    Ok(resident_steps)
190}
191
192fn resident_read_ranges_from<'a>(
193    read_ranges: &'a [(&'a vyre::backend::Resource, usize, usize)],
194) -> Result<Vec<vyre::backend::ResidentReadRange<'a>>, String> {
195    let mut resident_read_ranges =
196        crate::staging_reserve::reserved_vec(read_ranges.len(), "Vyre resident IFDS read range")?;
197    for &(resource, byte_offset, byte_len) in read_ranges {
198        resident_read_ranges.push(vyre::backend::ResidentReadRange {
199            resource,
200            byte_offset,
201            byte_len,
202        });
203    }
204    Ok(resident_read_ranges)
205}
206
207#[cfg(test)]
208#[path = "ifds_vyre_resident_tests.rs"]
209mod tests;