1use std::collections::{HashMap, HashSet};
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, Mutex, MutexGuard};
6
7use vyre::backend::{
8 private, BackendError, CompiledPipeline, DeviceBuffer, DispatchConfig, HostShimBuffer,
9 Resource, TimedDispatchResult,
10};
11use vyre::ir::{OpId, Program};
12use vyre::VyreBackend;
13
14#[derive(Clone, Debug, Eq, PartialEq)]
15pub struct AllocRecord {
16 pub handle: u64,
17 pub byte_len: usize,
18}
19
20#[derive(Clone, Debug, Eq, PartialEq)]
21pub struct UploadRecord {
22 pub handle: u64,
23 pub byte_len: usize,
24 pub offset: Option<usize>,
25}
26
27#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct FreeRecord {
29 pub handle: u64,
30}
31
32#[derive(Clone, Debug, Eq, PartialEq)]
33pub struct DispatchRecord {
34 pub resource_handles: Vec<u64>,
35 pub grid_override: Option<[u32; 3]>,
36 pub program_op_id: Option<String>,
37}
38
39#[derive(Clone, Debug)]
40pub enum InjectedFailure {
41 DeviceOutOfMemory { requested: u64, available: u64 },
42 UnsupportedFeature { name: String },
43 PoisonedLock { lock_error: String },
44 InvalidProgram { fix: String },
45 DispatchFailed { code: Option<i32>, message: String },
46 KernelCompileFailed { compiler_message: String },
47}
48
49impl InjectedFailure {
50 fn into_backend_error(self, backend: &'static str) -> BackendError {
51 match self {
52 Self::DeviceOutOfMemory {
53 requested,
54 available,
55 } => BackendError::DeviceOutOfMemory {
56 requested,
57 available,
58 },
59 Self::UnsupportedFeature { name } => BackendError::UnsupportedFeature {
60 name,
61 backend: backend.to_string(),
62 },
63 Self::PoisonedLock { lock_error } => BackendError::PoisonedLock { lock_error },
64 Self::InvalidProgram { fix } => BackendError::InvalidProgram { fix },
65 Self::DispatchFailed { code, message } => {
66 BackendError::DispatchFailed { code, message }
67 }
68 Self::KernelCompileFailed { compiler_message } => {
69 BackendError::KernelCompileFailed {
70 backend: backend.to_string(),
71 compiler_message,
72 }
73 }
74 }
75 }
76}
77
78#[derive(Default)]
79struct FailureSlots {
80 allocate: Option<InjectedFailure>,
81 upload: Option<InjectedFailure>,
82 download: Option<InjectedFailure>,
83 free: Option<InjectedFailure>,
84 dispatch: Option<InjectedFailure>,
85 compile: Option<InjectedFailure>,
86}
87
88#[derive(Debug)]
89struct ResidentAllocation {
90 bytes: Vec<u8>,
91}
92
93pub struct FakeResidentBackend {
94 id: &'static str,
95 version: &'static str,
96 next_handle: AtomicU64,
97 supported_ops: HashSet<OpId>,
98 allocations: Mutex<HashMap<u64, ResidentAllocation>>,
99 alloc_records: Mutex<Vec<AllocRecord>>,
100 upload_records: Mutex<Vec<UploadRecord>>,
101 free_records: Mutex<Vec<FreeRecord>>,
102 dispatch_records: Mutex<Vec<DispatchRecord>>,
103 failures: Mutex<FailureSlots>,
104}
105
106impl Default for FakeResidentBackend {
107 fn default() -> Self {
108 Self::new()
109 }
110}
111
112impl FakeResidentBackend {
113 pub fn new() -> Self {
114 Self {
115 id: "fake_resident",
116 version: "test-harness-v2",
117 next_handle: AtomicU64::new(1),
118 supported_ops: HashSet::new(),
119 allocations: Mutex::new(HashMap::new()),
120 alloc_records: Mutex::new(Vec::new()),
121 upload_records: Mutex::new(Vec::new()),
122 free_records: Mutex::new(Vec::new()),
123 dispatch_records: Mutex::new(Vec::new()),
124 failures: Mutex::new(FailureSlots::default()),
125 }
126 }
127
128 pub fn with_id(id: &'static str) -> Self {
129 Self { id, ..Self::new() }
130 }
131
132 pub fn with_supported_ops(mut self, ops: HashSet<OpId>) -> Self {
133 self.supported_ops = ops;
134 self
135 }
136
137 pub fn alloc_count(&self) -> usize {
138 self.alloc_records.lock().map(|records| records.len()).unwrap_or(0)
139 }
140
141 pub fn upload_count(&self) -> usize {
142 self.upload_records
143 .lock()
144 .map(|records| records.len())
145 .unwrap_or(0)
146 }
147
148 pub fn free_count(&self) -> usize {
149 self.free_records.lock().map(|records| records.len()).unwrap_or(0)
150 }
151
152 pub fn dispatch_count(&self) -> usize {
153 self.dispatch_records
154 .lock()
155 .map(|records| records.len())
156 .unwrap_or(0)
157 }
158
159 pub fn alive_resources(&self) -> Vec<u64> {
160 let mut handles = self
161 .allocations
162 .lock()
163 .map(|allocations| allocations.keys().copied().collect::<Vec<_>>())
164 .unwrap_or_default();
165 handles.sort_unstable();
166 handles
167 }
168
169 pub fn take_allocs(&self) -> Vec<AllocRecord> {
170 take_records(&self.alloc_records)
171 }
172
173 pub fn take_uploads(&self) -> Vec<UploadRecord> {
174 take_records(&self.upload_records)
175 }
176
177 pub fn take_frees(&self) -> Vec<FreeRecord> {
178 take_records(&self.free_records)
179 }
180
181 pub fn take_dispatches(&self) -> Vec<DispatchRecord> {
182 take_records(&self.dispatch_records)
183 }
184
185 pub fn inject_next_allocate(&self, failure: InjectedFailure) {
186 if let Ok(mut failures) = self.failures.lock() {
187 failures.allocate = Some(failure);
188 }
189 }
190
191 pub fn inject_next_upload(&self, failure: InjectedFailure) {
192 if let Ok(mut failures) = self.failures.lock() {
193 failures.upload = Some(failure);
194 }
195 }
196
197 pub fn inject_next_download(&self, failure: InjectedFailure) {
198 if let Ok(mut failures) = self.failures.lock() {
199 failures.download = Some(failure);
200 }
201 }
202
203 pub fn inject_next_free(&self, failure: InjectedFailure) {
204 if let Ok(mut failures) = self.failures.lock() {
205 failures.free = Some(failure);
206 }
207 }
208
209 pub fn inject_next_dispatch(&self, failure: InjectedFailure) {
210 if let Ok(mut failures) = self.failures.lock() {
211 failures.dispatch = Some(failure);
212 }
213 }
214
215 pub fn inject_next_compile(&self, failure: InjectedFailure) {
216 if let Ok(mut failures) = self.failures.lock() {
217 failures.compile = Some(failure);
218 }
219 }
220
221 fn take_failure(
222 &self,
223 select: impl FnOnce(&mut FailureSlots) -> &mut Option<InjectedFailure>,
224 ) -> Result<Option<InjectedFailure>, BackendError> {
225 let mut failures = lock(&self.failures, "fake backend failure slots")?;
226 Ok(select(&mut failures).take())
227 }
228
229 fn validate_range(
230 &self,
231 handle: u64,
232 allocation_len: usize,
233 offset: usize,
234 byte_len: usize,
235 operation: &str,
236 ) -> Result<std::ops::Range<usize>, BackendError> {
237 let end = offset.checked_add(byte_len).ok_or_else(|| {
238 BackendError::InvalidProgram {
239 fix: format!(
240 "Fix: {operation} range offset {offset} plus length {byte_len} overflows usize."
241 ),
242 }
243 })?;
244 if end > allocation_len {
245 return Err(BackendError::InvalidProgram {
246 fix: format!(
247 "Fix: {operation} range {offset}..{end} exceeds allocation {handle} length {allocation_len}."
248 ),
249 });
250 }
251 Ok(offset..end)
252 }
253
254 fn resident_handle(resource: &Resource) -> Result<u64, BackendError> {
255 match resource {
256 Resource::Resident(handle) => Ok(*handle),
257 Resource::Borrowed(_) => Err(BackendError::InvalidProgram {
258 fix: "Fix: fake resident backend expected Resource::Resident, got Resource::Borrowed."
259 .to_string(),
260 }),
261 }
262 }
263}
264
265impl private::Sealed for FakeResidentBackend {}
266
267impl VyreBackend for FakeResidentBackend {
268 fn id(&self) -> &'static str {
269 self.id
270 }
271
272 fn version(&self) -> &'static str {
273 self.version
274 }
275
276 fn supported_ops(&self) -> &HashSet<OpId> {
277 &self.supported_ops
278 }
279
280 fn dispatch(
281 &self,
282 _program: &Program,
283 _inputs: &[Vec<u8>],
284 _config: &DispatchConfig,
285 ) -> Result<Vec<Vec<u8>>, BackendError> {
286 if let Some(failure) = self.take_failure(|failures| &mut failures.dispatch)? {
287 return Err(failure.into_backend_error(self.id));
288 }
289 Err(BackendError::UnsupportedFeature {
290 name: "borrowed fake backend dispatch".to_string(),
291 backend: self.id.to_string(),
292 })
293 }
294
295 fn allocate_resident(&self, byte_len: usize) -> Result<Resource, BackendError> {
296 if let Some(failure) = self.take_failure(|failures| &mut failures.allocate)? {
297 return Err(failure.into_backend_error(self.id));
298 }
299
300 let handle = self.next_handle.fetch_add(1, Ordering::Relaxed);
301 let mut bytes = Vec::new();
302 bytes.try_reserve_exact(byte_len).map_err(|error| {
303 BackendError::InvalidProgram {
304 fix: format!(
305 "Fix: fake resident backend could not reserve {byte_len} byte(s): {error}."
306 ),
307 }
308 })?;
309 bytes.resize(byte_len, 0);
310
311 lock(&self.allocations, "fake backend allocations")?
312 .insert(handle, ResidentAllocation { bytes });
313 lock(&self.alloc_records, "fake backend allocation records")?.push(AllocRecord {
314 handle,
315 byte_len,
316 });
317 Ok(Resource::Resident(handle))
318 }
319
320 fn upload_resident(&self, resource: &Resource, bytes: &[u8]) -> Result<(), BackendError> {
321 if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
322 return Err(failure.into_backend_error(self.id));
323 }
324 self.upload_resident_at_no_injection(resource, 0, bytes, None)
325 }
326
327 fn upload_resident_many(&self, uploads: &[(&Resource, &[u8])]) -> Result<(), BackendError> {
328 if uploads.is_empty() {
329 return Ok(());
330 }
331 if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
332 return Err(failure.into_backend_error(self.id));
333 }
334 for &(resource, bytes) in uploads {
335 self.upload_resident_at_no_injection(resource, 0, bytes, None)?;
336 }
337 Ok(())
338 }
339
340 fn upload_resident_at(
341 &self,
342 resource: &Resource,
343 dst_offset_bytes: usize,
344 bytes: &[u8],
345 ) -> Result<(), BackendError> {
346 if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
347 return Err(failure.into_backend_error(self.id));
348 }
349 self.upload_resident_at_no_injection(resource, dst_offset_bytes, bytes, Some(dst_offset_bytes))
350 }
351
352 fn upload_resident_at_many(
353 &self,
354 uploads: &[(&Resource, usize, &[u8])],
355 ) -> Result<(), BackendError> {
356 if uploads.is_empty() {
357 return Ok(());
358 }
359 if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
360 return Err(failure.into_backend_error(self.id));
361 }
362 for &(resource, offset, bytes) in uploads {
363 self.upload_resident_at_no_injection(resource, offset, bytes, Some(offset))?;
364 }
365 Ok(())
366 }
367
368 fn download_resident_into(
369 &self,
370 resource: &Resource,
371 out: &mut Vec<u8>,
372 ) -> Result<(), BackendError> {
373 if let Some(failure) = self.take_failure(|failures| &mut failures.download)? {
374 return Err(failure.into_backend_error(self.id));
375 }
376 let handle = Self::resident_handle(resource)?;
377 let allocations = lock(&self.allocations, "fake backend allocations")?;
378 let Some(allocation) = allocations.get(&handle) else {
379 return Err(BackendError::InvalidProgram {
380 fix: format!("Fix: fake resident backend detected use-after-free resource {handle}."),
381 });
382 };
383 out.clear();
384 out.try_reserve_exact(allocation.bytes.len())
385 .map_err(|error| BackendError::InvalidProgram {
386 fix: format!(
387 "Fix: fake resident backend could not reserve download buffer: {error}."
388 ),
389 })?;
390 out.extend_from_slice(&allocation.bytes);
391 Ok(())
392 }
393
394 fn download_resident_range_into(
395 &self,
396 resource: &Resource,
397 byte_offset: usize,
398 byte_len: usize,
399 out: &mut Vec<u8>,
400 ) -> Result<(), BackendError> {
401 if let Some(failure) = self.take_failure(|failures| &mut failures.download)? {
402 return Err(failure.into_backend_error(self.id));
403 }
404 let handle = Self::resident_handle(resource)?;
405 let allocations = lock(&self.allocations, "fake backend allocations")?;
406 let Some(allocation) = allocations.get(&handle) else {
407 return Err(BackendError::InvalidProgram {
408 fix: format!("Fix: fake resident backend detected use-after-free resource {handle}."),
409 });
410 };
411 let range = self.validate_range(
412 handle,
413 allocation.bytes.len(),
414 byte_offset,
415 byte_len,
416 "resident download",
417 )?;
418 out.clear();
419 out.try_reserve_exact(byte_len)
420 .map_err(|error| BackendError::InvalidProgram {
421 fix: format!(
422 "Fix: fake resident backend could not reserve ranged download buffer: {error}."
423 ),
424 })?;
425 out.extend_from_slice(&allocation.bytes[range]);
426 Ok(())
427 }
428
429 fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
430 if let Some(failure) = self.take_failure(|failures| &mut failures.free)? {
431 return Err(failure.into_backend_error(self.id));
432 }
433 let handle = Self::resident_handle(&resource)?;
434 let removed = lock(&self.allocations, "fake backend allocations")?.remove(&handle);
435 if removed.is_none() {
436 return Err(BackendError::InvalidProgram {
437 fix: format!("Fix: fake resident backend detected double-free resource {handle}."),
438 });
439 }
440 lock(&self.free_records, "fake backend free records")?.push(FreeRecord { handle });
441 Ok(())
442 }
443
444 fn dispatch_resident_timed(
445 &self,
446 program: &Program,
447 resources: &[Resource],
448 config: &DispatchConfig,
449 ) -> Result<TimedDispatchResult, BackendError> {
450 if let Some(failure) = self.take_failure(|failures| &mut failures.dispatch)? {
451 return Err(failure.into_backend_error(self.id));
452 }
453 let allocations = lock(&self.allocations, "fake backend allocations")?;
454 let mut resource_handles = Vec::new();
455 resource_handles
456 .try_reserve_exact(resources.len())
457 .map_err(|error| BackendError::InvalidProgram {
458 fix: format!(
459 "Fix: fake resident backend could not reserve dispatch record handles: {error}."
460 ),
461 })?;
462 for resource in resources {
463 let handle = Self::resident_handle(resource)?;
464 if !allocations.contains_key(&handle) {
465 return Err(BackendError::InvalidProgram {
466 fix: format!(
467 "Fix: fake resident backend detected use-after-free resource {handle}."
468 ),
469 });
470 }
471 resource_handles.push(handle);
472 }
473 drop(allocations);
474
475 lock(&self.dispatch_records, "fake backend dispatch records")?.push(DispatchRecord {
476 resource_handles,
477 grid_override: config.grid_override,
478 program_op_id: program.entry_op_id.as_ref().map(ToString::to_string),
479 });
480 Ok(TimedDispatchResult {
481 outputs: Vec::new(),
482 wall_ns: 0,
483 device_ns: Some(0),
484 enqueue_ns: Some(0),
485 wait_ns: Some(0),
486 })
487 }
488
489 fn compile_native(
490 &self,
491 _program: &Program,
492 _config: &DispatchConfig,
493 ) -> Result<Option<Arc<dyn CompiledPipeline>>, BackendError> {
494 if let Some(failure) = self.take_failure(|failures| &mut failures.compile)? {
495 return Err(failure.into_backend_error(self.id));
496 }
497 Ok(None)
498 }
499
500 fn compile_native_shared(
501 &self,
502 _program: Arc<Program>,
503 config: &DispatchConfig,
504 ) -> Result<Option<Arc<dyn CompiledPipeline>>, BackendError> {
505 self.compile_native(&Program::default(), config)
506 }
507
508 fn prepare(&self) -> Result<(), BackendError> {
509 Ok(())
510 }
511
512 fn flush(&self) -> Result<(), BackendError> {
513 Ok(())
514 }
515
516 fn shutdown(&self) -> Result<(), BackendError> {
517 Ok(())
518 }
519
520 fn try_recover(&self) -> Result<(), BackendError> {
521 Ok(())
522 }
523
524 fn allocate_device_buffer(
525 &self,
526 byte_len: usize,
527 ) -> Result<Box<dyn DeviceBuffer>, BackendError> {
528 Ok(HostShimBuffer::allocate(self.id, byte_len))
529 }
530
531 fn upload_device_buffer(
532 &self,
533 buffer: &mut dyn DeviceBuffer,
534 bytes: &[u8],
535 ) -> Result<(), BackendError> {
536 let Some(host) = buffer
537 .as_any_mut()
538 .downcast_mut::<HostShimBuffer>()
539 else {
540 return Err(BackendError::InvalidProgram {
541 fix: "Fix: fake backend can only upload to HostShimBuffer device buffers."
542 .to_string(),
543 });
544 };
545 if host.byte_len() != bytes.len() {
546 return Err(BackendError::InvalidProgram {
547 fix: format!(
548 "Fix: fake backend device buffer upload expected {} byte(s), got {}.",
549 host.byte_len(),
550 bytes.len()
551 ),
552 });
553 }
554 host.as_mut_slice().copy_from_slice(bytes);
555 Ok(())
556 }
557
558 fn download_device_buffer(&self, buffer: &dyn DeviceBuffer) -> Result<Vec<u8>, BackendError> {
559 let Some(host) = buffer
560 .as_any()
561 .downcast_ref::<HostShimBuffer>()
562 else {
563 return Err(BackendError::InvalidProgram {
564 fix: "Fix: fake backend can only download HostShimBuffer device buffers."
565 .to_string(),
566 });
567 };
568 Ok(host.as_slice().to_vec())
569 }
570
571 fn free_device_buffer(&self, _buffer: Box<dyn DeviceBuffer>) -> Result<(), BackendError> {
572 Ok(())
573 }
574}
575
576impl FakeResidentBackend {
577 fn upload_resident_at_no_injection(
578 &self,
579 resource: &Resource,
580 dst_offset_bytes: usize,
581 bytes: &[u8],
582 record_offset: Option<usize>,
583 ) -> Result<(), BackendError> {
584 let handle = Self::resident_handle(resource)?;
585 let mut allocations = lock(&self.allocations, "fake backend allocations")?;
586 let Some(allocation) = allocations.get_mut(&handle) else {
587 return Err(BackendError::InvalidProgram {
588 fix: format!("Fix: fake resident backend detected use-after-free resource {handle}."),
589 });
590 };
591 let range = self.validate_range(
592 handle,
593 allocation.bytes.len(),
594 dst_offset_bytes,
595 bytes.len(),
596 "resident upload",
597 )?;
598 allocation.bytes[range].copy_from_slice(bytes);
599 lock(&self.upload_records, "fake backend upload records")?.push(UploadRecord {
600 handle,
601 byte_len: bytes.len(),
602 offset: record_offset,
603 });
604 Ok(())
605 }
606}
607
608fn lock<'a, T>(mutex: &'a Mutex<T>, label: &str) -> Result<MutexGuard<'a, T>, BackendError> {
609 mutex.lock().map_err(|error| BackendError::PoisonedLock {
610 lock_error: format!("{label}: {error}"),
611 })
612}
613
614fn take_records<T>(mutex: &Mutex<Vec<T>>) -> Vec<T> {
615 match mutex.lock() {
616 Ok(mut records) => std::mem::take(&mut *records),
617 Err(_) => Vec::new(),
618 }
619}