use crate::{
error::{Error, Result},
BufferIdentity, TensorMap, TensorMapTrait, TensorMemory, TensorTrait,
};
use log::trace;
use num_traits::Num;
use std::{
ffi::c_void,
fmt,
marker::PhantomData,
ops::{Deref, DerefMut},
ptr::NonNull,
sync::{Arc, Condvar, Mutex},
};
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);
}
struct PboHandle {
ops: Arc<dyn PboOps>,
buffer_id: u32,
size: usize,
map_state: Mutex<MapState>,
map_cv: Condvar,
}
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(),
}),
identity: BufferIdentity::new(),
view_offset: 0,
_marker: PhantomData,
})
}
pub fn buffer_id(&self) -> u32 {
self.handle.buffer_id
}
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(&self, access: crate::CpuAccess) -> Result<TensorMap<T>> {
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,
{
pub(crate) fn map_with_byte_size(
&self,
byte_size: usize,
access: crate::CpuAccess,
) -> Result<TensorMap<T>> {
self.map_internal(Some(byte_size), access)
}
fn map_internal(
&self,
byte_size_override: Option<usize>,
access: crate::CpuAccess,
) -> Result<TensorMap<T>> {
let writable = access.writes();
let ptr = self.handle.acquire_map(writable)?;
Ok(TensorMap::Pbo(PboMap {
ptr: Arc::new(Mutex::new(ptr)),
shape: self.shape.clone(),
handle: Arc::clone(&self.handle),
byte_size_override,
view_offset: self.view_offset,
writable,
released: false,
_marker: PhantomData,
}))
}
}
#[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 {}
pub struct PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
ptr: Arc<Mutex<PboPtr>>,
shape: Vec<usize>,
handle: Arc<PboHandle>,
byte_size_override: Option<usize>,
view_offset: usize,
writable: bool,
released: bool,
_marker: PhantomData<T>,
}
impl<T> fmt::Debug for PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PboMap")
.field("shape", &self.shape)
.field("buffer_id", &self.handle.buffer_id)
.finish()
}
}
impl<T> TensorMapTrait<T> for PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
fn shape(&self) -> &[usize] {
&self.shape
}
fn unmap(&mut self) {
if self.released {
return;
}
self.released = true;
self.handle.release_map();
}
fn as_slice(&self) -> &[T] {
let ptr = self.ptr.lock().expect("Failed to lock PboMap pointer");
let base = unsafe { (ptr.as_ptr() as *const u8).add(self.view_offset) as *const T };
unsafe { std::slice::from_raw_parts(base, self.slice_len_elems()) }
}
fn as_mut_slice(&mut self) -> &mut [T] {
crate::assert_map_writable(self.writable, "Pbo");
let ptr = self.ptr.lock().expect("Failed to lock PboMap pointer");
let base = unsafe { (ptr.as_ptr() as *mut u8).add(self.view_offset) as *mut T };
unsafe { std::slice::from_raw_parts_mut(base, self.slice_len_elems()) }
}
}
impl<T> PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
fn slice_len_elems(&self) -> usize {
match self.byte_size_override {
Some(bytes) => bytes / std::mem::size_of::<T>(),
None => self.shape.iter().product(),
}
}
}
impl<T> Deref for PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
type Target = [T];
fn deref(&self) -> &[T] {
self.as_slice()
}
}
impl<T> DerefMut for PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
fn deref_mut(&mut self) -> &mut [T] {
self.as_mut_slice()
}
}
impl<T> Drop for PboMap<T>
where
T: Num + Clone + fmt::Debug,
{
fn drop(&mut self) {
self.unmap();
}
}
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::*;
struct MockPboOps {
storage: Mutex<Vec<u8>>,
maps: std::sync::atomic::AtomicUsize,
unmaps: 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),
})
}
fn map_count(&self) -> usize {
self.maps.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) {}
}
#[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_ne!(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));
}
}