use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use vyre::backend::{
private, BackendError, CompiledPipeline, DeviceBuffer, DispatchConfig, HostShimBuffer,
Resource, TimedDispatchResult,
};
use vyre::ir::{OpId, Program};
use vyre::VyreBackend;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AllocRecord {
pub handle: u64,
pub byte_len: usize,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct UploadRecord {
pub handle: u64,
pub byte_len: usize,
pub offset: Option<usize>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FreeRecord {
pub handle: u64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DispatchRecord {
pub resource_handles: Vec<u64>,
pub grid_override: Option<[u32; 3]>,
pub program_op_id: Option<String>,
}
#[derive(Clone, Debug)]
pub enum InjectedFailure {
DeviceOutOfMemory { requested: u64, available: u64 },
UnsupportedFeature { name: String },
PoisonedLock { lock_error: String },
InvalidProgram { fix: String },
DispatchFailed { code: Option<i32>, message: String },
KernelCompileFailed { compiler_message: String },
}
impl InjectedFailure {
fn into_backend_error(self, backend: &'static str) -> BackendError {
match self {
Self::DeviceOutOfMemory {
requested,
available,
} => BackendError::DeviceOutOfMemory {
requested,
available,
},
Self::UnsupportedFeature { name } => BackendError::UnsupportedFeature {
name,
backend: backend.to_string(),
},
Self::PoisonedLock { lock_error } => BackendError::PoisonedLock { lock_error },
Self::InvalidProgram { fix } => BackendError::InvalidProgram { fix },
Self::DispatchFailed { code, message } => {
BackendError::DispatchFailed { code, message }
}
Self::KernelCompileFailed { compiler_message } => {
BackendError::KernelCompileFailed {
backend: backend.to_string(),
compiler_message,
}
}
}
}
}
#[derive(Default)]
struct FailureSlots {
allocate: Option<InjectedFailure>,
upload: Option<InjectedFailure>,
download: Option<InjectedFailure>,
free: Option<InjectedFailure>,
dispatch: Option<InjectedFailure>,
compile: Option<InjectedFailure>,
}
#[derive(Debug)]
struct ResidentAllocation {
bytes: Vec<u8>,
}
pub struct FakeResidentBackend {
id: &'static str,
version: &'static str,
next_handle: AtomicU64,
supported_ops: HashSet<OpId>,
allocations: Mutex<HashMap<u64, ResidentAllocation>>,
alloc_records: Mutex<Vec<AllocRecord>>,
upload_records: Mutex<Vec<UploadRecord>>,
free_records: Mutex<Vec<FreeRecord>>,
dispatch_records: Mutex<Vec<DispatchRecord>>,
failures: Mutex<FailureSlots>,
}
impl Default for FakeResidentBackend {
fn default() -> Self {
Self::new()
}
}
impl FakeResidentBackend {
pub fn new() -> Self {
Self {
id: "fake_resident",
version: "test-harness-v2",
next_handle: AtomicU64::new(1),
supported_ops: HashSet::new(),
allocations: Mutex::new(HashMap::new()),
alloc_records: Mutex::new(Vec::new()),
upload_records: Mutex::new(Vec::new()),
free_records: Mutex::new(Vec::new()),
dispatch_records: Mutex::new(Vec::new()),
failures: Mutex::new(FailureSlots::default()),
}
}
pub fn with_id(id: &'static str) -> Self {
Self { id, ..Self::new() }
}
pub fn with_supported_ops(mut self, ops: HashSet<OpId>) -> Self {
self.supported_ops = ops;
self
}
pub fn alloc_count(&self) -> usize {
self.alloc_records.lock().map(|records| records.len()).unwrap_or(0)
}
pub fn upload_count(&self) -> usize {
self.upload_records
.lock()
.map(|records| records.len())
.unwrap_or(0)
}
pub fn free_count(&self) -> usize {
self.free_records.lock().map(|records| records.len()).unwrap_or(0)
}
pub fn dispatch_count(&self) -> usize {
self.dispatch_records
.lock()
.map(|records| records.len())
.unwrap_or(0)
}
pub fn alive_resources(&self) -> Vec<u64> {
let mut handles = self
.allocations
.lock()
.map(|allocations| allocations.keys().copied().collect::<Vec<_>>())
.unwrap_or_default();
handles.sort_unstable();
handles
}
pub fn take_allocs(&self) -> Vec<AllocRecord> {
take_records(&self.alloc_records)
}
pub fn take_uploads(&self) -> Vec<UploadRecord> {
take_records(&self.upload_records)
}
pub fn take_frees(&self) -> Vec<FreeRecord> {
take_records(&self.free_records)
}
pub fn take_dispatches(&self) -> Vec<DispatchRecord> {
take_records(&self.dispatch_records)
}
pub fn inject_next_allocate(&self, failure: InjectedFailure) {
if let Ok(mut failures) = self.failures.lock() {
failures.allocate = Some(failure);
}
}
pub fn inject_next_upload(&self, failure: InjectedFailure) {
if let Ok(mut failures) = self.failures.lock() {
failures.upload = Some(failure);
}
}
pub fn inject_next_download(&self, failure: InjectedFailure) {
if let Ok(mut failures) = self.failures.lock() {
failures.download = Some(failure);
}
}
pub fn inject_next_free(&self, failure: InjectedFailure) {
if let Ok(mut failures) = self.failures.lock() {
failures.free = Some(failure);
}
}
pub fn inject_next_dispatch(&self, failure: InjectedFailure) {
if let Ok(mut failures) = self.failures.lock() {
failures.dispatch = Some(failure);
}
}
pub fn inject_next_compile(&self, failure: InjectedFailure) {
if let Ok(mut failures) = self.failures.lock() {
failures.compile = Some(failure);
}
}
fn take_failure(
&self,
select: impl FnOnce(&mut FailureSlots) -> &mut Option<InjectedFailure>,
) -> Result<Option<InjectedFailure>, BackendError> {
let mut failures = lock(&self.failures, "fake backend failure slots")?;
Ok(select(&mut failures).take())
}
fn validate_range(
&self,
handle: u64,
allocation_len: usize,
offset: usize,
byte_len: usize,
operation: &str,
) -> Result<std::ops::Range<usize>, BackendError> {
let end = offset.checked_add(byte_len).ok_or_else(|| {
BackendError::InvalidProgram {
fix: format!(
"Fix: {operation} range offset {offset} plus length {byte_len} overflows usize."
),
}
})?;
if end > allocation_len {
return Err(BackendError::InvalidProgram {
fix: format!(
"Fix: {operation} range {offset}..{end} exceeds allocation {handle} length {allocation_len}."
),
});
}
Ok(offset..end)
}
fn resident_handle(resource: &Resource) -> Result<u64, BackendError> {
match resource {
Resource::Resident(handle) => Ok(*handle),
Resource::Borrowed(_) => Err(BackendError::InvalidProgram {
fix: "Fix: fake resident backend expected Resource::Resident, got Resource::Borrowed."
.to_string(),
}),
}
}
}
impl private::Sealed for FakeResidentBackend {}
impl VyreBackend for FakeResidentBackend {
fn id(&self) -> &'static str {
self.id
}
fn version(&self) -> &'static str {
self.version
}
fn supported_ops(&self) -> &HashSet<OpId> {
&self.supported_ops
}
fn dispatch(
&self,
_program: &Program,
_inputs: &[Vec<u8>],
_config: &DispatchConfig,
) -> Result<Vec<Vec<u8>>, BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.dispatch)? {
return Err(failure.into_backend_error(self.id));
}
Err(BackendError::UnsupportedFeature {
name: "borrowed fake backend dispatch".to_string(),
backend: self.id.to_string(),
})
}
fn allocate_resident(&self, byte_len: usize) -> Result<Resource, BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.allocate)? {
return Err(failure.into_backend_error(self.id));
}
let handle = self.next_handle.fetch_add(1, Ordering::Relaxed);
let mut bytes = Vec::new();
bytes.try_reserve_exact(byte_len).map_err(|error| {
BackendError::InvalidProgram {
fix: format!(
"Fix: fake resident backend could not reserve {byte_len} byte(s): {error}."
),
}
})?;
bytes.resize(byte_len, 0);
lock(&self.allocations, "fake backend allocations")?
.insert(handle, ResidentAllocation { bytes });
lock(&self.alloc_records, "fake backend allocation records")?.push(AllocRecord {
handle,
byte_len,
});
Ok(Resource::Resident(handle))
}
fn upload_resident(&self, resource: &Resource, bytes: &[u8]) -> Result<(), BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
return Err(failure.into_backend_error(self.id));
}
self.upload_resident_at_no_injection(resource, 0, bytes, None)
}
fn upload_resident_many(&self, uploads: &[(&Resource, &[u8])]) -> Result<(), BackendError> {
if uploads.is_empty() {
return Ok(());
}
if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
return Err(failure.into_backend_error(self.id));
}
for &(resource, bytes) in uploads {
self.upload_resident_at_no_injection(resource, 0, bytes, None)?;
}
Ok(())
}
fn upload_resident_at(
&self,
resource: &Resource,
dst_offset_bytes: usize,
bytes: &[u8],
) -> Result<(), BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
return Err(failure.into_backend_error(self.id));
}
self.upload_resident_at_no_injection(resource, dst_offset_bytes, bytes, Some(dst_offset_bytes))
}
fn upload_resident_at_many(
&self,
uploads: &[(&Resource, usize, &[u8])],
) -> Result<(), BackendError> {
if uploads.is_empty() {
return Ok(());
}
if let Some(failure) = self.take_failure(|failures| &mut failures.upload)? {
return Err(failure.into_backend_error(self.id));
}
for &(resource, offset, bytes) in uploads {
self.upload_resident_at_no_injection(resource, offset, bytes, Some(offset))?;
}
Ok(())
}
fn download_resident_into(
&self,
resource: &Resource,
out: &mut Vec<u8>,
) -> Result<(), BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.download)? {
return Err(failure.into_backend_error(self.id));
}
let handle = Self::resident_handle(resource)?;
let allocations = lock(&self.allocations, "fake backend allocations")?;
let Some(allocation) = allocations.get(&handle) else {
return Err(BackendError::InvalidProgram {
fix: format!("Fix: fake resident backend detected use-after-free resource {handle}."),
});
};
out.clear();
out.try_reserve_exact(allocation.bytes.len())
.map_err(|error| BackendError::InvalidProgram {
fix: format!(
"Fix: fake resident backend could not reserve download buffer: {error}."
),
})?;
out.extend_from_slice(&allocation.bytes);
Ok(())
}
fn download_resident_range_into(
&self,
resource: &Resource,
byte_offset: usize,
byte_len: usize,
out: &mut Vec<u8>,
) -> Result<(), BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.download)? {
return Err(failure.into_backend_error(self.id));
}
let handle = Self::resident_handle(resource)?;
let allocations = lock(&self.allocations, "fake backend allocations")?;
let Some(allocation) = allocations.get(&handle) else {
return Err(BackendError::InvalidProgram {
fix: format!("Fix: fake resident backend detected use-after-free resource {handle}."),
});
};
let range = self.validate_range(
handle,
allocation.bytes.len(),
byte_offset,
byte_len,
"resident download",
)?;
out.clear();
out.try_reserve_exact(byte_len)
.map_err(|error| BackendError::InvalidProgram {
fix: format!(
"Fix: fake resident backend could not reserve ranged download buffer: {error}."
),
})?;
out.extend_from_slice(&allocation.bytes[range]);
Ok(())
}
fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.free)? {
return Err(failure.into_backend_error(self.id));
}
let handle = Self::resident_handle(&resource)?;
let removed = lock(&self.allocations, "fake backend allocations")?.remove(&handle);
if removed.is_none() {
return Err(BackendError::InvalidProgram {
fix: format!("Fix: fake resident backend detected double-free resource {handle}."),
});
}
lock(&self.free_records, "fake backend free records")?.push(FreeRecord { handle });
Ok(())
}
fn dispatch_resident_timed(
&self,
program: &Program,
resources: &[Resource],
config: &DispatchConfig,
) -> Result<TimedDispatchResult, BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.dispatch)? {
return Err(failure.into_backend_error(self.id));
}
let allocations = lock(&self.allocations, "fake backend allocations")?;
let mut resource_handles = Vec::new();
resource_handles
.try_reserve_exact(resources.len())
.map_err(|error| BackendError::InvalidProgram {
fix: format!(
"Fix: fake resident backend could not reserve dispatch record handles: {error}."
),
})?;
for resource in resources {
let handle = Self::resident_handle(resource)?;
if !allocations.contains_key(&handle) {
return Err(BackendError::InvalidProgram {
fix: format!(
"Fix: fake resident backend detected use-after-free resource {handle}."
),
});
}
resource_handles.push(handle);
}
drop(allocations);
lock(&self.dispatch_records, "fake backend dispatch records")?.push(DispatchRecord {
resource_handles,
grid_override: config.grid_override,
program_op_id: program.entry_op_id.as_ref().map(ToString::to_string),
});
Ok(TimedDispatchResult {
outputs: Vec::new(),
wall_ns: 0,
device_ns: Some(0),
enqueue_ns: Some(0),
wait_ns: Some(0),
})
}
fn compile_native(
&self,
_program: &Program,
_config: &DispatchConfig,
) -> Result<Option<Arc<dyn CompiledPipeline>>, BackendError> {
if let Some(failure) = self.take_failure(|failures| &mut failures.compile)? {
return Err(failure.into_backend_error(self.id));
}
Ok(None)
}
fn compile_native_shared(
&self,
_program: Arc<Program>,
config: &DispatchConfig,
) -> Result<Option<Arc<dyn CompiledPipeline>>, BackendError> {
self.compile_native(&Program::default(), config)
}
fn prepare(&self) -> Result<(), BackendError> {
Ok(())
}
fn flush(&self) -> Result<(), BackendError> {
Ok(())
}
fn shutdown(&self) -> Result<(), BackendError> {
Ok(())
}
fn try_recover(&self) -> Result<(), BackendError> {
Ok(())
}
fn allocate_device_buffer(
&self,
byte_len: usize,
) -> Result<Box<dyn DeviceBuffer>, BackendError> {
Ok(HostShimBuffer::allocate(self.id, byte_len))
}
fn upload_device_buffer(
&self,
buffer: &mut dyn DeviceBuffer,
bytes: &[u8],
) -> Result<(), BackendError> {
let Some(host) = buffer
.as_any_mut()
.downcast_mut::<HostShimBuffer>()
else {
return Err(BackendError::InvalidProgram {
fix: "Fix: fake backend can only upload to HostShimBuffer device buffers."
.to_string(),
});
};
if host.byte_len() != bytes.len() {
return Err(BackendError::InvalidProgram {
fix: format!(
"Fix: fake backend device buffer upload expected {} byte(s), got {}.",
host.byte_len(),
bytes.len()
),
});
}
host.as_mut_slice().copy_from_slice(bytes);
Ok(())
}
fn download_device_buffer(&self, buffer: &dyn DeviceBuffer) -> Result<Vec<u8>, BackendError> {
let Some(host) = buffer
.as_any()
.downcast_ref::<HostShimBuffer>()
else {
return Err(BackendError::InvalidProgram {
fix: "Fix: fake backend can only download HostShimBuffer device buffers."
.to_string(),
});
};
Ok(host.as_slice().to_vec())
}
fn free_device_buffer(&self, _buffer: Box<dyn DeviceBuffer>) -> Result<(), BackendError> {
Ok(())
}
}
impl FakeResidentBackend {
fn upload_resident_at_no_injection(
&self,
resource: &Resource,
dst_offset_bytes: usize,
bytes: &[u8],
record_offset: Option<usize>,
) -> Result<(), BackendError> {
let handle = Self::resident_handle(resource)?;
let mut allocations = lock(&self.allocations, "fake backend allocations")?;
let Some(allocation) = allocations.get_mut(&handle) else {
return Err(BackendError::InvalidProgram {
fix: format!("Fix: fake resident backend detected use-after-free resource {handle}."),
});
};
let range = self.validate_range(
handle,
allocation.bytes.len(),
dst_offset_bytes,
bytes.len(),
"resident upload",
)?;
allocation.bytes[range].copy_from_slice(bytes);
lock(&self.upload_records, "fake backend upload records")?.push(UploadRecord {
handle,
byte_len: bytes.len(),
offset: record_offset,
});
Ok(())
}
}
fn lock<'a, T>(mutex: &'a Mutex<T>, label: &str) -> Result<MutexGuard<'a, T>, BackendError> {
mutex.lock().map_err(|error| BackendError::PoisonedLock {
lock_error: format!("{label}: {error}"),
})
}
fn take_records<T>(mutex: &Mutex<Vec<T>>) -> Vec<T> {
match mutex.lock() {
Ok(mut records) => std::mem::take(&mut *records),
Err(_) => Vec::new(),
}
}