1use std::fmt;
2use std::marker::PhantomData;
3use std::mem::{align_of, size_of};
4use std::ptr::NonNull;
5
6use crate::types::{shape_vec, stride_vec, tensor_from_group, tensor_view_from_group};
7use crate::{DType, DynRank, Placement, TensorLayout, TensorRank, TensorRead, TensorScalar};
8use tinyvec::TinyVec;
9
10use super::identity::RootResourceId;
11use super::prepared::{
12 prepare_read, prepare_write, validate_descriptor, AccessError, AccessTarget, CheckedDescriptor,
13 CheckedRead, CheckedWrite, PreparedRead, PreparedWrite, ProviderReadMapping,
14 ProviderWriteMapping, WriteInjectivityProof,
15};
16use super::root::{BackendAllocation, OwnedStorage, ProviderKind};
17use super::span::{ByteRange, RootBoundSpan};
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
21pub(crate) struct AllocationSlot(u32);
22
23impl AllocationSlot {
24 pub(crate) const fn index(self) -> usize {
25 self.0 as usize
26 }
27
28 #[cfg(test)]
29 pub(crate) const fn test_raw(raw: u32) -> Self {
30 Self(raw)
31 }
32}
33
34#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
45pub struct DescriptorSlot(u32);
46
47impl DescriptorSlot {
48 pub const fn index(self) -> usize {
59 self.0 as usize
60 }
61
62 pub fn from_index(index: usize) -> Option<Self> {
73 match u32::try_from(index) {
74 Ok(index) => Some(Self(index)),
75 Err(_) => None,
76 }
77 }
78
79 #[cfg(test)]
80 pub(crate) const fn test_raw(raw: u32) -> Self {
81 Self(raw)
82 }
83}
84
85#[derive(Clone, Debug, PartialEq, Eq)]
87pub(crate) struct DescriptorInput<R: TensorRank> {
88 relative: ByteRange,
89 shape: R::Shape,
90 strides: R::Strides,
91 offset: isize,
92 require_injective: bool,
93}
94
95impl<R: TensorRank> DescriptorInput<R> {
96 pub(crate) fn new(
97 relative: ByteRange,
98 shape: R::Shape,
99 strides: R::Strides,
100 offset: isize,
101 require_injective: bool,
102 ) -> Self {
103 Self {
104 relative,
105 shape,
106 strides,
107 offset,
108 require_injective,
109 }
110 }
111}
112
113#[derive(Clone, Debug, PartialEq, Eq)]
119pub(crate) struct DescriptorRecord {
120 allocation: AllocationSlot,
121 element_count: usize,
122 provider: ProviderKind,
123 placement: Placement,
124 envelope: Option<ByteRange>,
125 write_injective: bool,
126 checked: CheckedDescriptor<DynRank>,
127}
128
129impl DescriptorRecord {
130 pub(crate) const fn allocation(&self) -> AllocationSlot {
131 self.allocation
132 }
133
134 pub(crate) const fn span(&self) -> RootBoundSpan {
135 self.checked.span()
136 }
137
138 pub(crate) const fn dtype(&self) -> DType {
139 self.checked.dtype()
140 }
141
142 pub(crate) const fn element_size(&self) -> usize {
143 self.checked.element_size()
144 }
145
146 pub(crate) const fn element_count(&self) -> usize {
147 self.element_count
148 }
149
150 pub(crate) fn layout(&self) -> &TensorLayout<DynRank> {
151 self.checked.logical_layout()
152 }
153
154 pub(crate) const fn provider(&self) -> ProviderKind {
155 self.provider
156 }
157
158 pub(crate) fn placement(&self) -> &Placement {
159 &self.placement
160 }
161
162 pub(crate) const fn envelope(&self) -> Option<ByteRange> {
163 self.envelope
164 }
165
166 pub(crate) const fn write_injective(&self) -> bool {
167 self.write_injective
168 }
169}
170
171fn into_group_parts(
179 tensor: crate::Tensor,
180) -> Result<(AllocationGroup, DescriptorSlot), GroupError> {
181 if tensor.external_payload().is_some() {
182 return Err(GroupError::InvalidDescriptor {
183 message: "a caller-owned payload has no allocation group".to_owned(),
184 });
185 }
186 Ok(tensor.into_group_parts())
187}
188
189#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
200pub enum GroupError {
201 #[error("group index overflows u32")]
202 IndexOverflow,
203 #[error("descriptor slot {slot} is outside the group")]
204 DescriptorSlotOutOfBounds { slot: usize },
205 #[error("descriptor slot {slot} is vacant")]
206 DescriptorSlotVacant { slot: usize },
207 #[error("allocation slot {slot} is outside the group")]
208 AllocationSlotOutOfBounds { slot: usize },
209 #[error("allocation slot {slot} is vacant")]
210 AllocationSlotVacant { slot: usize },
211 #[error("descriptor validation failed: {message}")]
212 InvalidDescriptor { message: String },
213 #[error("descriptor dtype mismatch: expected {expected:?}, actual {actual:?}")]
214 DTypeMismatch { expected: DType, actual: DType },
215 #[error("descriptor rank mismatch: expected {expected}, actual {actual}")]
216 RankMismatch { expected: usize, actual: usize },
217 #[error("allocation slot {allocation} has more than one live descriptor")]
218 AliasedAllocation { allocation: usize },
219 #[error(
223 "descriptor slot {slot} is not a compact whole-allocation layout; materialize it instead of extracting"
224 )]
225 NonCompactDescriptor { slot: usize },
226}
227
228#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
230pub(crate) enum DisjointViewError {
231 #[error(transparent)]
232 Group(#[from] GroupError),
233 #[error("descriptor slot {slot} appears more than once")]
234 DuplicateSlot { slot: usize },
235 #[error("descriptor slot {slot} has a non-injective mutable layout")]
236 NonInjective { slot: usize },
237 #[error("requested mutable descriptor envelopes overlap")]
238 PairwiseOverlap,
239 #[error("requested mutable descriptors are not provably disjoint")]
240 NotProvablyDisjoint,
241}
242
243#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
245pub(crate) enum ExtractError {
246 #[error(transparent)]
247 Group(#[from] GroupError),
248 #[error("allocation slot {allocation} still has another descriptor")]
249 AliasedAllocation { allocation: usize },
250}
251
252#[derive(Default)]
265pub struct AllocationGroup {
266 allocations: TinyVec<[Option<OwnedStorage>; 1]>,
271 descriptors: TinyVec<[Option<DescriptorRecord>; 1]>,
272}
273
274pub(crate) struct GroupReadView<'a, T, R: TensorRank> {
276 owner: NonNull<OwnedStorage>,
277 descriptor: DescriptorRecord,
278 _borrow: PhantomData<(&'a OwnedStorage, T, R)>,
279}
280
281unsafe impl<'a, T: Send, R: TensorRank> Send for GroupReadView<'a, T, R> {}
284unsafe impl<'a, T: Sync, R: TensorRank> Sync for GroupReadView<'a, T, R> {}
285
286impl<'a, T, R: TensorRank> Clone for GroupReadView<'a, T, R> {
287 fn clone(&self) -> Self {
288 Self {
289 owner: self.owner,
290 descriptor: self.descriptor.clone(),
291 _borrow: PhantomData,
292 }
293 }
294}
295
296impl<'a, T, R: TensorRank> GroupReadView<'a, T, R> {
297 pub(crate) fn clone_dyn(&self) -> GroupReadView<'a, T, crate::DynRank> {
298 GroupReadView {
299 owner: self.owner,
300 descriptor: self.descriptor.clone(),
301 _borrow: PhantomData,
302 }
303 }
304
305 pub(crate) fn into_dyn(self) -> GroupReadView<'a, T, crate::DynRank> {
306 GroupReadView {
307 owner: self.owner,
308 descriptor: self.descriptor,
309 _borrow: PhantomData,
310 }
311 }
312}
313
314impl<T, R: TensorRank> std::fmt::Debug for GroupReadView<'_, T, R> {
315 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
316 formatter
317 .debug_struct("GroupReadView")
318 .field("descriptor", &self.descriptor)
319 .finish_non_exhaustive()
320 }
321}
322
323impl<'a, T: TensorScalar, R: TensorRank> GroupReadView<'a, T, R> {
324 pub(crate) fn descriptor(&self) -> &DescriptorRecord {
325 &self.descriptor
326 }
327
328 pub(crate) fn provider_kind(&self) -> ProviderKind {
329 self.descriptor.provider
330 }
331
332 pub(crate) fn backend_identity(
333 &self,
334 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
335 if self.descriptor.provider == ProviderKind::Cpu {
336 return None;
337 }
338 let key = unsafe { self.owner.as_ref().root_identity().extent().key() };
339 Some((key.domain(), key.local()))
340 }
341
342 pub(crate) fn storage_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
343 let buffer = unsafe {
346 self.owner
347 .as_ref()
348 .host_buffer::<T>()
349 .or_else(|| self.owner.as_ref().backend_buffer::<T>())
350 }?;
351 Some(unsafe {
352 std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
353 })
354 }
355
356 pub(crate) fn host_root_slice(&self) -> Option<&'a [T]> {
365 let owner = unsafe { self.owner.as_ref() };
367 owner
368 .as_ref()
369 .host_slice::<T>(owner.root_span(), self.descriptor.dtype())
370 .ok()
371 }
372
373 pub(crate) fn map_read(&self) -> Result<ProviderReadMapping<'_>, AccessError> {
374 unsafe {
377 self.owner
378 .as_ref()
379 .as_ref()
380 .map_read(self.descriptor.span(), self.descriptor.dtype())
381 }
382 }
383
384 pub(crate) fn prepare_device_read_for_layout(
385 &self,
386 layout: &TensorLayout<R>,
387 ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
388 self.prepare_read_for_layout(layout, AccessTarget::Device)?
389 .into_device_state()
390 }
391
392 pub(crate) fn prepare_host_read_for_layout(
393 &self,
394 layout: &TensorLayout<R>,
395 ) -> Result<PreparedRead<'a, T, R>, AccessError> {
396 self.prepare_read_for_layout(layout, AccessTarget::Host)
397 }
398
399 fn prepare_read_for_layout(
400 &self,
401 layout: &TensorLayout<R>,
402 target: AccessTarget,
403 ) -> Result<PreparedRead<'a, T, R>, AccessError> {
404 let owner: crate::storage::root::StorageRef<'a> = unsafe { self.owner.as_ref().as_ref() };
406 let checked: CheckedRead<'a, R> = CheckedRead::new::<T>(
407 owner,
409 self.descriptor.span(),
410 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
411 AccessError::InvalidLayout {
412 message: error.to_string(),
413 }
414 })?,
415 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
416 AccessError::InvalidLayout {
417 message: error.to_string(),
418 }
419 })?,
420 layout.offset(),
421 )?;
422 prepare_read::<T, R>(checked, target).map_err(|failure| failure.1)
423 }
424
425 pub(crate) fn prepare_host_read(&self) -> Result<PreparedRead<'_, T, DynRank>, AccessError> {
426 prepare_read(
427 CheckedRead::from_validated(
428 unsafe { self.owner.as_ref().as_ref() },
430 self.descriptor.checked.clone(),
431 ),
432 AccessTarget::Host,
433 )
434 .map_err(|failure| failure.1)
435 }
436
437 pub(crate) fn host_slice(&self) -> Result<&'a [T], AccessError> {
438 unsafe {
440 self.owner
441 .as_ref()
442 .as_ref()
443 .host_slice(self.descriptor.span(), self.descriptor.dtype())
444 }
445 }
446}
447
448impl<'a, T: 'static, R: TensorRank> GroupReadView<'a, T, R> {
449 pub(crate) fn backend_allocation(&self) -> Option<&'a dyn BackendAllocation> {
450 let allocation = unsafe { self.owner.as_ref().backend_allocation() }?;
452 Some(unsafe {
453 std::mem::transmute::<&dyn BackendAllocation, &'a dyn BackendAllocation>(allocation)
454 })
455 }
456
457 pub(crate) fn backend_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
458 let buffer = unsafe { self.owner.as_ref().backend_buffer::<T>() }?;
461 Some(unsafe {
462 std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
463 })
464 }
465}
466
467pub(crate) struct GroupWriteView<'a, T, R: TensorRank> {
471 owner: NonNull<OwnedStorage>,
472 descriptor: DescriptorRecord,
473 _borrow: PhantomData<(&'a mut [u8], T, R)>,
474}
475
476unsafe impl<'a, T: Send, R: TensorRank> Send for GroupWriteView<'a, T, R> {}
479
480impl<T, R: TensorRank> std::fmt::Debug for GroupWriteView<'_, T, R> {
481 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
482 formatter
483 .debug_struct("GroupWriteView")
484 .field("descriptor", &self.descriptor)
485 .finish_non_exhaustive()
486 }
487}
488
489impl<'a, T: TensorScalar, R: TensorRank> GroupWriteView<'a, T, R> {
490 pub(crate) fn descriptor(&self) -> &DescriptorRecord {
491 &self.descriptor
492 }
493
494 pub(crate) fn map_write(&mut self) -> Result<ProviderWriteMapping<'_>, AccessError> {
495 unsafe {
499 self.owner
500 .as_mut()
501 .as_mut()
502 .map_write(self.descriptor.span(), self.descriptor.dtype())
503 }
504 }
505
506 pub(crate) fn backend_buffer_mut(&mut self) -> Option<&'a mut crate::StorageBuffer<T>> {
507 let owner = unsafe { &mut *self.owner.as_ptr() };
510 let buffer = owner.backend_buffer_mut::<T>()?;
511 Some(unsafe {
512 std::mem::transmute::<&mut crate::StorageBuffer<T>, &'a mut crate::StorageBuffer<T>>(
513 buffer,
514 )
515 })
516 }
517
518 pub(crate) fn prepare_device_write_for_layout(
519 &mut self,
520 layout: &TensorLayout<R>,
521 ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
522 let owner: crate::storage::root::StorageMut<'a> =
523 unsafe { (&mut *self.owner.as_ptr()).as_mut() };
524 let checked: CheckedWrite<'a, R> = CheckedWrite::new::<T>(
525 owner,
527 self.descriptor.span(),
528 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
529 AccessError::InvalidLayout {
530 message: error.to_string(),
531 }
532 })?,
533 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
534 AccessError::InvalidLayout {
535 message: error.to_string(),
536 }
537 })?,
538 layout.offset(),
539 )?;
540 prepare_write::<T, R>(checked, AccessTarget::Device)
541 .map_err(|failure| failure.1)?
542 .into_device_state()
543 }
544
545 pub(crate) fn prepare_host_write(
546 &mut self,
547 ) -> Result<PreparedWrite<'_, T, DynRank>, AccessError> {
548 let checked = CheckedWrite::from_validated(
549 unsafe { self.owner.as_mut().as_mut() },
551 self.descriptor.checked.clone(),
552 WriteInjectivityProof,
553 );
554 prepare_write(checked, AccessTarget::Host).map_err(|failure| failure.1)
555 }
556
557 pub(crate) fn host_slice_mut(&mut self) -> Result<&'a mut [T], AccessError> {
558 unsafe {
561 self.owner
562 .as_mut()
563 .as_mut()
564 .host_slice_mut(self.descriptor.span(), self.descriptor.dtype())
565 }
566 }
567}
568
569impl<'a, T: 'static, R: TensorRank> GroupWriteView<'a, T, R> {
570 pub(crate) fn backend_buffer(&self) -> Option<&crate::StorageBuffer<T>> {
571 unsafe { self.owner.as_ref().backend_buffer::<T>() }
574 }
575}
576
577impl fmt::Debug for AllocationGroup {
578 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
579 formatter
580 .debug_struct("AllocationGroup")
581 .field("allocation_count", &self.allocations.len())
582 .field("descriptor_count", &self.descriptors.len())
583 .finish()
584 }
585}
586
587impl AllocationGroup {
588 pub(crate) fn new() -> Self {
589 Self::default()
590 }
591
592 pub fn from_tensors(
611 tensors: Vec<crate::Tensor>,
612 ) -> Result<(Self, Box<[DescriptorSlot]>), GroupError> {
613 let mut group = Self::new();
614 let mut bindings = Vec::with_capacity(tensors.len());
615 for tensor in tensors {
616 let (source, source_slot) = into_group_parts(tensor)?;
617 bindings.push(group.append_group(source, source_slot)?);
618 }
619 Ok((group, bindings.into_boxed_slice()))
620 }
621
622 pub fn read_views<'a>(
631 &'a self,
632 bindings: &[DescriptorSlot],
633 ) -> Result<Vec<TensorRead<'a>>, GroupError> {
634 bindings
635 .iter()
636 .map(|&slot| self.tensor_read(slot))
637 .collect()
638 }
639
640 pub fn read_view<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
649 self.tensor_read(slot)
650 }
651
652 fn tensor_read<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
653 let dtype = self.resolve_descriptor(slot)?.1.dtype();
654 let view = match dtype {
655 DType::F32 => tensor_view_from_group(self.view::<f32, DynRank>(slot)?),
656 DType::F64 => tensor_view_from_group(self.view::<f64, DynRank>(slot)?),
657 DType::I32 => tensor_view_from_group(self.view::<i32, DynRank>(slot)?),
658 DType::I64 => tensor_view_from_group(self.view::<i64, DynRank>(slot)?),
659 DType::Bool => tensor_view_from_group(self.view::<bool, DynRank>(slot)?),
660 DType::C32 => {
661 tensor_view_from_group(self.view::<num_complex::Complex32, DynRank>(slot)?)
662 }
663 DType::C64 => {
664 tensor_view_from_group(self.view::<num_complex::Complex64, DynRank>(slot)?)
665 }
666 DType::External(_) => unreachable!("descriptors are preset-typed"),
669 }
670 .map_err(|error| GroupError::InvalidDescriptor {
671 message: error.to_string(),
672 })?;
673 Ok(TensorRead::from_view(view))
674 }
675
676 pub fn append_tensor(&mut self, tensor: crate::Tensor) -> Result<DescriptorSlot, GroupError> {
685 let (source, source_slot) = into_group_parts(tensor)?;
686 self.append_group(source, source_slot)
687 }
688
689 pub fn append_group(
696 &mut self,
697 mut source: AllocationGroup,
698 source_slot: DescriptorSlot,
699 ) -> Result<DescriptorSlot, GroupError> {
700 let allocation_offset =
701 u32::try_from(self.allocations.len()).map_err(|_| GroupError::IndexOverflow)?;
702 let descriptor_offset =
703 u32::try_from(self.descriptors.len()).map_err(|_| GroupError::IndexOverflow)?;
704 source.resolve_descriptor(source_slot)?;
705
706 for owner in source.allocations.drain(..) {
707 self.allocations.push(owner);
708 }
709 for descriptor in source.descriptors.drain(..) {
710 let descriptor = match descriptor {
711 Some(mut descriptor) => {
712 let allocation = descriptor
713 .allocation
714 .0
715 .checked_add(allocation_offset)
716 .ok_or(GroupError::IndexOverflow)?;
717 descriptor.allocation = AllocationSlot(allocation);
718 Some(descriptor)
719 }
720 None => None,
721 };
722 self.descriptors.push(descriptor);
723 }
724
725 let source_descriptor_index =
726 u32::try_from(source_slot.index()).map_err(|_| GroupError::IndexOverflow)?;
727 Ok(DescriptorSlot(
728 source_descriptor_index
729 .checked_add(descriptor_offset)
730 .ok_or(GroupError::IndexOverflow)?,
731 ))
732 }
733
734 pub(crate) fn set_descriptor_placement(
735 &mut self,
736 slot: DescriptorSlot,
737 placement: Placement,
738 ) -> Result<(), GroupError> {
739 let descriptor = self
740 .descriptors
741 .get_mut(slot.index())
742 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: slot.index() })?
743 .as_mut()
744 .ok_or(GroupError::DescriptorSlotVacant { slot: slot.index() })?;
745 descriptor.placement = placement;
746 Ok(())
747 }
748
749 pub(crate) fn publish_live_descriptor_placement(
758 &mut self,
759 slot: DescriptorSlot,
760 placement: Placement,
761 ) {
762 let descriptor = self
763 .descriptors
764 .get_mut(slot.index())
765 .and_then(Option::as_mut)
766 .unwrap_or_else(|| unreachable!("an owned group always carries its descriptor"));
767 descriptor.placement = placement;
768 }
769
770 pub(crate) fn from_host_vec<T: TensorScalar, R: TensorRank>(
771 shape: R::Shape,
772 data: Vec<T>,
773 ) -> Result<(Self, DescriptorSlot), GroupError> {
774 let owner =
775 super::root::import_host_vec(data).map_err(|error| GroupError::InvalidDescriptor {
776 message: error.to_string(),
777 })?;
778 let span = owner.root_span();
779 let mut group = Self::new();
780 let allocation = group.insert_owner(owner)?;
781 let layout =
782 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
783 message: error.to_string(),
784 })?;
785 let input = DescriptorInput::new(
786 ByteRange::new(0, span.byte_len()),
787 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
788 GroupError::InvalidDescriptor {
789 message: error.to_string(),
790 }
791 })?,
792 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
793 GroupError::InvalidDescriptor {
794 message: error.to_string(),
795 }
796 })?,
797 layout.offset(),
798 true,
799 );
800 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
801 Ok((group, slot))
802 }
803
804 #[doc(hidden)]
806 pub fn from_backend_allocation<T: TensorScalar, R: TensorRank>(
807 shape: R::Shape,
808 allocation: Box<dyn BackendAllocation>,
809 ) -> Result<(Self, DescriptorSlot), GroupError> {
810 let owner = super::root::import_unique_root(allocation).map_err(|error| {
811 GroupError::InvalidDescriptor {
812 message: error.to_string(),
813 }
814 })?;
815 let span = owner.root_span();
816 let mut group = Self::new();
817 let allocation = group.insert_owner(owner)?;
818 let layout =
819 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
820 message: error.to_string(),
821 })?;
822 let input = DescriptorInput::new(
823 ByteRange::new(0, span.byte_len()),
824 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
825 GroupError::InvalidDescriptor {
826 message: error.to_string(),
827 }
828 })?,
829 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
830 GroupError::InvalidDescriptor {
831 message: error.to_string(),
832 }
833 })?,
834 layout.offset(),
835 true,
836 );
837 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
838 Ok((group, slot))
839 }
840
841 pub(crate) fn from_backend_buffer<T: TensorScalar, R: TensorRank>(
842 shape: R::Shape,
843 buffer: crate::StorageBuffer<T>,
844 ) -> Result<(Self, DescriptorSlot), GroupError> {
845 let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
846 GroupError::InvalidDescriptor {
847 message: error.to_string(),
848 }
849 })?;
850 let span = owner.root_span();
851 let mut group = Self::new();
852 let allocation = group.insert_owner(owner)?;
853 let layout =
854 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
855 message: error.to_string(),
856 })?;
857 let input = DescriptorInput::new(
858 ByteRange::new(0, span.byte_len()),
859 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
860 GroupError::InvalidDescriptor {
861 message: error.to_string(),
862 }
863 })?,
864 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
865 GroupError::InvalidDescriptor {
866 message: error.to_string(),
867 }
868 })?,
869 layout.offset(),
870 true,
871 );
872 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
873 Ok((group, slot))
874 }
875
876 pub(crate) fn from_backend_root<T: Send + Sync + 'static>(
877 buffer: crate::StorageBuffer<T>,
878 ) -> Result<Self, GroupError> {
879 let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
880 GroupError::InvalidDescriptor {
881 message: error.to_string(),
882 }
883 })?;
884 let mut group = Self::new();
885 group.insert_owner(owner)?;
886 Ok(group)
887 }
888
889 pub(crate) fn insert_owner(
890 &mut self,
891 owner: OwnedStorage,
892 ) -> Result<AllocationSlot, GroupError> {
893 let index = self.allocations.len();
894 let slot = u32::try_from(index).map_err(|_| GroupError::IndexOverflow)?;
895 self.allocations.push(Some(owner));
896 Ok(AllocationSlot(slot))
897 }
898
899 pub(crate) fn insert_descriptor<T: TensorScalar, R: TensorRank>(
900 &mut self,
901 allocation: AllocationSlot,
902 input: DescriptorInput<R>,
903 ) -> Result<DescriptorSlot, GroupError> {
904 let allocation_index = allocation.index();
905 let owner = self
906 .allocations
907 .get(allocation_index)
908 .ok_or(GroupError::AllocationSlotOutOfBounds {
909 slot: allocation_index,
910 })?
911 .as_ref()
912 .ok_or(GroupError::AllocationSlotVacant {
913 slot: allocation_index,
914 })?;
915
916 let root = owner.as_ref().root_identity();
917 let span = root.bind_relative_range(input.relative).map_err(|error| {
918 GroupError::InvalidDescriptor {
919 message: error.to_string(),
920 }
921 })?;
922 let element_size = size_of::<T>();
923 if element_size == 0 || !span.byte_len().is_multiple_of(element_size) {
924 return Err(GroupError::InvalidDescriptor {
925 message: format!(
926 "byte span {} is not divisible by element size {}",
927 span.byte_len(),
928 element_size
929 ),
930 });
931 }
932 if !span
933 .guaranteed_alignment()
934 .get()
935 .is_multiple_of(align_of::<T>())
936 {
937 return Err(GroupError::InvalidDescriptor {
938 message: format!(
939 "span alignment {} is insufficient for {}-byte alignment",
940 span.guaranteed_alignment().get(),
941 align_of::<T>()
942 ),
943 });
944 }
945
946 let shape = R::shape_into_vec(input.shape);
947 let strides = R::strides_into_vec(input.strides);
948 let layout = TensorLayout::<DynRank>::from_parts(
949 shape,
950 strides,
951 input.offset,
952 span.byte_len() / element_size,
953 )
954 .map_err(|error| GroupError::InvalidDescriptor {
955 message: error.to_string(),
956 })?;
957 let element_count = logical_element_count(layout.shape())?;
958 let write_injective = if input.require_injective {
959 layout.validate_mutable_no_overlap().map_err(|error| {
960 GroupError::InvalidDescriptor {
961 message: error.to_string(),
962 }
963 })?;
964 true
965 } else {
966 false
967 };
968 let envelope = reachable_envelope(&span, &layout, element_size)?;
969 let (checked, _) = validate_descriptor::<T, DynRank>(
970 root.root_resource(),
971 span,
972 layout.shape().iter().copied().collect(),
973 layout.strides().iter().copied().collect(),
974 layout.offset(),
975 false,
976 )
977 .map_err(|error| GroupError::InvalidDescriptor {
978 message: error.to_string(),
979 })?;
980 let record = DescriptorRecord {
981 allocation,
982 element_count,
983 provider: owner.as_ref().provider_kind(),
984 placement: Placement::default(),
985 envelope,
986 write_injective,
987 checked,
988 };
989
990 let descriptor_index = self.descriptors.len();
991 let slot = u32::try_from(descriptor_index).map_err(|_| GroupError::IndexOverflow)?;
992 self.descriptors.push(Some(record));
993 Ok(DescriptorSlot(slot))
994 }
995
996 #[allow(clippy::result_large_err)]
1001 pub(crate) fn update_descriptor_layout(
1002 mut self,
1003 slot: DescriptorSlot,
1004 shape: Vec<usize>,
1005 strides: Vec<isize>,
1006 offset: isize,
1007 ) -> Result<Self, (Self, GroupError)> {
1008 let dtype = match self.resolve_descriptor(slot) {
1009 Ok((_, descriptor)) => descriptor.dtype(),
1010 Err(error) => return Err((self, error)),
1011 };
1012 if let Some(Some(descriptor)) = self.descriptors.get_mut(slot.index()) {
1013 descriptor.write_injective = false;
1016 }
1017 match dtype {
1018 DType::F32 => self.reinterpret_descriptor::<f32, f32>(slot, shape, strides, offset),
1019 DType::F64 => self.reinterpret_descriptor::<f64, f64>(slot, shape, strides, offset),
1020 DType::I32 => self.reinterpret_descriptor::<i32, i32>(slot, shape, strides, offset),
1021 DType::I64 => self.reinterpret_descriptor::<i64, i64>(slot, shape, strides, offset),
1022 DType::Bool => self.reinterpret_descriptor::<bool, bool>(slot, shape, strides, offset),
1023 DType::C32 => self
1024 .reinterpret_descriptor::<num_complex::Complex32, num_complex::Complex32>(
1025 slot, shape, strides, offset,
1026 ),
1027 DType::C64 => self
1028 .reinterpret_descriptor::<num_complex::Complex64, num_complex::Complex64>(
1029 slot, shape, strides, offset,
1030 ),
1031 DType::External(_) => unreachable!("descriptors are preset-typed"),
1034 }
1035 }
1036
1037 #[allow(clippy::result_large_err)]
1045 pub(crate) fn reinterpret_descriptor<T: TensorScalar, U: TensorScalar>(
1046 mut self,
1047 slot: DescriptorSlot,
1048 shape: Vec<usize>,
1049 strides: Vec<isize>,
1050 offset: isize,
1051 ) -> Result<Self, (Self, GroupError)> {
1052 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1053 Ok((index, descriptor)) => (index, descriptor.clone()),
1054 Err(error) => return Err((self, error)),
1055 };
1056 if descriptor.dtype() != T::dtype() {
1057 return Err((
1058 self,
1059 GroupError::DTypeMismatch {
1060 expected: descriptor.dtype(),
1061 actual: T::dtype(),
1062 },
1063 ));
1064 }
1065 let allocation = descriptor.allocation;
1066 let references = self
1067 .descriptors
1068 .iter()
1069 .flatten()
1070 .filter(|candidate| candidate.allocation == allocation)
1071 .count();
1072 if references != 1 {
1073 return Err((
1074 self,
1075 GroupError::AliasedAllocation {
1076 allocation: allocation.index(),
1077 },
1078 ));
1079 }
1080
1081 let root = match self.allocation_root_resource(allocation) {
1082 Some(root) => root,
1083 None => {
1084 return Err((
1085 self,
1086 GroupError::AllocationSlotOutOfBounds {
1087 slot: allocation.index(),
1088 },
1089 ))
1090 }
1091 };
1092 let span = descriptor.span();
1093 let target_element_size = std::mem::size_of::<U>();
1094 if target_element_size == 0 || !span.byte_len().is_multiple_of(target_element_size) {
1095 return Err((
1096 self,
1097 GroupError::InvalidDescriptor {
1098 message: format!(
1099 "byte span {} is not divisible by target element size {}",
1100 span.byte_len(),
1101 target_element_size
1102 ),
1103 },
1104 ));
1105 }
1106 let layout = match TensorLayout::<DynRank>::from_parts(
1107 shape_vec(&shape),
1108 stride_vec(&strides),
1109 offset,
1110 span.byte_len() / target_element_size,
1111 ) {
1112 Ok(layout) => layout,
1113 Err(error) => {
1114 return Err((
1115 self,
1116 GroupError::InvalidDescriptor {
1117 message: error.to_string(),
1118 },
1119 ))
1120 }
1121 };
1122 let write_injective = descriptor.write_injective;
1123 let envelope = match reachable_envelope(&span, &layout, target_element_size) {
1124 Ok(envelope) => envelope,
1125 Err(error) => return Err((self, error)),
1126 };
1127 let (checked, _) = match validate_descriptor::<U, DynRank>(
1128 root,
1129 span,
1130 layout.shape().iter().copied().collect(),
1131 layout.strides().iter().copied().collect(),
1132 layout.offset(),
1133 write_injective,
1134 ) {
1135 Ok(value) => value,
1136 Err(error) => {
1137 return Err((
1138 self,
1139 GroupError::InvalidDescriptor {
1140 message: error.to_string(),
1141 },
1142 ))
1143 }
1144 };
1145 let element_count = match logical_element_count(layout.shape()) {
1146 Ok(count) => count,
1147 Err(error) => return Err((self, error)),
1148 };
1149 self.descriptors[descriptor_index] = Some(DescriptorRecord {
1150 allocation,
1151 element_count,
1152 provider: descriptor.provider,
1153 placement: descriptor.placement.clone(),
1154 envelope,
1155 write_injective,
1156 checked,
1157 });
1158 Ok(self)
1159 }
1160
1161 pub(crate) fn view<T: TensorScalar, R: TensorRank>(
1162 &self,
1163 slot: DescriptorSlot,
1164 ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1165 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1166 check_typed::<T, R>(descriptor)?;
1167 let owner = self
1168 .allocations
1169 .get(descriptor.allocation.index())
1170 .ok_or(GroupError::AllocationSlotOutOfBounds {
1171 slot: descriptor.allocation.index(),
1172 })?
1173 .as_ref()
1174 .ok_or(GroupError::AllocationSlotVacant {
1175 slot: descriptor.allocation.index(),
1176 })?;
1177 let _ = descriptor_index;
1178 Ok(GroupReadView {
1179 owner: NonNull::from(owner),
1180 descriptor: descriptor.clone(),
1181 _borrow: PhantomData,
1182 })
1183 }
1184
1185 pub(crate) fn view_raw<T: 'static, R: TensorRank>(
1186 &self,
1187 slot: DescriptorSlot,
1188 ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1189 let (_, descriptor) = self.resolve_descriptor(slot)?;
1190 let owner = self
1191 .allocations
1192 .get(descriptor.allocation.index())
1193 .ok_or(GroupError::AllocationSlotOutOfBounds {
1194 slot: descriptor.allocation.index(),
1195 })?
1196 .as_ref()
1197 .ok_or(GroupError::AllocationSlotVacant {
1198 slot: descriptor.allocation.index(),
1199 })?;
1200 Ok(GroupReadView {
1201 owner: NonNull::from(owner),
1202 descriptor: descriptor.clone(),
1203 _borrow: PhantomData,
1204 })
1205 }
1206
1207 pub(crate) fn prepare_device_read_for_layout<T: TensorScalar, R: TensorRank>(
1208 &self,
1209 slot: DescriptorSlot,
1210 layout: &TensorLayout<R>,
1211 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1212 self.view_raw::<T, R>(slot)
1213 .map_err(|error| AccessError::InvalidLayout {
1214 message: error.to_string(),
1215 })?
1216 .prepare_device_read_for_layout(layout)
1217 }
1218
1219 pub(crate) fn allocation_index(
1220 &self,
1221 slot: DescriptorSlot,
1222 ) -> Result<AllocationSlot, GroupError> {
1223 Ok(self.resolve_descriptor(slot)?.1.allocation)
1224 }
1225
1226 pub(crate) fn set_host_recycler<T: TensorScalar>(
1227 &mut self,
1228 allocation_index: usize,
1229 recycler: std::sync::Weak<dyn super::root::HostBufferRecycler<T>>,
1230 ) -> Result<(), AccessError> {
1231 let owner = self
1232 .allocations
1233 .get_mut(allocation_index)
1234 .and_then(Option::as_mut)
1235 .ok_or(AccessError::Unsupported {
1236 backend: "missing host allocation",
1237 })?;
1238 owner.set_host_recycler(recycler)
1239 }
1240
1241 pub(crate) fn host_buffer_at<T: 'static>(
1242 &self,
1243 allocation_index: AllocationSlot,
1244 ) -> Option<&crate::StorageBuffer<T>> {
1245 self.allocations
1246 .get(allocation_index.index())?
1247 .as_ref()?
1248 .host_buffer::<T>()
1249 }
1250
1251 pub(crate) fn host_root_metadata<T: TensorScalar>(
1252 &self,
1253 slot: DescriptorSlot,
1254 ) -> Option<(usize, usize)> {
1255 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1256 let owner = self
1257 .allocations
1258 .get(descriptor.allocation.index())?
1259 .as_ref()?;
1260 if descriptor.span() != owner.root_span() {
1261 return None;
1262 }
1263 let (pointer, byte_len) = match owner.host_buffer::<T>() {
1264 Some(crate::StorageBuffer::Host(data)) => (
1265 data.as_ptr() as usize,
1266 data.len().checked_mul(size_of::<T>())?,
1267 ),
1268 _ => {
1273 let slice = owner
1274 .as_ref()
1275 .host_slice::<T>(owner.root_span(), descriptor.dtype())
1276 .ok()?;
1277 (slice.as_ptr() as usize, size_of_val(slice))
1278 }
1279 };
1280 Some((pointer, byte_len))
1281 }
1282
1283 pub(crate) fn backend_buffer<T: 'static>(
1284 &self,
1285 slot: DescriptorSlot,
1286 ) -> Option<&crate::StorageBuffer<T>> {
1287 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1288 self.allocations
1289 .get(descriptor.allocation.index())?
1290 .as_ref()?
1291 .backend_buffer::<T>()
1292 }
1293
1294 pub(crate) fn descriptor_dtype(&self, slot: DescriptorSlot) -> Option<DType> {
1295 self.resolve_descriptor(slot)
1296 .ok()
1297 .map(|(_, descriptor)| descriptor.dtype())
1298 }
1299
1300 fn allocation_root_resource(&self, allocation: AllocationSlot) -> Option<RootResourceId> {
1306 self.allocations
1307 .get(allocation.index())?
1308 .as_ref()
1309 .map(|owner| owner.root_span().root_resource())
1310 }
1311
1312 pub(crate) fn descriptor_layout(&self, slot: DescriptorSlot) -> Option<TensorLayout<DynRank>> {
1313 self.resolve_descriptor(slot)
1314 .ok()
1315 .map(|(_, descriptor)| descriptor.layout().clone())
1316 }
1317
1318 pub(crate) fn descriptor_backing_len(&self, slot: DescriptorSlot) -> Option<usize> {
1323 self.resolve_descriptor(slot)
1324 .ok()
1325 .and_then(|(_, descriptor)| {
1326 descriptor
1327 .span()
1328 .byte_len()
1329 .checked_div(descriptor.element_size())
1330 })
1331 }
1332
1333 pub(crate) fn descriptor_logical_range(
1339 &self,
1340 slot: DescriptorSlot,
1341 ) -> Option<std::ops::Range<usize>> {
1342 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1343 if !descriptor.layout().is_compact_col_major().ok()? {
1344 return None;
1345 }
1346 let start = usize::try_from(descriptor.layout().offset()).ok()?;
1347 start
1348 .checked_add(descriptor.element_count())
1349 .map(|end| start..end)
1350 }
1351
1352 pub(crate) fn backend_identity(
1353 &self,
1354 slot: DescriptorSlot,
1355 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
1356 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1357 if descriptor.provider == ProviderKind::Cpu {
1358 return None;
1359 }
1360 let owner = self
1361 .allocations
1362 .get(descriptor.allocation.index())?
1363 .as_ref()?;
1364 let key = owner.root_identity().extent().key();
1365 Some((key.domain(), key.local()))
1366 }
1367
1368 pub(crate) fn provider_kind(&self, slot: DescriptorSlot) -> Option<ProviderKind> {
1369 self.resolve_descriptor(slot)
1370 .ok()
1371 .map(|(_, descriptor)| descriptor.provider)
1372 }
1373
1374 pub(crate) fn backend_root_buffer<T: 'static>(&self) -> Option<&crate::StorageBuffer<T>> {
1375 self.allocations
1376 .first()
1377 .and_then(|owner| owner.as_ref())
1378 .and_then(|owner| owner.backend_buffer::<T>())
1379 }
1380
1381 pub(crate) fn backend_buffer_mut<T: 'static>(
1382 &mut self,
1383 slot: DescriptorSlot,
1384 ) -> Option<&mut crate::StorageBuffer<T>> {
1385 let allocation = self.resolve_descriptor(slot).ok()?.1.allocation;
1386 self.allocations
1387 .get_mut(allocation.index())?
1388 .as_mut()?
1389 .backend_buffer_mut::<T>()
1390 }
1391
1392 pub(crate) fn backend_root_buffer_mut<T: 'static>(
1393 &mut self,
1394 ) -> Option<&mut crate::StorageBuffer<T>> {
1395 self.allocations
1396 .first_mut()
1397 .and_then(|owner| owner.as_mut())
1398 .and_then(|owner| owner.backend_buffer_mut::<T>())
1399 }
1400
1401 pub(crate) fn view_mut_raw<T: 'static, R: TensorRank>(
1402 &mut self,
1403 slot: DescriptorSlot,
1404 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1405 let (_, descriptor) = self.resolve_descriptor(slot)?;
1406 let descriptor = descriptor.clone();
1407 let owner = self
1408 .allocations
1409 .get_mut(descriptor.allocation.index())
1410 .ok_or(GroupError::AllocationSlotOutOfBounds {
1411 slot: descriptor.allocation.index(),
1412 })?
1413 .as_mut()
1414 .ok_or(GroupError::AllocationSlotVacant {
1415 slot: descriptor.allocation.index(),
1416 })?;
1417 Ok(GroupWriteView {
1418 owner: NonNull::from(owner),
1419 descriptor: descriptor.clone(),
1420 _borrow: PhantomData,
1421 })
1422 }
1423
1424 pub(crate) fn prepare_device_write_for_layout<T: TensorScalar, R: TensorRank>(
1425 &mut self,
1426 slot: DescriptorSlot,
1427 layout: &TensorLayout<R>,
1428 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1429 self.view_mut_raw::<T, R>(slot)
1430 .map_err(|error| AccessError::InvalidLayout {
1431 message: error.to_string(),
1432 })?
1433 .prepare_device_write_for_layout(layout)
1434 }
1435
1436 pub(crate) fn view_mut<T: TensorScalar, R: TensorRank>(
1437 &mut self,
1438 slot: DescriptorSlot,
1439 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1440 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1441 let mut descriptor = descriptor.clone();
1442 check_typed::<T, R>(&descriptor)?;
1443 if !descriptor.write_injective {
1444 descriptor
1445 .layout()
1446 .validate_mutable_no_overlap()
1447 .map_err(|error| GroupError::InvalidDescriptor {
1448 message: error.to_string(),
1449 })?;
1450 if let Some(Some(retained)) = self.descriptors.get_mut(descriptor_index) {
1451 retained.write_injective = true;
1452 }
1453 descriptor.write_injective = true;
1454 }
1455 let owner = self
1456 .allocations
1457 .get_mut(descriptor.allocation.index())
1458 .ok_or(GroupError::AllocationSlotOutOfBounds {
1459 slot: descriptor.allocation.index(),
1460 })?
1461 .as_mut()
1462 .ok_or(GroupError::AllocationSlotVacant {
1463 slot: descriptor.allocation.index(),
1464 })?;
1465 Ok(GroupWriteView {
1466 owner: NonNull::from(owner),
1467 descriptor,
1468 _borrow: PhantomData,
1469 })
1470 }
1471
1472 pub(crate) fn split_mut<T: TensorScalar, R: TensorRank>(
1473 &mut self,
1474 slots: &[DescriptorSlot],
1475 ) -> Result<Vec<GroupWriteView<'_, T, R>>, DisjointViewError> {
1476 let mut selected = Vec::with_capacity(slots.len());
1477 for &slot in slots {
1478 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1479 if selected
1480 .iter()
1481 .any(|(seen, _, _): &(DescriptorSlot, usize, DescriptorRecord)| *seen == slot)
1482 {
1483 return Err(DisjointViewError::DuplicateSlot { slot: slot.index() });
1484 }
1485 check_typed::<T, R>(descriptor)?;
1486 if !descriptor.write_injective {
1487 descriptor
1488 .layout()
1489 .validate_mutable_no_overlap()
1490 .map_err(|_| DisjointViewError::NonInjective { slot: slot.index() })?;
1491 }
1492 selected.push((slot, descriptor_index, descriptor.clone()));
1493 }
1494
1495 for left in 0..selected.len() {
1496 for right in (left + 1)..selected.len() {
1497 let first = &selected[left].2;
1498 let second = &selected[right].2;
1499 if first.span().root_resource() != second.span().root_resource() {
1500 continue;
1501 }
1502 match (first.envelope, second.envelope) {
1503 (None, _) | (_, None) => {}
1504 (Some(first), Some(second)) => {
1505 if first
1506 .overlaps(second)
1507 .map_err(|_| DisjointViewError::NotProvablyDisjoint)?
1508 {
1509 return Err(DisjointViewError::PairwiseOverlap);
1510 }
1511 }
1512 }
1513 }
1514 }
1515
1516 for (_, descriptor_index, _) in &selected {
1517 if let Some(Some(descriptor)) = self.descriptors.get_mut(*descriptor_index) {
1518 descriptor.write_injective = true;
1519 }
1520 }
1521
1522 let mut children = Vec::with_capacity(selected.len());
1523 for (_, _, mut descriptor) in selected {
1524 descriptor.write_injective = true;
1525 let owner = self
1526 .allocations
1527 .get_mut(descriptor.allocation.index())
1528 .ok_or(GroupError::AllocationSlotOutOfBounds {
1529 slot: descriptor.allocation.index(),
1530 })?
1531 .as_mut()
1532 .ok_or(GroupError::AllocationSlotVacant {
1533 slot: descriptor.allocation.index(),
1534 })?;
1535 children.push(GroupWriteView {
1536 owner: NonNull::from(owner),
1537 descriptor,
1538 _borrow: PhantomData,
1539 });
1540 }
1541 Ok(children)
1542 }
1543
1544 pub(crate) fn try_extract(
1545 &mut self,
1546 slot: DescriptorSlot,
1547 ) -> Result<OwnedStorage, ExtractError> {
1548 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1549 let allocation = descriptor.allocation;
1550 let references = self
1551 .descriptors
1552 .iter()
1553 .flatten()
1554 .filter(|candidate| candidate.allocation == allocation)
1555 .count();
1556 if references != 1 {
1557 return Err(ExtractError::AliasedAllocation {
1558 allocation: allocation.index(),
1559 });
1560 }
1561 let _ = self.descriptors[descriptor_index].take();
1562 self.allocations
1563 .get_mut(allocation.index())
1564 .ok_or(GroupError::AllocationSlotOutOfBounds {
1565 slot: allocation.index(),
1566 })?
1567 .take()
1568 .ok_or(GroupError::AllocationSlotVacant {
1569 slot: allocation.index(),
1570 })
1571 .map_err(ExtractError::from)
1572 }
1573
1574 #[allow(clippy::result_large_err)]
1587 pub fn take_tensor(&mut self, slot: DescriptorSlot) -> Result<crate::Tensor, GroupError> {
1588 let (_, descriptor) = self.resolve_descriptor(slot)?;
1589 self.ensure_whole_compact(slot, descriptor)?;
1590 let descriptor = descriptor.clone();
1591 let dtype = descriptor.dtype();
1592 let layout = descriptor.layout().clone();
1593 let placement = descriptor.placement.clone();
1594 let owner = self
1595 .try_extract(slot)
1596 .map_err(|error| GroupError::InvalidDescriptor {
1597 message: error.to_string(),
1598 })?;
1599 let mut extracted = Self::new();
1600 extracted.allocations.push(Some(owner));
1601 let mut descriptor = descriptor.clone();
1602 descriptor.allocation = AllocationSlot(0);
1603 extracted.descriptors.push(Some(descriptor));
1604 Ok(tensor_from_group(
1605 extracted,
1606 DescriptorSlot(0),
1607 AllocationSlot(0),
1608 dtype,
1609 layout,
1610 placement,
1611 ))
1612 }
1613
1614 #[allow(clippy::result_large_err)]
1617 pub(crate) fn into_owner(
1618 mut self,
1619 slot: DescriptorSlot,
1620 ) -> Result<OwnedStorage, (Self, ExtractError)> {
1621 let result = self.try_extract(slot);
1622 match result {
1623 Ok(owner) => Ok(owner),
1624 Err(error) => Err((self, error)),
1625 }
1626 }
1627
1628 #[allow(clippy::result_large_err)]
1644 pub fn into_tensor(self, slot: DescriptorSlot) -> Result<crate::Tensor, (Self, GroupError)> {
1645 let (_, descriptor) = match self.resolve_descriptor(slot) {
1646 Ok(value) => value,
1647 Err(error) => return Err((self, error)),
1648 };
1649 let allocation = descriptor.allocation;
1650 let references = self
1651 .descriptors
1652 .iter()
1653 .flatten()
1654 .filter(|candidate| candidate.allocation == allocation)
1655 .count();
1656 if references != 1 {
1657 return Err((
1658 self,
1659 GroupError::AliasedAllocation {
1660 allocation: allocation.index(),
1661 },
1662 ));
1663 }
1664 if let Err(error) = self.ensure_whole_compact(slot, descriptor) {
1665 return Err((self, error));
1666 }
1667 let dtype = descriptor.dtype();
1668 let layout = descriptor.layout().clone();
1669 let placement = descriptor.placement.clone();
1670 let (group, slot) = match self.into_single_descriptor(slot) {
1671 Ok(value) => value,
1672 Err((group, error)) => return Err((group, error)),
1673 };
1674 Ok(tensor_from_group(
1675 group,
1676 slot,
1677 AllocationSlot(0),
1678 dtype,
1679 layout,
1680 placement,
1681 ))
1682 }
1683
1684 #[allow(clippy::result_large_err)]
1687 fn into_single_descriptor(
1688 mut self,
1689 slot: DescriptorSlot,
1690 ) -> Result<(Self, DescriptorSlot), (Self, GroupError)> {
1691 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1692 Ok(value) => value,
1693 Err(error) => return Err((self, error)),
1694 };
1695 let allocation = descriptor.allocation;
1696 let mut descriptor = match self.descriptors[descriptor_index].take() {
1697 Some(descriptor) => descriptor,
1698 None => {
1699 return Err((
1700 self,
1701 GroupError::DescriptorSlotVacant { slot: slot.index() },
1702 ))
1703 }
1704 };
1705 let owner = match self.allocations.get_mut(allocation.index()) {
1706 Some(owner) => match owner.take() {
1707 Some(owner) => owner,
1708 None => {
1709 self.descriptors[descriptor_index] = Some(descriptor);
1710 return Err((
1711 self,
1712 GroupError::AllocationSlotVacant {
1713 slot: allocation.index(),
1714 },
1715 ));
1716 }
1717 },
1718 None => {
1719 self.descriptors[descriptor_index] = Some(descriptor);
1720 return Err((
1721 self,
1722 GroupError::AllocationSlotOutOfBounds {
1723 slot: allocation.index(),
1724 },
1725 ));
1726 }
1727 };
1728 let mut group = Self::new();
1729 descriptor.allocation = AllocationSlot(0);
1730 group.allocations.push(Some(owner));
1731 group.descriptors.push(Some(descriptor));
1732 Ok((group, DescriptorSlot(0)))
1733 }
1734
1735 #[allow(clippy::result_large_err)]
1741 pub(crate) fn into_host_vec<T: 'static>(
1742 mut self,
1743 slot: DescriptorSlot,
1744 ) -> Result<Vec<T>, (Self, String)> {
1745 let allocation = match self.resolve_descriptor(slot) {
1749 Ok((_, descriptor)) => descriptor.allocation,
1750 Err(error) => return Err((self, error.to_string())),
1751 };
1752 let owner = match self
1753 .allocations
1754 .get_mut(allocation.index())
1755 .map(Option::take)
1756 {
1757 Some(Some(owner)) => owner,
1758 Some(None) => {
1759 return Err((
1760 self,
1761 GroupError::AllocationSlotVacant {
1762 slot: allocation.index(),
1763 }
1764 .to_string(),
1765 ))
1766 }
1767 None => {
1768 return Err((
1769 self,
1770 GroupError::AllocationSlotOutOfBounds {
1771 slot: allocation.index(),
1772 }
1773 .to_string(),
1774 ))
1775 }
1776 };
1777 match owner.into_host_vec::<T>() {
1778 Ok(data) => Ok(data),
1779 Err((owner, error)) => {
1780 self.allocations[allocation.index()] = Some(owner);
1781 Err((self, error.to_string()))
1782 }
1783 }
1784 }
1785
1786 fn ensure_whole_compact(
1800 &self,
1801 slot: DescriptorSlot,
1802 descriptor: &DescriptorRecord,
1803 ) -> Result<(), GroupError> {
1804 let index = descriptor.allocation.index();
1805 let owner = self
1806 .allocations
1807 .get(index)
1808 .ok_or(GroupError::AllocationSlotOutOfBounds { slot: index })?
1809 .as_ref()
1810 .ok_or(GroupError::AllocationSlotVacant { slot: index })?;
1811 let root = owner.as_ref().root_identity().root_span();
1812 let span = descriptor.span();
1813 let layout = descriptor.layout();
1814 let compact =
1815 layout
1816 .is_compact_col_major()
1817 .map_err(|error| GroupError::InvalidDescriptor {
1818 message: error.to_string(),
1819 })?;
1820 let root_spanning =
1821 span.byte_offset() == root.byte_offset() && span.byte_len() == root.byte_len();
1822 let whole = root_spanning
1823 && descriptor
1824 .element_count()
1825 .checked_mul(descriptor.element_size())
1826 == Some(span.byte_len());
1827 let host_owned = owner.backend_allocation().is_none();
1828 if (compact && layout.offset() == 0 && whole) || (compact && root_spanning && host_owned) {
1829 Ok(())
1830 } else {
1831 Err(GroupError::NonCompactDescriptor { slot: slot.index() })
1832 }
1833 }
1834
1835 fn resolve_descriptor(
1836 &self,
1837 slot: DescriptorSlot,
1838 ) -> Result<(usize, &DescriptorRecord), GroupError> {
1839 let index = slot.index();
1840 let descriptor = self
1841 .descriptors
1842 .get(index)
1843 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: index })?
1844 .as_ref()
1845 .ok_or(GroupError::DescriptorSlotVacant { slot: index })?;
1846 Ok((index, descriptor))
1847 }
1848
1849 #[cfg(test)]
1850 pub(crate) fn test_vacate_allocation(&mut self, slot: AllocationSlot) {
1851 if let Some(entry) = self.allocations.get_mut(slot.index()) {
1852 *entry = None;
1853 }
1854 }
1855}
1856
1857fn check_typed<T: TensorScalar, R: TensorRank>(
1858 descriptor: &DescriptorRecord,
1859) -> Result<(), GroupError> {
1860 if descriptor.dtype() != T::dtype() {
1861 return Err(GroupError::DTypeMismatch {
1862 expected: descriptor.dtype(),
1863 actual: T::dtype(),
1864 });
1865 }
1866 if let Some(expected) = R::RANK {
1867 let actual = descriptor.layout().shape().len();
1868 if expected != actual {
1869 return Err(GroupError::RankMismatch { expected, actual });
1870 }
1871 }
1872 Ok(())
1873}
1874
1875fn logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1876 shape.iter().try_fold(1usize, |count, &extent| {
1877 count
1878 .checked_mul(extent)
1879 .ok_or_else(|| GroupError::InvalidDescriptor {
1880 message: "logical element count overflows".to_owned(),
1881 })
1882 })
1883}
1884
1885fn reachable_envelope(
1886 span: &RootBoundSpan,
1887 layout: &TensorLayout<DynRank>,
1888 element_size: usize,
1889) -> Result<Option<ByteRange>, GroupError> {
1890 if layout.shape().contains(&0) {
1891 return Ok(None);
1892 }
1893 let mut minimum = layout.offset() as i128;
1894 let mut maximum = minimum;
1895 for (&extent, &stride) in layout.shape().iter().zip(layout.strides()) {
1896 let steps = i128::try_from(extent - 1).map_err(|_| GroupError::InvalidDescriptor {
1897 message: "layout extent does not fit i128".to_owned(),
1898 })?;
1899 let contribution =
1900 (stride as i128)
1901 .checked_mul(steps)
1902 .ok_or_else(|| GroupError::InvalidDescriptor {
1903 message: "reachable layout arithmetic overflows".to_owned(),
1904 })?;
1905 if contribution < 0 {
1906 minimum =
1907 minimum
1908 .checked_add(contribution)
1909 .ok_or_else(|| GroupError::InvalidDescriptor {
1910 message: "reachable layout minimum overflows".to_owned(),
1911 })?;
1912 } else {
1913 maximum =
1914 maximum
1915 .checked_add(contribution)
1916 .ok_or_else(|| GroupError::InvalidDescriptor {
1917 message: "reachable layout maximum overflows".to_owned(),
1918 })?;
1919 }
1920 }
1921 let minimum = usize::try_from(minimum).map_err(|_| GroupError::InvalidDescriptor {
1922 message: "reachable layout minimum is negative".to_owned(),
1923 })?;
1924 let maximum = usize::try_from(maximum).map_err(|_| GroupError::InvalidDescriptor {
1925 message: "reachable layout maximum is negative or too large".to_owned(),
1926 })?;
1927 let byte_offset = span
1928 .byte_offset()
1929 .checked_add(minimum.checked_mul(element_size).ok_or_else(|| {
1930 GroupError::InvalidDescriptor {
1931 message: "reachable byte offset overflows".to_owned(),
1932 }
1933 })?)
1934 .ok_or_else(|| GroupError::InvalidDescriptor {
1935 message: "reachable byte offset overflows".to_owned(),
1936 })?;
1937 let byte_len = maximum
1938 .checked_sub(minimum)
1939 .and_then(|length| length.checked_add(1))
1940 .and_then(|length| length.checked_mul(element_size))
1941 .ok_or_else(|| GroupError::InvalidDescriptor {
1942 message: "reachable byte length overflows".to_owned(),
1943 })?;
1944 let range = ByteRange::new(byte_offset, byte_len);
1945 range
1946 .checked_end()
1947 .map_err(|error| GroupError::InvalidDescriptor {
1948 message: error.to_string(),
1949 })?;
1950 Ok(Some(range))
1951}
1952
1953#[cfg(test)]
1954pub(crate) fn test_logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1955 logical_element_count(shape)
1956}
1957
1958#[cfg(test)]
1959pub(crate) fn test_reachable_envelope(
1960 span: RootBoundSpan,
1961 layout: TensorLayout<DynRank>,
1962 element_size: usize,
1963) -> Result<Option<ByteRange>, GroupError> {
1964 reachable_envelope(&span, &layout, element_size)
1965}