use crate::{
error::{Error, Result},
BufferIdentity, IdentityKind, TensorMemory, TensorTrait,
};
use log::trace;
use num_traits::Num;
use std::{
ffi::c_void,
fmt,
marker::PhantomData,
ops::Deref,
ptr::NonNull,
sync::{Arc, Condvar, Mutex},
};
use std::{os::raw::c_int, sync::OnceLock};
pub struct PboMapping {
pub ptr: *mut u8,
pub size: usize,
}
unsafe impl Send for PboMapping {}
pub unsafe trait PboOps: Send + Sync {
fn map_buffer(&self, buffer_id: u32, size: usize) -> Result<PboMapping>;
fn unmap_buffer(&self, buffer_id: u32) -> Result<()>;
fn delete_buffer(&self, buffer_id: u32);
}
#[repr(C)]
pub struct PboOpsVtable {
ctx: *const c_void,
map_buffer_fn: unsafe extern "C" fn(
ctx: *const c_void,
buffer_id: u32,
size: usize,
out_ptr: *mut *mut u8,
out_len: *mut usize,
) -> c_int,
unmap_buffer_fn: unsafe extern "C" fn(ctx: *const c_void, buffer_id: u32) -> c_int,
}
unsafe impl Send for PboOpsVtable {}
unsafe impl Sync for PboOpsVtable {}
impl PboOpsVtable {
fn new(ops: &Arc<dyn PboOps>) -> Self {
PboOpsVtable {
ctx: ops as *const Arc<dyn PboOps> as *const c_void,
map_buffer_fn: vt_map_buffer,
unmap_buffer_fn: vt_unmap_buffer,
}
}
unsafe fn map_buffer(&self, buffer_id: u32, size: usize) -> Result<PboMapping> {
let mut ptr: *mut u8 = std::ptr::null_mut();
let mut len: usize = 0;
let rc = unsafe { (self.map_buffer_fn)(self.ctx, buffer_id, size, &mut ptr, &mut len) };
if rc != 0 {
return Err(vtable_errno_to_error(rc));
}
Ok(PboMapping { ptr, size: len })
}
unsafe fn unmap_buffer(&self, buffer_id: u32) -> Result<()> {
let rc = unsafe { (self.unmap_buffer_fn)(self.ctx, buffer_id) };
if rc != 0 {
return Err(vtable_errno_to_error(rc));
}
Ok(())
}
}
fn vtable_errno_to_error(rc: c_int) -> Error {
match rc {
-1 => Error::PboDisconnected,
n => Error::NotImplemented(format!("cross-package PBO operation failed: errno {n}")),
}
}
fn error_to_vtable_errno(e: &Error) -> c_int {
match e {
Error::PboDisconnected => -1,
_ => 1,
}
}
unsafe extern "C" fn vt_map_buffer(
ctx: *const c_void,
buffer_id: u32,
size: usize,
out_ptr: *mut *mut u8,
out_len: *mut usize,
) -> c_int {
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let ops = unsafe { &*(ctx as *const Arc<dyn PboOps>) };
ops.map_buffer(buffer_id, size)
}));
match result {
Ok(Ok(mapping)) => {
unsafe {
*out_ptr = mapping.ptr;
*out_len = mapping.size;
}
0
}
Ok(Err(e)) => error_to_vtable_errno(&e),
Err(_) => 1, }
}
unsafe extern "C" fn vt_unmap_buffer(ctx: *const c_void, buffer_id: u32) -> c_int {
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let ops = unsafe { &*(ctx as *const Arc<dyn PboOps>) };
ops.unmap_buffer(buffer_id)
}));
match result {
Ok(Ok(())) => 0,
Ok(Err(e)) => error_to_vtable_errno(&e),
Err(_) => 1,
}
}
struct ImportedPboOps {
vtable: NonNull<PboOpsVtable>,
}
unsafe impl Send for ImportedPboOps {}
unsafe impl Sync for ImportedPboOps {}
unsafe impl PboOps for ImportedPboOps {
fn map_buffer(&self, buffer_id: u32, size: usize) -> Result<PboMapping> {
unsafe { self.vtable.as_ref().map_buffer(buffer_id, size) }
}
fn unmap_buffer(&self, buffer_id: u32) -> Result<()> {
unsafe { self.vtable.as_ref().unmap_buffer(buffer_id) }
}
fn delete_buffer(&self, _buffer_id: u32) {
}
}
pub(crate) unsafe fn import_pbo_ops(vtable_ptr: *const c_void) -> Result<Arc<dyn PboOps>> {
let Some(vtable) = NonNull::new(vtable_ptr as *mut PboOpsVtable) else {
return Err(Error::InvalidArgument(
"PBO descriptor carries no ops vtable".into(),
));
};
Ok(Arc::new(ImportedPboOps { vtable }))
}
struct PboHandle {
ops: Arc<dyn PboOps>,
buffer_id: u32,
size: usize,
map_state: Mutex<MapState>,
map_cv: Condvar,
abi_vtable: OnceLock<PboOpsVtable>,
}
enum MapState {
Unmapped,
Mapping,
Exclusive,
Shared {
readers: usize,
ptr: PboPtr,
},
Unmapping,
}
impl PboHandle {
fn acquire_map(&self, writable: bool) -> Result<PboPtr> {
loop {
let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
while matches!(*guard, MapState::Mapping | MapState::Unmapping) {
guard = self
.map_cv
.wait(guard)
.expect("PBO map state mutex poisoned");
}
match &mut *guard {
MapState::Shared { readers, ptr } if !writable => {
*readers += 1;
return Ok(PboPtr(ptr.0));
}
MapState::Shared { .. } | MapState::Exclusive => return Err(Error::PboMapped),
MapState::Mapping | MapState::Unmapping => continue,
MapState::Unmapped => *guard = MapState::Mapping,
}
drop(guard);
let mapped = self.ops.map_buffer(self.buffer_id, self.size);
let base = match mapped {
Ok(mapping) => NonNull::new(mapping.ptr as *mut c_void),
Err(e) => {
self.finish_map(None, writable);
return Err(e);
}
};
let Some(base) = base else {
let _ = self.ops.unmap_buffer(self.buffer_id);
self.finish_map(None, writable);
return Err(Error::InvalidSize(self.size));
};
self.finish_map(Some(base), writable);
return Ok(PboPtr(base));
}
}
fn finish_map(&self, base: Option<NonNull<c_void>>, writable: bool) {
let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
*guard = match base {
None => MapState::Unmapped,
Some(_) if writable => MapState::Exclusive,
Some(base) => MapState::Shared {
readers: 1,
ptr: PboPtr(base),
},
};
self.map_cv.notify_all();
}
fn release_map(&self) {
let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
let last_holder = match &mut *guard {
MapState::Shared { readers, .. } => {
*readers = readers.saturating_sub(1);
*readers == 0
}
MapState::Exclusive => true,
MapState::Unmapped | MapState::Mapping | MapState::Unmapping => false,
};
if !last_holder {
return;
}
*guard = MapState::Unmapping;
drop(guard);
trace!("Unmapping PBO buffer_id={}", self.buffer_id);
if let Err(e) = self.ops.unmap_buffer(self.buffer_id) {
log::warn!("Failed to unmap PBO buffer {}: {e}", self.buffer_id);
}
let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
*guard = MapState::Unmapped;
self.map_cv.notify_all();
}
fn is_mapped(&self) -> bool {
!matches!(
*self.map_state.lock().expect("PBO map state mutex poisoned"),
MapState::Unmapped
)
}
}
impl Drop for PboHandle {
fn drop(&mut self) {
self.ops.delete_buffer(self.buffer_id);
}
}
pub struct PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
pub name: String,
pub shape: Vec<usize>,
handle: Arc<PboHandle>,
identity: BufferIdentity,
pub(crate) view_offset: usize,
_marker: PhantomData<T>,
}
impl<T> fmt::Debug for PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PboTensor")
.field("name", &self.name)
.field("shape", &self.shape)
.field("buffer_id", &self.handle.buffer_id)
.field("size", &self.handle.size)
.finish()
}
}
unsafe impl<T> Send for PboTensor<T> where T: Num + Clone + fmt::Debug + Send + Sync {}
unsafe impl<T> Sync for PboTensor<T> where T: Num + Clone + fmt::Debug + Send + Sync {}
impl<T> PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
pub fn from_pbo(
buffer_id: u32,
size: usize,
shape: &[usize],
name: Option<&str>,
ops: Arc<dyn PboOps>,
) -> Result<Self> {
if size == 0 {
return Err(Error::InvalidSize(0));
}
let expected = shape.iter().product::<usize>() * std::mem::size_of::<T>();
if size < expected {
return Err(Error::ShapeMismatch(format!(
"PBO size {size} is smaller than shape {shape:?} * sizeof({}) = {expected}",
std::any::type_name::<T>(),
)));
}
let name = name.unwrap_or("pbo_tensor").to_owned();
Ok(Self {
name,
shape: shape.to_vec(),
handle: Arc::new(PboHandle {
ops,
buffer_id,
size,
map_state: Mutex::new(MapState::Unmapped),
map_cv: Condvar::new(),
abi_vtable: OnceLock::new(),
}),
identity: BufferIdentity::derived(IdentityKind::Pbo, buffer_id as u64),
view_offset: 0,
_marker: PhantomData,
})
}
pub fn buffer_id(&self) -> u32 {
self.handle.buffer_id
}
pub(crate) fn pbo_vtable(&self) -> &PboOpsVtable {
self.handle
.abi_vtable
.get_or_init(|| PboOpsVtable::new(&self.handle.ops))
}
pub(crate) fn pbo_keepalive(&self) -> Arc<dyn Send + Sync> {
self.handle.clone()
}
pub fn is_mapped(&self) -> bool {
self.handle.is_mapped()
}
}
impl<T> TensorTrait<T> for PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn new(_shape: &[usize], _name: Option<&str>) -> Result<Self> {
Err(Error::NotImplemented(
"PboTensor cannot be created directly — use ImageProcessor::create_image()".to_owned(),
))
}
#[cfg(unix)]
fn from_fd(_fd: std::os::fd::OwnedFd, _shape: &[usize], _name: Option<&str>) -> Result<Self> {
Err(Error::NotImplemented(
"PboTensor does not support from_fd".to_owned(),
))
}
#[cfg(unix)]
fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
Err(Error::NotImplemented(
"PboTensor does not support clone_fd".to_owned(),
))
}
fn memory(&self) -> TensorMemory {
TensorMemory::Pbo
}
fn name(&self) -> String {
self.name.clone()
}
fn shape(&self) -> &[usize] {
&self.shape
}
fn reshape(&mut self, shape: &[usize]) -> Result<()> {
if shape.is_empty() {
return Err(Error::InvalidSize(0));
}
let new_size = shape.iter().product::<usize>() * std::mem::size_of::<T>();
if new_size != self.handle.size {
return Err(Error::ShapeMismatch(format!(
"Cannot reshape incompatible shape: {:?} to {:?}",
self.shape, shape
)));
}
self.shape = shape.to_vec();
self.view_offset = 0;
Ok(())
}
fn capacity_bytes(&self) -> usize {
self.handle.size
}
fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
if shape.is_empty() {
return Err(Error::InvalidSize(0));
}
let needed = shape.iter().product::<usize>() * std::mem::size_of::<T>();
if needed > self.handle.size {
return Err(Error::InsufficientCapacity {
needed,
capacity: self.handle.size,
});
}
self.shape = shape.to_vec();
Ok(())
}
fn map_with<'a>(&self, access: crate::CpuAccess) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_internal(None, access)
}
fn buffer_identity(&self) -> &BufferIdentity {
&self.identity
}
fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
if !offset_bytes.is_multiple_of(std::mem::align_of::<T>()) {
return Err(Error::InvalidOperation(format!(
"PboTensor::view: offset {offset_bytes} not aligned to align_of::<T>()={}",
std::mem::align_of::<T>()
)));
}
let abs_offset = self
.view_offset
.checked_add(offset_bytes)
.ok_or(Error::InvalidSize(offset_bytes))?;
let logical = shape.iter().product::<usize>() * std::mem::size_of::<T>();
let needed = abs_offset
.checked_add(logical)
.ok_or(Error::InvalidSize(logical))?;
if needed > self.handle.size {
return Err(Error::InsufficientCapacity {
needed,
capacity: self.handle.size,
});
}
Ok(Self {
name: self.name.clone(),
shape: shape.to_vec(),
handle: Arc::clone(&self.handle),
identity: self.identity.clone(),
view_offset: abs_offset,
_marker: PhantomData,
})
}
}
impl<T> PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
#[cfg(feature = "static")]
pub(crate) fn map_with_byte_size<'a>(
&self,
byte_size: usize,
access: crate::CpuAccess,
) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
self.map_internal(Some(byte_size), access)
}
fn map_internal<'a>(
&self,
byte_size_override: Option<usize>,
access: crate::CpuAccess,
) -> Result<crate::view::HostView<'a, T>>
where
T: 'a,
{
Ok(crate::view::HostView::new(
self.scoped_pin(access)?,
self.shape.clone(),
byte_size_override,
access,
))
}
}
#[derive(Debug)]
struct PboPtr(NonNull<c_void>);
impl Deref for PboPtr {
type Target = NonNull<c_void>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
unsafe impl Send for PboPtr {}
#[allow(dead_code)]
pub(crate) struct PboMapLock {
handle: Arc<PboHandle>,
}
impl Drop for PboMapLock {
fn drop(&mut self) {
self.handle.release_map();
}
}
impl<T> PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
#[allow(dead_code)]
pub(crate) fn scoped_pin<'a>(&self, access: crate::CpuAccess) -> Result<crate::pin::HostPin<'a>>
where
T: 'a,
{
let ptr = self.handle.acquire_map(access.writes())?;
let base = unsafe { (ptr.as_ptr() as *mut u8).add(self.view_offset) };
let len = self.handle.size.saturating_sub(self.view_offset);
Ok(crate::pin::HostPin::new(
Arc::new(PboMapLock {
handle: Arc::clone(&self.handle),
}),
base,
len,
))
}
}
impl<T> Clone for PboTensor<T>
where
T: Num + Clone + fmt::Debug + Send + Sync,
{
fn clone(&self) -> Self {
Self {
name: self.name.clone(),
shape: self.shape.clone(),
handle: Arc::clone(&self.handle),
identity: self.identity.clone(),
view_offset: self.view_offset,
_marker: PhantomData,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::TensorMapTrait;
struct MockPboOps {
storage: Mutex<Vec<u8>>,
maps: std::sync::atomic::AtomicUsize,
unmaps: std::sync::atomic::AtomicUsize,
deletes: std::sync::atomic::AtomicUsize,
}
impl MockPboOps {
fn new(size: usize) -> Arc<Self> {
Arc::new(Self {
storage: Mutex::new(vec![0u8; size]),
maps: std::sync::atomic::AtomicUsize::new(0),
unmaps: std::sync::atomic::AtomicUsize::new(0),
deletes: std::sync::atomic::AtomicUsize::new(0),
})
}
fn map_count(&self) -> usize {
self.maps.load(std::sync::atomic::Ordering::Acquire)
}
fn delete_count(&self) -> usize {
self.deletes.load(std::sync::atomic::Ordering::Acquire)
}
fn unmap_count(&self) -> usize {
self.unmaps.load(std::sync::atomic::Ordering::Acquire)
}
}
unsafe impl PboOps for MockPboOps {
fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
self.maps.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
let storage = self.storage.lock().expect("lock");
assert_eq!(storage.len(), size);
Ok(PboMapping {
ptr: storage.as_ptr() as *mut u8,
size,
})
}
fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
self.unmaps
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
Ok(())
}
fn delete_buffer(&self, _buffer_id: u32) {
self.deletes
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}
}
#[test]
fn test_pbo_tensor_create_and_metadata() {
let ops = MockPboOps::new(24);
let tensor = PboTensor::<u8>::from_pbo(42, 24, &[2, 3, 4], Some("test_pbo"), ops).unwrap();
assert_eq!(tensor.memory(), TensorMemory::Pbo);
assert_eq!(tensor.name(), "test_pbo");
assert_eq!(tensor.shape(), &[2, 3, 4]);
assert_eq!(tensor.buffer_id(), 42);
assert!(!tensor.is_mapped());
}
#[test]
fn test_pbo_tensor_map_write_read() {
let ops = MockPboOps::new(12);
let tensor = PboTensor::<u8>::from_pbo(1, 12, &[3, 4], Some("rw_test"), ops).unwrap();
{
let mut map = tensor.map().expect("map should succeed");
assert_eq!(map.shape(), &[3, 4]);
assert!(tensor.is_mapped());
map.as_mut_slice().fill(0xAB);
assert!(map.as_slice().iter().all(|&b| b == 0xAB));
}
assert!(!tensor.is_mapped());
}
#[test]
fn test_pbo_tensor_double_map_fails() {
let ops = MockPboOps::new(8);
let tensor = PboTensor::<u8>::from_pbo(2, 8, &[8], None, ops).unwrap();
let _map1 = tensor.map().expect("first map should succeed");
assert!(tensor.is_mapped());
let result = tensor.map();
assert!(result.is_err(), "second map while mapped should fail");
}
#[derive(Default)]
struct NullMappingOps {
unmaps: std::sync::atomic::AtomicUsize,
}
impl NullMappingOps {
fn unmap_count(&self) -> usize {
self.unmaps.load(std::sync::atomic::Ordering::Acquire)
}
}
unsafe impl PboOps for NullMappingOps {
fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
Ok(PboMapping {
ptr: std::ptr::null_mut(),
size,
})
}
fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
self.unmaps
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
Ok(())
}
fn delete_buffer(&self, _buffer_id: u32) {}
}
#[test]
fn null_mapping_releases_the_buffer_and_resets_state() {
let ops = Arc::new(NullMappingOps::default());
let dyn_ops: Arc<dyn PboOps> = ops.clone();
let tensor = PboTensor::<u8>::from_pbo(12, 8, &[8], None, dyn_ops).unwrap();
let Err(err) = tensor.map_read() else {
panic!("a null mapping must not succeed");
};
assert!(
matches!(err, Error::InvalidSize(_)),
"unexpected error: {err:?}"
);
assert_eq!(
ops.unmap_count(),
1,
"GL still holds the buffer mapped — the failed map must release it"
);
assert!(
!tensor.is_mapped(),
"state must return to Unmapped so later maps can proceed"
);
let Err(err) = tensor.map_read() else {
panic!("still null, so it must still fail");
};
assert!(matches!(err, Error::InvalidSize(_)));
assert_eq!(ops.unmap_count(), 2, "and released again");
}
#[test]
fn read_maps_share_one_gl_mapping() {
let ops = MockPboOps::new(8);
let dyn_ops: Arc<dyn PboOps> = ops.clone();
let tensor = PboTensor::<u8>::from_pbo(9, 8, &[8], None, dyn_ops).unwrap();
let r1 = tensor.map_read().expect("first read map");
let r2 = tensor
.map_read()
.expect("second read map must share, not fail");
let r3 = tensor
.map_read()
.expect("third read map must share, not fail");
assert_eq!(ops.map_count(), 1, "one GL mapping serves every reader");
assert_eq!(ops.unmap_count(), 0, "nothing unmapped while readers live");
assert!(tensor.is_mapped());
assert_eq!(r1.as_slice(), r2.as_slice());
assert_eq!(r2.as_slice(), r3.as_slice());
drop(r1);
drop(r2);
assert_eq!(
ops.unmap_count(),
0,
"the mapping outlives every reader but the last"
);
assert!(tensor.is_mapped());
drop(r3);
assert_eq!(ops.unmap_count(), 1, "last reader out unmaps, exactly once");
assert!(!tensor.is_mapped());
}
#[test]
fn writers_and_readers_still_exclude_each_other() {
let ops = MockPboOps::new(8);
let dyn_ops: Arc<dyn PboOps> = ops.clone();
let tensor = PboTensor::<u8>::from_pbo(10, 8, &[8], None, dyn_ops).unwrap();
let reader = tensor.map_read().expect("read map");
assert!(
tensor.map().is_err(),
"a writer must not join an existing reader set"
);
drop(reader);
let writer = tensor.map().expect("write map after readers drained");
assert!(
tensor.map_read().is_err(),
"a reader must not join an exclusive writer"
);
assert!(
tensor.map().is_err(),
"a second writer must not join an exclusive writer"
);
drop(writer);
assert!(tensor.map_read().is_ok());
}
#[test]
fn explicit_unmap_then_drop_releases_only_once() {
let ops = MockPboOps::new(8);
let dyn_ops: Arc<dyn PboOps> = ops.clone();
let tensor = PboTensor::<u8>::from_pbo(11, 8, &[8], None, dyn_ops).unwrap();
let keeper = tensor.map_read().expect("reader that must survive");
let mut early = tensor.map_read().expect("reader unmapped by hand");
early.unmap();
drop(early);
assert_eq!(
ops.unmap_count(),
0,
"one reader still holds the mapping — it must not be unmapped"
);
assert!(tensor.is_mapped(), "the surviving reader still holds it");
assert_eq!(keeper.as_slice().len(), 8);
drop(keeper);
assert_eq!(ops.unmap_count(), 1);
assert!(!tensor.is_mapped());
}
#[test]
fn test_pbo_tensor_reshape() {
let ops = MockPboOps::new(24);
let mut tensor = PboTensor::<u8>::from_pbo(3, 24, &[2, 3, 4], None, ops).unwrap();
tensor
.reshape(&[4, 6])
.expect("compatible reshape should succeed");
assert_eq!(tensor.shape(), &[4, 6]);
let result = tensor.reshape(&[100]);
assert!(result.is_err(), "incompatible reshape should fail");
}
#[test]
fn test_pbo_tensor_set_logical_shape_capacity_based() {
let ops = MockPboOps::new(24);
let mut tensor = PboTensor::<u8>::from_pbo(7, 24, &[24], None, ops).unwrap();
tensor
.set_logical_shape(&[4, 5])
.expect("shape within capacity should succeed");
assert_eq!(tensor.shape(), &[4, 5]);
tensor.set_logical_shape(&[24]).unwrap();
assert!(
tensor.set_logical_shape(&[5, 5]).is_err(),
"shape exceeding PBO capacity must be rejected"
);
}
#[test]
fn test_pbo_tensor_buffer_identity() {
let ops1 = MockPboOps::new(8);
let ops2 = MockPboOps::new(8);
let t1 = PboTensor::<u8>::from_pbo(1, 8, &[8], None, ops1).unwrap();
let t2 = PboTensor::<u8>::from_pbo(2, 8, &[8], None, ops2).unwrap();
assert_eq!(t1.buffer_identity().kind(), IdentityKind::Pbo);
assert_ne!(t1.buffer_identity().id(), t2.buffer_identity().id());
}
#[test]
fn two_pbos_with_the_same_gl_name_share_an_identity() {
let ops1 = MockPboOps::new(8);
let ops2 = MockPboOps::new(8);
let t1 = PboTensor::<u8>::from_pbo(7, 8, &[8], None, ops1).unwrap();
let t2 = PboTensor::<u8>::from_pbo(7, 8, &[8], None, ops2).unwrap();
assert_eq!(t1.buffer_identity().id(), t2.buffer_identity().id());
}
#[test]
fn test_pbo_tensor_new_returns_error() {
let result = PboTensor::<u8>::new(&[8], None);
assert!(result.is_err(), "PboTensor::new() should fail");
}
#[cfg(unix)]
#[test]
fn test_pbo_tensor_fd_ops_return_error() {
let ops = MockPboOps::new(8);
let tensor = PboTensor::<u8>::from_pbo(1, 8, &[8], None, ops).unwrap();
assert!(tensor.clone_fd().is_err());
}
#[test]
fn test_pbo_tensor_from_pbo_size_mismatch() {
let ops = MockPboOps::new(24);
let result = PboTensor::<u8>::from_pbo(1, 24, &[2, 3, 5], None, ops);
assert!(result.is_err(), "mismatched size/shape should fail");
}
#[test]
fn test_pbo_tensor_from_pbo_zero_size() {
let ops = MockPboOps::new(0);
let result = PboTensor::<u8>::from_pbo(1, 0, &[0], None, ops);
assert!(result.is_err(), "zero size should fail");
}
#[test]
fn test_pbo_via_tensor_enum() {
let ops = MockPboOps::new(12);
let pbo = PboTensor::<u8>::from_pbo(10, 12, &[3, 4], Some("enum_test"), ops).unwrap();
let tensor = crate::Tensor::wrap(crate::TensorStorage::Pbo(pbo));
assert_eq!(tensor.memory(), TensorMemory::Pbo);
assert_eq!(tensor.name(), "enum_test");
assert_eq!(tensor.shape(), &[3, 4]);
let mut map = tensor.map().expect("map via enum");
map.as_mut_slice().fill(42);
assert!(map.as_slice().iter().all(|&b| b == 42));
}
#[test]
fn cross_package_pbo_round_trips_through_tensor_desc_with_real_bytes() {
let ops = MockPboOps::new(8);
let dyn_ops: Arc<dyn PboOps> = ops.clone();
let producer = PboTensor::<u8>::from_pbo(42, 8, &[8], None, dyn_ops).unwrap();
{
let mut m = producer.map().expect("map producer");
m.as_mut_slice().copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
}
let producer_dyn =
crate::TensorDyn::U8(crate::Tensor::wrap(crate::TensorStorage::Pbo(producer)));
let desc = producer_dyn.descriptor();
assert_eq!(desc.kind, crate::protocol::kind::PBO);
assert_eq!(desc.handle, 42);
assert!(
!desc.ptr.is_null(),
"a PBO descriptor must carry an ops vtable -- see TensorDesc::ptr's doc comment"
);
let imported = crate::TensorDyn::import_descriptor(&desc)
.expect("import_descriptor must reconstruct a PBO-backed tensor from its own vtable");
assert_eq!(imported.memory(), TensorMemory::Pbo);
let m = imported
.map_bytes(crate::CpuAccess::Read)
.expect("map the reconstructed tensor");
assert_eq!(
m.as_slice(),
&[1u8, 2, 3, 4, 5, 6, 7, 8],
"the reconstructed tensor must read the SAME bytes the producer wrote through the \
SAME GL buffer, not a copy or garbage"
);
}
#[test]
fn two_pbo_tensors_sharing_a_buffer_id_do_not_cross_resolve() {
let ops_a = MockPboOps::new(4);
let ops_b = MockPboOps::new(4);
ops_a.storage.lock().unwrap().copy_from_slice(&[0xAA; 4]);
ops_b.storage.lock().unwrap().copy_from_slice(&[0xBB; 4]);
let dyn_a: Arc<dyn PboOps> = ops_a.clone();
let dyn_b: Arc<dyn PboOps> = ops_b.clone();
let tensor_a = PboTensor::<u8>::from_pbo(7, 4, &[4], None, dyn_a).unwrap();
let tensor_b = PboTensor::<u8>::from_pbo(7, 4, &[4], None, dyn_b).unwrap();
let dyn_a_wrapped =
crate::TensorDyn::U8(crate::Tensor::wrap(crate::TensorStorage::Pbo(tensor_a)));
let dyn_b_wrapped =
crate::TensorDyn::U8(crate::Tensor::wrap(crate::TensorStorage::Pbo(tensor_b)));
let desc_a = dyn_a_wrapped.descriptor();
let desc_b = dyn_b_wrapped.descriptor();
assert_eq!(
desc_a.handle, desc_b.handle,
"both descriptors share the same buffer_id by construction"
);
assert_ne!(
desc_a.ptr.0, desc_b.ptr.0,
"but each carries its own, distinct vtable address -- there is no shared key \
for them to collide on"
);
let imported_a = crate::TensorDyn::import_descriptor(&desc_a).unwrap();
let imported_b = crate::TensorDyn::import_descriptor(&desc_b).unwrap();
assert_eq!(
imported_a
.map_bytes(crate::CpuAccess::Read)
.unwrap()
.as_slice(),
&[0xAA; 4],
"importing desc_a must reach ops_a's buffer, never ops_b's"
);
assert_eq!(
imported_b
.map_bytes(crate::CpuAccess::Read)
.unwrap()
.as_slice(),
&[0xBB; 4],
"importing desc_b must reach ops_b's buffer, never ops_a's"
);
}
#[test]
fn dropping_the_producing_pbo_tensor_leaks_nothing_the_vtable_touched() {
let ops = MockPboOps::new(8);
let dyn_ops: Arc<dyn PboOps> = ops.clone();
assert_eq!(
Arc::strong_count(&ops),
2,
"this test's own handle, plus the clone about to move into from_pbo"
);
let producer = PboTensor::<u8>::from_pbo(99, 8, &[8], None, dyn_ops).unwrap();
assert_eq!(
Arc::strong_count(&ops),
2,
"from_pbo must not clone ops again"
);
let producer_dyn =
crate::TensorDyn::U8(crate::Tensor::wrap(crate::TensorStorage::Pbo(producer)));
let _desc = producer_dyn.descriptor();
assert_eq!(
Arc::strong_count(&ops),
2,
"building/exporting the vtable must not clone ops -- see PboOpsVtable's own \
doc comment for why ctx is a raw pointer, not an Arc clone"
);
assert_eq!(ops.delete_count(), 0, "not dropped yet");
drop(producer_dyn);
assert_eq!(
Arc::strong_count(&ops),
1,
"dropping the producing tensor must release its ops reference immediately -- \
nothing the vtable machinery built kept it alive past that"
);
assert_eq!(
ops.delete_count(),
1,
"the real delete_buffer must still fire exactly once, from PboHandle::Drop, \
completely unaffected by whether a descriptor was ever exported"
);
}
}