1use std::fmt;
2use std::marker::PhantomData;
3use std::mem::{align_of, size_of};
4use std::ptr::NonNull;
5
6use crate::types::{tensor_from_group, tensor_view_from_group};
7use crate::{DType, DynRank, Placement, TensorLayout, TensorRank, TensorRead, TensorScalar};
8use smallvec::SmallVec;
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}
220
221#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
223pub(crate) enum DisjointViewError {
224 #[error(transparent)]
225 Group(#[from] GroupError),
226 #[error("descriptor slot {slot} appears more than once")]
227 DuplicateSlot { slot: usize },
228 #[error("descriptor slot {slot} has a non-injective mutable layout")]
229 NonInjective { slot: usize },
230 #[error("requested mutable descriptor envelopes overlap")]
231 PairwiseOverlap,
232 #[error("requested mutable descriptors are not provably disjoint")]
233 NotProvablyDisjoint,
234}
235
236#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
238pub(crate) enum ExtractError {
239 #[error(transparent)]
240 Group(#[from] GroupError),
241 #[error("allocation slot {allocation} still has another descriptor")]
242 AliasedAllocation { allocation: usize },
243}
244
245#[derive(Default)]
258pub struct AllocationGroup {
259 allocations: SmallVec<[Option<OwnedStorage>; 1]>,
264 descriptors: SmallVec<[Option<DescriptorRecord>; 1]>,
265}
266
267pub(crate) struct GroupReadView<'a, T, R: TensorRank> {
269 owner: NonNull<OwnedStorage>,
270 descriptor: DescriptorRecord,
271 _borrow: PhantomData<(&'a OwnedStorage, T, R)>,
272}
273
274unsafe impl<'a, T: Send, R: TensorRank> Send for GroupReadView<'a, T, R> {}
277unsafe impl<'a, T: Sync, R: TensorRank> Sync for GroupReadView<'a, T, R> {}
278
279impl<'a, T, R: TensorRank> Clone for GroupReadView<'a, T, R> {
280 fn clone(&self) -> Self {
281 Self {
282 owner: self.owner,
283 descriptor: self.descriptor.clone(),
284 _borrow: PhantomData,
285 }
286 }
287}
288
289impl<'a, T, R: TensorRank> GroupReadView<'a, T, R> {
290 pub(crate) fn clone_dyn(&self) -> GroupReadView<'a, T, crate::DynRank> {
291 GroupReadView {
292 owner: self.owner,
293 descriptor: self.descriptor.clone(),
294 _borrow: PhantomData,
295 }
296 }
297}
298
299impl<T, R: TensorRank> std::fmt::Debug for GroupReadView<'_, T, R> {
300 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
301 formatter
302 .debug_struct("GroupReadView")
303 .field("descriptor", &self.descriptor)
304 .finish_non_exhaustive()
305 }
306}
307
308impl<'a, T: TensorScalar, R: TensorRank> GroupReadView<'a, T, R> {
309 pub(crate) fn descriptor(&self) -> &DescriptorRecord {
310 &self.descriptor
311 }
312
313 pub(crate) fn provider_kind(&self) -> ProviderKind {
314 self.descriptor.provider
315 }
316
317 pub(crate) fn backend_identity(
318 &self,
319 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
320 if self.descriptor.provider == ProviderKind::Cpu {
321 return None;
322 }
323 let key = unsafe { self.owner.as_ref().root_identity().extent().key() };
324 Some((key.domain(), key.local()))
325 }
326
327 pub(crate) fn storage_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
328 let buffer = unsafe {
331 self.owner
332 .as_ref()
333 .host_buffer::<T>()
334 .or_else(|| self.owner.as_ref().backend_buffer::<T>())
335 }?;
336 Some(unsafe {
337 std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
338 })
339 }
340
341 pub(crate) fn map_read(&self) -> Result<ProviderReadMapping<'_>, AccessError> {
342 unsafe {
345 self.owner
346 .as_ref()
347 .as_ref()
348 .map_read(self.descriptor.span(), self.descriptor.dtype())
349 }
350 }
351
352 pub(crate) fn backend_allocation(&self) -> Option<&'a dyn BackendAllocation> {
353 let allocation = unsafe { self.owner.as_ref().backend_allocation() }?;
354 Some(unsafe {
355 std::mem::transmute::<&dyn BackendAllocation, &'a dyn BackendAllocation>(allocation)
356 })
357 }
358
359 pub(crate) fn prepare_device_read_for_layout(
360 &self,
361 layout: &TensorLayout<R>,
362 ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
363 self.prepare_read_for_layout(layout, AccessTarget::Device)?
364 .into_device_state()
365 }
366
367 pub(crate) fn prepare_host_read_for_layout(
368 &self,
369 layout: &TensorLayout<R>,
370 ) -> Result<PreparedRead<'a, T, R>, AccessError> {
371 self.prepare_read_for_layout(layout, AccessTarget::Host)
372 }
373
374 fn prepare_read_for_layout(
375 &self,
376 layout: &TensorLayout<R>,
377 target: AccessTarget,
378 ) -> Result<PreparedRead<'a, T, R>, AccessError> {
379 let owner: crate::storage::root::StorageRef<'a> = unsafe { self.owner.as_ref().as_ref() };
381 let checked: CheckedRead<'a, R> = CheckedRead::new::<T>(
382 owner,
384 self.descriptor.span(),
385 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
386 AccessError::InvalidLayout {
387 message: error.to_string(),
388 }
389 })?,
390 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
391 AccessError::InvalidLayout {
392 message: error.to_string(),
393 }
394 })?,
395 layout.offset(),
396 )?;
397 prepare_read::<T, R>(checked, target).map_err(|failure| failure.1)
398 }
399
400 pub(crate) fn prepare_host_read(&self) -> Result<PreparedRead<'_, T, DynRank>, AccessError> {
401 prepare_read(
402 CheckedRead::from_validated(
403 unsafe { self.owner.as_ref().as_ref() },
405 self.descriptor.checked.clone(),
406 ),
407 AccessTarget::Host,
408 )
409 .map_err(|failure| failure.1)
410 }
411
412 pub(crate) fn host_slice(&self) -> Result<&'a [T], AccessError> {
413 unsafe {
415 self.owner
416 .as_ref()
417 .as_ref()
418 .host_slice(self.descriptor.span(), self.descriptor.dtype())
419 }
420 }
421}
422
423impl<'a, T: 'static, R: TensorRank> GroupReadView<'a, T, R> {
424 pub(crate) fn backend_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
425 let buffer = unsafe { self.owner.as_ref().backend_buffer::<T>() }?;
428 Some(unsafe {
429 std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
430 })
431 }
432}
433
434pub(crate) struct GroupWriteView<'a, T, R: TensorRank> {
438 owner: NonNull<OwnedStorage>,
439 descriptor: DescriptorRecord,
440 _borrow: PhantomData<(&'a mut [u8], T, R)>,
441}
442
443unsafe impl<'a, T: Send, R: TensorRank> Send for GroupWriteView<'a, T, R> {}
446
447impl<T, R: TensorRank> std::fmt::Debug for GroupWriteView<'_, T, R> {
448 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
449 formatter
450 .debug_struct("GroupWriteView")
451 .field("descriptor", &self.descriptor)
452 .finish_non_exhaustive()
453 }
454}
455
456impl<'a, T: TensorScalar, R: TensorRank> GroupWriteView<'a, T, R> {
457 pub(crate) fn descriptor(&self) -> &DescriptorRecord {
458 &self.descriptor
459 }
460
461 pub(crate) fn map_write(&mut self) -> Result<ProviderWriteMapping<'_>, AccessError> {
462 unsafe {
466 self.owner
467 .as_mut()
468 .as_mut()
469 .map_write(self.descriptor.span(), self.descriptor.dtype())
470 }
471 }
472
473 pub(crate) fn backend_buffer_mut(&mut self) -> Option<&'a mut crate::StorageBuffer<T>> {
474 let owner = unsafe { &mut *self.owner.as_ptr() };
477 let buffer = owner.backend_buffer_mut::<T>()?;
478 Some(unsafe {
479 std::mem::transmute::<&mut crate::StorageBuffer<T>, &'a mut crate::StorageBuffer<T>>(
480 buffer,
481 )
482 })
483 }
484
485 pub(crate) fn prepare_device_write_for_layout(
486 &mut self,
487 layout: &TensorLayout<R>,
488 ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
489 let owner: crate::storage::root::StorageMut<'a> =
490 unsafe { (&mut *self.owner.as_ptr()).as_mut() };
491 let checked: CheckedWrite<'a, R> = CheckedWrite::new::<T>(
492 owner,
494 self.descriptor.span(),
495 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
496 AccessError::InvalidLayout {
497 message: error.to_string(),
498 }
499 })?,
500 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
501 AccessError::InvalidLayout {
502 message: error.to_string(),
503 }
504 })?,
505 layout.offset(),
506 )?;
507 prepare_write::<T, R>(checked, AccessTarget::Device)
508 .map_err(|failure| failure.1)?
509 .into_device_state()
510 }
511
512 pub(crate) fn prepare_host_write(
513 &mut self,
514 ) -> Result<PreparedWrite<'_, T, DynRank>, AccessError> {
515 let checked = CheckedWrite::from_validated(
516 unsafe { self.owner.as_mut().as_mut() },
518 self.descriptor.checked.clone(),
519 WriteInjectivityProof,
520 );
521 prepare_write(checked, AccessTarget::Host).map_err(|failure| failure.1)
522 }
523
524 pub(crate) fn host_slice_mut(&mut self) -> Result<&'a mut [T], AccessError> {
525 unsafe {
528 self.owner
529 .as_mut()
530 .as_mut()
531 .host_slice_mut(self.descriptor.span(), self.descriptor.dtype())
532 }
533 }
534}
535
536impl<'a, T: 'static, R: TensorRank> GroupWriteView<'a, T, R> {
537 pub(crate) fn backend_buffer(&self) -> Option<&crate::StorageBuffer<T>> {
538 unsafe { self.owner.as_ref().backend_buffer::<T>() }
541 }
542}
543
544impl fmt::Debug for AllocationGroup {
545 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
546 formatter
547 .debug_struct("AllocationGroup")
548 .field("allocation_count", &self.allocations.len())
549 .field("descriptor_count", &self.descriptors.len())
550 .finish()
551 }
552}
553
554impl AllocationGroup {
555 pub(crate) fn new() -> Self {
556 Self::default()
557 }
558
559 pub fn from_tensors(
578 tensors: Vec<crate::Tensor>,
579 ) -> Result<(Self, Box<[DescriptorSlot]>), GroupError> {
580 let mut group = Self::new();
581 let mut bindings = Vec::with_capacity(tensors.len());
582 for tensor in tensors {
583 let (source, source_slot) = into_group_parts(tensor)?;
584 bindings.push(group.append_group(source, source_slot)?);
585 }
586 Ok((group, bindings.into_boxed_slice()))
587 }
588
589 pub fn read_views<'a>(
598 &'a self,
599 bindings: &[DescriptorSlot],
600 ) -> Result<Vec<TensorRead<'a>>, GroupError> {
601 bindings
602 .iter()
603 .map(|&slot| self.tensor_read(slot))
604 .collect()
605 }
606
607 pub fn read_view<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
616 self.tensor_read(slot)
617 }
618
619 fn tensor_read<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
620 let dtype = self.resolve_descriptor(slot)?.1.dtype();
621 let view = match dtype {
622 DType::F32 => tensor_view_from_group(self.view::<f32, DynRank>(slot)?),
623 DType::F64 => tensor_view_from_group(self.view::<f64, DynRank>(slot)?),
624 DType::I32 => tensor_view_from_group(self.view::<i32, DynRank>(slot)?),
625 DType::I64 => tensor_view_from_group(self.view::<i64, DynRank>(slot)?),
626 DType::Bool => tensor_view_from_group(self.view::<bool, DynRank>(slot)?),
627 DType::C32 => {
628 tensor_view_from_group(self.view::<num_complex::Complex32, DynRank>(slot)?)
629 }
630 DType::C64 => {
631 tensor_view_from_group(self.view::<num_complex::Complex64, DynRank>(slot)?)
632 }
633 DType::External(_) => unreachable!("descriptors are preset-typed"),
636 }
637 .map_err(|error| GroupError::InvalidDescriptor {
638 message: error.to_string(),
639 })?;
640 Ok(TensorRead::from_view(view))
641 }
642
643 pub fn append_tensor(&mut self, tensor: crate::Tensor) -> Result<DescriptorSlot, GroupError> {
652 let (source, source_slot) = into_group_parts(tensor)?;
653 self.append_group(source, source_slot)
654 }
655
656 pub fn append_group(
663 &mut self,
664 mut source: AllocationGroup,
665 source_slot: DescriptorSlot,
666 ) -> Result<DescriptorSlot, GroupError> {
667 let allocation_offset =
668 u32::try_from(self.allocations.len()).map_err(|_| GroupError::IndexOverflow)?;
669 let descriptor_offset =
670 u32::try_from(self.descriptors.len()).map_err(|_| GroupError::IndexOverflow)?;
671 source.resolve_descriptor(source_slot)?;
672
673 for owner in source.allocations.drain(..) {
674 self.allocations.push(owner);
675 }
676 for descriptor in source.descriptors.drain(..) {
677 let descriptor = match descriptor {
678 Some(mut descriptor) => {
679 let allocation = descriptor
680 .allocation
681 .0
682 .checked_add(allocation_offset)
683 .ok_or(GroupError::IndexOverflow)?;
684 descriptor.allocation = AllocationSlot(allocation);
685 Some(descriptor)
686 }
687 None => None,
688 };
689 self.descriptors.push(descriptor);
690 }
691
692 let source_descriptor_index =
693 u32::try_from(source_slot.index()).map_err(|_| GroupError::IndexOverflow)?;
694 Ok(DescriptorSlot(
695 source_descriptor_index
696 .checked_add(descriptor_offset)
697 .ok_or(GroupError::IndexOverflow)?,
698 ))
699 }
700
701 pub(crate) fn set_descriptor_placement(
702 &mut self,
703 slot: DescriptorSlot,
704 placement: Placement,
705 ) -> Result<(), GroupError> {
706 let descriptor = self
707 .descriptors
708 .get_mut(slot.index())
709 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: slot.index() })?
710 .as_mut()
711 .ok_or(GroupError::DescriptorSlotVacant { slot: slot.index() })?;
712 descriptor.placement = placement;
713 Ok(())
714 }
715
716 pub(crate) fn publish_live_descriptor_placement(
725 &mut self,
726 slot: DescriptorSlot,
727 placement: Placement,
728 ) {
729 let descriptor = self
730 .descriptors
731 .get_mut(slot.index())
732 .and_then(Option::as_mut)
733 .unwrap_or_else(|| unreachable!("an owned group always carries its descriptor"));
734 descriptor.placement = placement;
735 }
736
737 pub(crate) fn from_host_vec<T: TensorScalar, R: TensorRank>(
738 shape: R::Shape,
739 data: Vec<T>,
740 ) -> Result<(Self, DescriptorSlot), GroupError> {
741 let owner =
742 super::root::import_host_vec(data).map_err(|error| GroupError::InvalidDescriptor {
743 message: error.to_string(),
744 })?;
745 let span = owner.root_span();
746 let mut group = Self::new();
747 let allocation = group.insert_owner(owner)?;
748 let layout =
749 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
750 message: error.to_string(),
751 })?;
752 let input = DescriptorInput::new(
753 ByteRange::new(0, span.byte_len()),
754 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
755 GroupError::InvalidDescriptor {
756 message: error.to_string(),
757 }
758 })?,
759 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
760 GroupError::InvalidDescriptor {
761 message: error.to_string(),
762 }
763 })?,
764 layout.offset(),
765 true,
766 );
767 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
768 Ok((group, slot))
769 }
770
771 #[doc(hidden)]
773 pub fn from_backend_allocation<T: TensorScalar, R: TensorRank>(
774 shape: R::Shape,
775 allocation: Box<dyn BackendAllocation>,
776 ) -> Result<(Self, DescriptorSlot), GroupError> {
777 let owner = super::root::import_unique_root(allocation).map_err(|error| {
778 GroupError::InvalidDescriptor {
779 message: error.to_string(),
780 }
781 })?;
782 let span = owner.root_span();
783 let mut group = Self::new();
784 let allocation = group.insert_owner(owner)?;
785 let layout =
786 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
787 message: error.to_string(),
788 })?;
789 let input = DescriptorInput::new(
790 ByteRange::new(0, span.byte_len()),
791 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
792 GroupError::InvalidDescriptor {
793 message: error.to_string(),
794 }
795 })?,
796 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
797 GroupError::InvalidDescriptor {
798 message: error.to_string(),
799 }
800 })?,
801 layout.offset(),
802 true,
803 );
804 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
805 Ok((group, slot))
806 }
807
808 pub(crate) fn from_backend_buffer<T: TensorScalar, R: TensorRank>(
809 shape: R::Shape,
810 buffer: crate::StorageBuffer<T>,
811 ) -> Result<(Self, DescriptorSlot), GroupError> {
812 let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
813 GroupError::InvalidDescriptor {
814 message: error.to_string(),
815 }
816 })?;
817 let span = owner.root_span();
818 let mut group = Self::new();
819 let allocation = group.insert_owner(owner)?;
820 let layout =
821 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
822 message: error.to_string(),
823 })?;
824 let input = DescriptorInput::new(
825 ByteRange::new(0, span.byte_len()),
826 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
827 GroupError::InvalidDescriptor {
828 message: error.to_string(),
829 }
830 })?,
831 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
832 GroupError::InvalidDescriptor {
833 message: error.to_string(),
834 }
835 })?,
836 layout.offset(),
837 true,
838 );
839 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
840 Ok((group, slot))
841 }
842
843 pub(crate) fn from_backend_root<T: Send + Sync + 'static>(
844 buffer: crate::StorageBuffer<T>,
845 ) -> Result<Self, GroupError> {
846 let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
847 GroupError::InvalidDescriptor {
848 message: error.to_string(),
849 }
850 })?;
851 let mut group = Self::new();
852 group.insert_owner(owner)?;
853 Ok(group)
854 }
855
856 pub(crate) fn insert_owner(
857 &mut self,
858 owner: OwnedStorage,
859 ) -> Result<AllocationSlot, GroupError> {
860 let index = self.allocations.len();
861 let slot = u32::try_from(index).map_err(|_| GroupError::IndexOverflow)?;
862 self.allocations.push(Some(owner));
863 Ok(AllocationSlot(slot))
864 }
865
866 pub(crate) fn insert_descriptor<T: TensorScalar, R: TensorRank>(
867 &mut self,
868 allocation: AllocationSlot,
869 input: DescriptorInput<R>,
870 ) -> Result<DescriptorSlot, GroupError> {
871 let allocation_index = allocation.index();
872 let owner = self
873 .allocations
874 .get(allocation_index)
875 .ok_or(GroupError::AllocationSlotOutOfBounds {
876 slot: allocation_index,
877 })?
878 .as_ref()
879 .ok_or(GroupError::AllocationSlotVacant {
880 slot: allocation_index,
881 })?;
882
883 let root = owner.as_ref().root_identity();
884 let span = root.bind_relative_range(input.relative).map_err(|error| {
885 GroupError::InvalidDescriptor {
886 message: error.to_string(),
887 }
888 })?;
889 let element_size = size_of::<T>();
890 if element_size == 0 || !span.byte_len().is_multiple_of(element_size) {
891 return Err(GroupError::InvalidDescriptor {
892 message: format!(
893 "byte span {} is not divisible by element size {}",
894 span.byte_len(),
895 element_size
896 ),
897 });
898 }
899 if !span
900 .guaranteed_alignment()
901 .get()
902 .is_multiple_of(align_of::<T>())
903 {
904 return Err(GroupError::InvalidDescriptor {
905 message: format!(
906 "span alignment {} is insufficient for {}-byte alignment",
907 span.guaranteed_alignment().get(),
908 align_of::<T>()
909 ),
910 });
911 }
912
913 let shape = R::shape_into_vec(input.shape);
914 let strides = R::strides_into_vec(input.strides);
915 let layout = TensorLayout::<DynRank>::from_parts(
916 shape,
917 strides,
918 input.offset,
919 span.byte_len() / element_size,
920 )
921 .map_err(|error| GroupError::InvalidDescriptor {
922 message: error.to_string(),
923 })?;
924 let element_count = logical_element_count(layout.shape())?;
925 let write_injective = if input.require_injective {
926 layout.validate_mutable_no_overlap().map_err(|error| {
927 GroupError::InvalidDescriptor {
928 message: error.to_string(),
929 }
930 })?;
931 true
932 } else {
933 false
934 };
935 let envelope = reachable_envelope(&span, &layout, element_size)?;
936 let (checked, _) = validate_descriptor::<T, DynRank>(
937 root.root_resource(),
938 span,
939 layout.shape().iter().copied().collect(),
940 layout.strides().iter().copied().collect(),
941 layout.offset(),
942 false,
943 )
944 .map_err(|error| GroupError::InvalidDescriptor {
945 message: error.to_string(),
946 })?;
947 let record = DescriptorRecord {
948 allocation,
949 element_count,
950 provider: owner.as_ref().provider_kind(),
951 placement: Placement::default(),
952 envelope,
953 write_injective,
954 checked,
955 };
956
957 let descriptor_index = self.descriptors.len();
958 let slot = u32::try_from(descriptor_index).map_err(|_| GroupError::IndexOverflow)?;
959 self.descriptors.push(Some(record));
960 Ok(DescriptorSlot(slot))
961 }
962
963 #[allow(clippy::result_large_err)]
968 pub(crate) fn update_descriptor_layout(
969 mut self,
970 slot: DescriptorSlot,
971 shape: Vec<usize>,
972 strides: Vec<isize>,
973 offset: isize,
974 ) -> Result<Self, (Self, GroupError)> {
975 let dtype = match self.resolve_descriptor(slot) {
976 Ok((_, descriptor)) => descriptor.dtype(),
977 Err(error) => return Err((self, error)),
978 };
979 if let Some(Some(descriptor)) = self.descriptors.get_mut(slot.index()) {
980 descriptor.write_injective = false;
983 }
984 match dtype {
985 DType::F32 => self.reinterpret_descriptor::<f32, f32>(slot, shape, strides, offset),
986 DType::F64 => self.reinterpret_descriptor::<f64, f64>(slot, shape, strides, offset),
987 DType::I32 => self.reinterpret_descriptor::<i32, i32>(slot, shape, strides, offset),
988 DType::I64 => self.reinterpret_descriptor::<i64, i64>(slot, shape, strides, offset),
989 DType::Bool => self.reinterpret_descriptor::<bool, bool>(slot, shape, strides, offset),
990 DType::C32 => self
991 .reinterpret_descriptor::<num_complex::Complex32, num_complex::Complex32>(
992 slot, shape, strides, offset,
993 ),
994 DType::C64 => self
995 .reinterpret_descriptor::<num_complex::Complex64, num_complex::Complex64>(
996 slot, shape, strides, offset,
997 ),
998 DType::External(_) => unreachable!("descriptors are preset-typed"),
1001 }
1002 }
1003
1004 #[allow(clippy::result_large_err)]
1012 pub(crate) fn reinterpret_descriptor<T: TensorScalar, U: TensorScalar>(
1013 mut self,
1014 slot: DescriptorSlot,
1015 shape: Vec<usize>,
1016 strides: Vec<isize>,
1017 offset: isize,
1018 ) -> Result<Self, (Self, GroupError)> {
1019 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1020 Ok((index, descriptor)) => (index, descriptor.clone()),
1021 Err(error) => return Err((self, error)),
1022 };
1023 if descriptor.dtype() != T::dtype() {
1024 return Err((
1025 self,
1026 GroupError::DTypeMismatch {
1027 expected: descriptor.dtype(),
1028 actual: T::dtype(),
1029 },
1030 ));
1031 }
1032 let allocation = descriptor.allocation;
1033 let references = self
1034 .descriptors
1035 .iter()
1036 .flatten()
1037 .filter(|candidate| candidate.allocation == allocation)
1038 .count();
1039 if references != 1 {
1040 return Err((
1041 self,
1042 GroupError::AliasedAllocation {
1043 allocation: allocation.index(),
1044 },
1045 ));
1046 }
1047
1048 let root = match self.allocation_root_resource(allocation) {
1049 Some(root) => root,
1050 None => {
1051 return Err((
1052 self,
1053 GroupError::AllocationSlotOutOfBounds {
1054 slot: allocation.index(),
1055 },
1056 ))
1057 }
1058 };
1059 let span = descriptor.span();
1060 let target_element_size = std::mem::size_of::<U>();
1061 if target_element_size == 0 || !span.byte_len().is_multiple_of(target_element_size) {
1062 return Err((
1063 self,
1064 GroupError::InvalidDescriptor {
1065 message: format!(
1066 "byte span {} is not divisible by target element size {}",
1067 span.byte_len(),
1068 target_element_size
1069 ),
1070 },
1071 ));
1072 }
1073 let layout = match TensorLayout::<DynRank>::from_parts(
1074 shape.into(),
1075 strides.into(),
1076 offset,
1077 span.byte_len() / target_element_size,
1078 ) {
1079 Ok(layout) => layout,
1080 Err(error) => {
1081 return Err((
1082 self,
1083 GroupError::InvalidDescriptor {
1084 message: error.to_string(),
1085 },
1086 ))
1087 }
1088 };
1089 let write_injective = descriptor.write_injective;
1090 let envelope = match reachable_envelope(&span, &layout, target_element_size) {
1091 Ok(envelope) => envelope,
1092 Err(error) => return Err((self, error)),
1093 };
1094 let (checked, _) = match validate_descriptor::<U, DynRank>(
1095 root,
1096 span,
1097 layout.shape().iter().copied().collect(),
1098 layout.strides().iter().copied().collect(),
1099 layout.offset(),
1100 write_injective,
1101 ) {
1102 Ok(value) => value,
1103 Err(error) => {
1104 return Err((
1105 self,
1106 GroupError::InvalidDescriptor {
1107 message: error.to_string(),
1108 },
1109 ))
1110 }
1111 };
1112 let element_count = match logical_element_count(layout.shape()) {
1113 Ok(count) => count,
1114 Err(error) => return Err((self, error)),
1115 };
1116 self.descriptors[descriptor_index] = Some(DescriptorRecord {
1117 allocation,
1118 element_count,
1119 provider: descriptor.provider,
1120 placement: descriptor.placement.clone(),
1121 envelope,
1122 write_injective,
1123 checked,
1124 });
1125 Ok(self)
1126 }
1127
1128 pub(crate) fn view<T: TensorScalar, R: TensorRank>(
1129 &self,
1130 slot: DescriptorSlot,
1131 ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1132 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1133 check_typed::<T, R>(descriptor)?;
1134 let owner = self
1135 .allocations
1136 .get(descriptor.allocation.index())
1137 .ok_or(GroupError::AllocationSlotOutOfBounds {
1138 slot: descriptor.allocation.index(),
1139 })?
1140 .as_ref()
1141 .ok_or(GroupError::AllocationSlotVacant {
1142 slot: descriptor.allocation.index(),
1143 })?;
1144 let _ = descriptor_index;
1145 Ok(GroupReadView {
1146 owner: NonNull::from(owner),
1147 descriptor: descriptor.clone(),
1148 _borrow: PhantomData,
1149 })
1150 }
1151
1152 pub(crate) fn view_raw<T: 'static, R: TensorRank>(
1153 &self,
1154 slot: DescriptorSlot,
1155 ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1156 let (_, descriptor) = self.resolve_descriptor(slot)?;
1157 let owner = self
1158 .allocations
1159 .get(descriptor.allocation.index())
1160 .ok_or(GroupError::AllocationSlotOutOfBounds {
1161 slot: descriptor.allocation.index(),
1162 })?
1163 .as_ref()
1164 .ok_or(GroupError::AllocationSlotVacant {
1165 slot: descriptor.allocation.index(),
1166 })?;
1167 Ok(GroupReadView {
1168 owner: NonNull::from(owner),
1169 descriptor: descriptor.clone(),
1170 _borrow: PhantomData,
1171 })
1172 }
1173
1174 pub(crate) fn prepare_device_read_for_layout<T: TensorScalar, R: TensorRank>(
1175 &self,
1176 slot: DescriptorSlot,
1177 layout: &TensorLayout<R>,
1178 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1179 self.view_raw::<T, R>(slot)
1180 .map_err(|error| AccessError::InvalidLayout {
1181 message: error.to_string(),
1182 })?
1183 .prepare_device_read_for_layout(layout)
1184 }
1185
1186 pub(crate) fn allocation_index(
1187 &self,
1188 slot: DescriptorSlot,
1189 ) -> Result<AllocationSlot, GroupError> {
1190 Ok(self.resolve_descriptor(slot)?.1.allocation)
1191 }
1192
1193 pub(crate) fn set_host_recycler<T: TensorScalar>(
1194 &mut self,
1195 allocation_index: usize,
1196 recycler: std::sync::Weak<dyn super::root::HostBufferRecycler<T>>,
1197 ) -> Result<(), AccessError> {
1198 let owner = self
1199 .allocations
1200 .get_mut(allocation_index)
1201 .and_then(Option::as_mut)
1202 .ok_or(AccessError::Unsupported {
1203 backend: "missing host allocation",
1204 })?;
1205 owner.set_host_recycler(recycler)
1206 }
1207
1208 pub(crate) fn host_buffer_at<T: 'static>(
1209 &self,
1210 allocation_index: AllocationSlot,
1211 ) -> Option<&crate::StorageBuffer<T>> {
1212 self.allocations
1213 .get(allocation_index.index())?
1214 .as_ref()?
1215 .host_buffer::<T>()
1216 }
1217
1218 pub(crate) fn host_root_metadata<T: 'static>(
1219 &self,
1220 slot: DescriptorSlot,
1221 ) -> Option<(usize, usize)> {
1222 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1223 let owner = self
1224 .allocations
1225 .get(descriptor.allocation.index())?
1226 .as_ref()?;
1227 if descriptor.span() != owner.root_span() {
1228 return None;
1229 }
1230 let crate::StorageBuffer::Host(data) = owner.host_buffer::<T>()? else {
1231 return None;
1232 };
1233 let pointer = data.as_ptr() as usize;
1234 let byte_len = data.len().checked_mul(size_of::<T>())?;
1235 Some((pointer, byte_len))
1236 }
1237
1238 pub(crate) fn backend_buffer<T: 'static>(
1239 &self,
1240 slot: DescriptorSlot,
1241 ) -> Option<&crate::StorageBuffer<T>> {
1242 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1243 self.allocations
1244 .get(descriptor.allocation.index())?
1245 .as_ref()?
1246 .backend_buffer::<T>()
1247 }
1248
1249 pub(crate) fn descriptor_dtype(&self, slot: DescriptorSlot) -> Option<DType> {
1250 self.resolve_descriptor(slot)
1251 .ok()
1252 .map(|(_, descriptor)| descriptor.dtype())
1253 }
1254
1255 fn allocation_root_resource(&self, allocation: AllocationSlot) -> Option<RootResourceId> {
1261 self.allocations
1262 .get(allocation.index())?
1263 .as_ref()
1264 .map(|owner| owner.root_span().root_resource())
1265 }
1266
1267 pub(crate) fn descriptor_len(&self, slot: DescriptorSlot) -> Option<usize> {
1268 self.resolve_descriptor(slot)
1269 .ok()
1270 .map(|(_, descriptor)| descriptor.element_count)
1271 }
1272
1273 pub(crate) fn backend_identity(
1274 &self,
1275 slot: DescriptorSlot,
1276 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
1277 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1278 if descriptor.provider == ProviderKind::Cpu {
1279 return None;
1280 }
1281 let owner = self
1282 .allocations
1283 .get(descriptor.allocation.index())?
1284 .as_ref()?;
1285 let key = owner.root_identity().extent().key();
1286 Some((key.domain(), key.local()))
1287 }
1288
1289 pub(crate) fn provider_kind(&self, slot: DescriptorSlot) -> Option<ProviderKind> {
1290 self.resolve_descriptor(slot)
1291 .ok()
1292 .map(|(_, descriptor)| descriptor.provider)
1293 }
1294
1295 pub(crate) fn backend_root_buffer<T: 'static>(&self) -> Option<&crate::StorageBuffer<T>> {
1296 self.allocations
1297 .first()
1298 .and_then(|owner| owner.as_ref())
1299 .and_then(|owner| owner.backend_buffer::<T>())
1300 }
1301
1302 pub(crate) fn backend_buffer_mut<T: 'static>(
1303 &mut self,
1304 slot: DescriptorSlot,
1305 ) -> Option<&mut crate::StorageBuffer<T>> {
1306 let allocation = self.resolve_descriptor(slot).ok()?.1.allocation;
1307 self.allocations
1308 .get_mut(allocation.index())?
1309 .as_mut()?
1310 .backend_buffer_mut::<T>()
1311 }
1312
1313 pub(crate) fn backend_root_buffer_mut<T: 'static>(
1314 &mut self,
1315 ) -> Option<&mut crate::StorageBuffer<T>> {
1316 self.allocations
1317 .first_mut()
1318 .and_then(|owner| owner.as_mut())
1319 .and_then(|owner| owner.backend_buffer_mut::<T>())
1320 }
1321
1322 pub(crate) fn view_mut_raw<T: 'static, R: TensorRank>(
1323 &mut self,
1324 slot: DescriptorSlot,
1325 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1326 let (_, descriptor) = self.resolve_descriptor(slot)?;
1327 let descriptor = descriptor.clone();
1328 let owner = self
1329 .allocations
1330 .get_mut(descriptor.allocation.index())
1331 .ok_or(GroupError::AllocationSlotOutOfBounds {
1332 slot: descriptor.allocation.index(),
1333 })?
1334 .as_mut()
1335 .ok_or(GroupError::AllocationSlotVacant {
1336 slot: descriptor.allocation.index(),
1337 })?;
1338 Ok(GroupWriteView {
1339 owner: NonNull::from(owner),
1340 descriptor: descriptor.clone(),
1341 _borrow: PhantomData,
1342 })
1343 }
1344
1345 pub(crate) fn prepare_device_write_for_layout<T: TensorScalar, R: TensorRank>(
1346 &mut self,
1347 slot: DescriptorSlot,
1348 layout: &TensorLayout<R>,
1349 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1350 self.view_mut_raw::<T, R>(slot)
1351 .map_err(|error| AccessError::InvalidLayout {
1352 message: error.to_string(),
1353 })?
1354 .prepare_device_write_for_layout(layout)
1355 }
1356
1357 pub(crate) fn view_mut<T: TensorScalar, R: TensorRank>(
1358 &mut self,
1359 slot: DescriptorSlot,
1360 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1361 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1362 let mut descriptor = descriptor.clone();
1363 check_typed::<T, R>(&descriptor)?;
1364 if !descriptor.write_injective {
1365 descriptor
1366 .layout()
1367 .validate_mutable_no_overlap()
1368 .map_err(|error| GroupError::InvalidDescriptor {
1369 message: error.to_string(),
1370 })?;
1371 if let Some(Some(retained)) = self.descriptors.get_mut(descriptor_index) {
1372 retained.write_injective = true;
1373 }
1374 descriptor.write_injective = true;
1375 }
1376 let owner = self
1377 .allocations
1378 .get_mut(descriptor.allocation.index())
1379 .ok_or(GroupError::AllocationSlotOutOfBounds {
1380 slot: descriptor.allocation.index(),
1381 })?
1382 .as_mut()
1383 .ok_or(GroupError::AllocationSlotVacant {
1384 slot: descriptor.allocation.index(),
1385 })?;
1386 Ok(GroupWriteView {
1387 owner: NonNull::from(owner),
1388 descriptor,
1389 _borrow: PhantomData,
1390 })
1391 }
1392
1393 pub(crate) fn split_mut<T: TensorScalar, R: TensorRank>(
1394 &mut self,
1395 slots: &[DescriptorSlot],
1396 ) -> Result<Vec<GroupWriteView<'_, T, R>>, DisjointViewError> {
1397 let mut selected = Vec::with_capacity(slots.len());
1398 for &slot in slots {
1399 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1400 if selected
1401 .iter()
1402 .any(|(seen, _, _): &(DescriptorSlot, usize, DescriptorRecord)| *seen == slot)
1403 {
1404 return Err(DisjointViewError::DuplicateSlot { slot: slot.index() });
1405 }
1406 check_typed::<T, R>(descriptor)?;
1407 if !descriptor.write_injective {
1408 descriptor
1409 .layout()
1410 .validate_mutable_no_overlap()
1411 .map_err(|_| DisjointViewError::NonInjective { slot: slot.index() })?;
1412 }
1413 selected.push((slot, descriptor_index, descriptor.clone()));
1414 }
1415
1416 for left in 0..selected.len() {
1417 for right in (left + 1)..selected.len() {
1418 let first = &selected[left].2;
1419 let second = &selected[right].2;
1420 if first.span().root_resource() != second.span().root_resource() {
1421 continue;
1422 }
1423 match (first.envelope, second.envelope) {
1424 (None, _) | (_, None) => {}
1425 (Some(first), Some(second)) => {
1426 if first
1427 .overlaps(second)
1428 .map_err(|_| DisjointViewError::NotProvablyDisjoint)?
1429 {
1430 return Err(DisjointViewError::PairwiseOverlap);
1431 }
1432 }
1433 }
1434 }
1435 }
1436
1437 for (_, descriptor_index, _) in &selected {
1438 if let Some(Some(descriptor)) = self.descriptors.get_mut(*descriptor_index) {
1439 descriptor.write_injective = true;
1440 }
1441 }
1442
1443 let mut children = Vec::with_capacity(selected.len());
1444 for (_, _, mut descriptor) in selected {
1445 descriptor.write_injective = true;
1446 let owner = self
1447 .allocations
1448 .get_mut(descriptor.allocation.index())
1449 .ok_or(GroupError::AllocationSlotOutOfBounds {
1450 slot: descriptor.allocation.index(),
1451 })?
1452 .as_mut()
1453 .ok_or(GroupError::AllocationSlotVacant {
1454 slot: descriptor.allocation.index(),
1455 })?;
1456 children.push(GroupWriteView {
1457 owner: NonNull::from(owner),
1458 descriptor,
1459 _borrow: PhantomData,
1460 });
1461 }
1462 Ok(children)
1463 }
1464
1465 pub(crate) fn try_extract(
1466 &mut self,
1467 slot: DescriptorSlot,
1468 ) -> Result<OwnedStorage, ExtractError> {
1469 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1470 let allocation = descriptor.allocation;
1471 let references = self
1472 .descriptors
1473 .iter()
1474 .flatten()
1475 .filter(|candidate| candidate.allocation == allocation)
1476 .count();
1477 if references != 1 {
1478 return Err(ExtractError::AliasedAllocation {
1479 allocation: allocation.index(),
1480 });
1481 }
1482 let _ = self.descriptors[descriptor_index].take();
1483 self.allocations
1484 .get_mut(allocation.index())
1485 .ok_or(GroupError::AllocationSlotOutOfBounds {
1486 slot: allocation.index(),
1487 })?
1488 .take()
1489 .ok_or(GroupError::AllocationSlotVacant {
1490 slot: allocation.index(),
1491 })
1492 .map_err(ExtractError::from)
1493 }
1494
1495 #[allow(clippy::result_large_err)]
1507 pub fn take_tensor(&mut self, slot: DescriptorSlot) -> Result<crate::Tensor, GroupError> {
1508 let (_, descriptor) = self.resolve_descriptor(slot)?;
1509 let descriptor = descriptor.clone();
1510 let dtype = descriptor.dtype();
1511 let layout = descriptor.layout().clone();
1512 let placement = descriptor.placement.clone();
1513 let owner = self
1514 .try_extract(slot)
1515 .map_err(|error| GroupError::InvalidDescriptor {
1516 message: error.to_string(),
1517 })?;
1518 let mut extracted = Self::new();
1519 extracted.allocations.push(Some(owner));
1520 let mut descriptor = descriptor.clone();
1521 descriptor.allocation = AllocationSlot(0);
1522 extracted.descriptors.push(Some(descriptor));
1523 Ok(tensor_from_group(
1524 extracted,
1525 DescriptorSlot(0),
1526 AllocationSlot(0),
1527 dtype,
1528 layout,
1529 placement,
1530 ))
1531 }
1532
1533 #[allow(clippy::result_large_err)]
1536 pub(crate) fn into_owner(
1537 mut self,
1538 slot: DescriptorSlot,
1539 ) -> Result<OwnedStorage, (Self, ExtractError)> {
1540 let result = self.try_extract(slot);
1541 match result {
1542 Ok(owner) => Ok(owner),
1543 Err(error) => Err((self, error)),
1544 }
1545 }
1546
1547 #[allow(clippy::result_large_err)]
1561 pub fn into_tensor(self, slot: DescriptorSlot) -> Result<crate::Tensor, (Self, GroupError)> {
1562 let (_, descriptor) = match self.resolve_descriptor(slot) {
1563 Ok(value) => value,
1564 Err(error) => return Err((self, error)),
1565 };
1566 let allocation = descriptor.allocation;
1567 let references = self
1568 .descriptors
1569 .iter()
1570 .flatten()
1571 .filter(|candidate| candidate.allocation == allocation)
1572 .count();
1573 if references != 1 {
1574 return Err((
1575 self,
1576 GroupError::AliasedAllocation {
1577 allocation: allocation.index(),
1578 },
1579 ));
1580 }
1581 let dtype = descriptor.dtype();
1582 let layout = descriptor.layout().clone();
1583 let placement = descriptor.placement.clone();
1584 let (group, slot) = match self.into_single_descriptor(slot) {
1585 Ok(value) => value,
1586 Err((group, error)) => return Err((group, error)),
1587 };
1588 Ok(tensor_from_group(
1589 group,
1590 slot,
1591 AllocationSlot(0),
1592 dtype,
1593 layout,
1594 placement,
1595 ))
1596 }
1597
1598 #[allow(clippy::result_large_err)]
1601 fn into_single_descriptor(
1602 mut self,
1603 slot: DescriptorSlot,
1604 ) -> Result<(Self, DescriptorSlot), (Self, GroupError)> {
1605 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1606 Ok(value) => value,
1607 Err(error) => return Err((self, error)),
1608 };
1609 let allocation = descriptor.allocation;
1610 let mut descriptor = match self.descriptors[descriptor_index].take() {
1611 Some(descriptor) => descriptor,
1612 None => {
1613 return Err((
1614 self,
1615 GroupError::DescriptorSlotVacant { slot: slot.index() },
1616 ))
1617 }
1618 };
1619 let owner = match self.allocations.get_mut(allocation.index()) {
1620 Some(owner) => match owner.take() {
1621 Some(owner) => owner,
1622 None => {
1623 self.descriptors[descriptor_index] = Some(descriptor);
1624 return Err((
1625 self,
1626 GroupError::AllocationSlotVacant {
1627 slot: allocation.index(),
1628 },
1629 ));
1630 }
1631 },
1632 None => {
1633 self.descriptors[descriptor_index] = Some(descriptor);
1634 return Err((
1635 self,
1636 GroupError::AllocationSlotOutOfBounds {
1637 slot: allocation.index(),
1638 },
1639 ));
1640 }
1641 };
1642 let mut group = Self::new();
1643 descriptor.allocation = AllocationSlot(0);
1644 group.allocations.push(Some(owner));
1645 group.descriptors.push(Some(descriptor));
1646 Ok((group, DescriptorSlot(0)))
1647 }
1648
1649 pub(crate) fn into_host_vec<T: TensorScalar>(
1650 self,
1651 slot: DescriptorSlot,
1652 ) -> Result<Vec<T>, String> {
1653 let owner = self
1654 .into_owner(slot)
1655 .map_err(|(_, error)| error.to_string())?;
1656 owner
1657 .into_host_vec::<T>()
1658 .map_err(|error| error.to_string())
1659 }
1660
1661 fn resolve_descriptor(
1662 &self,
1663 slot: DescriptorSlot,
1664 ) -> Result<(usize, &DescriptorRecord), GroupError> {
1665 let index = slot.index();
1666 let descriptor = self
1667 .descriptors
1668 .get(index)
1669 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: index })?
1670 .as_ref()
1671 .ok_or(GroupError::DescriptorSlotVacant { slot: index })?;
1672 Ok((index, descriptor))
1673 }
1674
1675 #[cfg(test)]
1676 pub(crate) fn test_vacate_allocation(&mut self, slot: AllocationSlot) {
1677 if let Some(entry) = self.allocations.get_mut(slot.index()) {
1678 *entry = None;
1679 }
1680 }
1681}
1682
1683fn check_typed<T: TensorScalar, R: TensorRank>(
1684 descriptor: &DescriptorRecord,
1685) -> Result<(), GroupError> {
1686 if descriptor.dtype() != T::dtype() {
1687 return Err(GroupError::DTypeMismatch {
1688 expected: descriptor.dtype(),
1689 actual: T::dtype(),
1690 });
1691 }
1692 if let Some(expected) = R::RANK {
1693 let actual = descriptor.layout().shape().len();
1694 if expected != actual {
1695 return Err(GroupError::RankMismatch { expected, actual });
1696 }
1697 }
1698 Ok(())
1699}
1700
1701fn logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1702 shape.iter().try_fold(1usize, |count, &extent| {
1703 count
1704 .checked_mul(extent)
1705 .ok_or_else(|| GroupError::InvalidDescriptor {
1706 message: "logical element count overflows".to_owned(),
1707 })
1708 })
1709}
1710
1711fn reachable_envelope(
1712 span: &RootBoundSpan,
1713 layout: &TensorLayout<DynRank>,
1714 element_size: usize,
1715) -> Result<Option<ByteRange>, GroupError> {
1716 if layout.shape().contains(&0) {
1717 return Ok(None);
1718 }
1719 let mut minimum = layout.offset() as i128;
1720 let mut maximum = minimum;
1721 for (&extent, &stride) in layout.shape().iter().zip(layout.strides()) {
1722 let steps = i128::try_from(extent - 1).map_err(|_| GroupError::InvalidDescriptor {
1723 message: "layout extent does not fit i128".to_owned(),
1724 })?;
1725 let contribution =
1726 (stride as i128)
1727 .checked_mul(steps)
1728 .ok_or_else(|| GroupError::InvalidDescriptor {
1729 message: "reachable layout arithmetic overflows".to_owned(),
1730 })?;
1731 if contribution < 0 {
1732 minimum =
1733 minimum
1734 .checked_add(contribution)
1735 .ok_or_else(|| GroupError::InvalidDescriptor {
1736 message: "reachable layout minimum overflows".to_owned(),
1737 })?;
1738 } else {
1739 maximum =
1740 maximum
1741 .checked_add(contribution)
1742 .ok_or_else(|| GroupError::InvalidDescriptor {
1743 message: "reachable layout maximum overflows".to_owned(),
1744 })?;
1745 }
1746 }
1747 let minimum = usize::try_from(minimum).map_err(|_| GroupError::InvalidDescriptor {
1748 message: "reachable layout minimum is negative".to_owned(),
1749 })?;
1750 let maximum = usize::try_from(maximum).map_err(|_| GroupError::InvalidDescriptor {
1751 message: "reachable layout maximum is negative or too large".to_owned(),
1752 })?;
1753 let byte_offset = span
1754 .byte_offset()
1755 .checked_add(minimum.checked_mul(element_size).ok_or_else(|| {
1756 GroupError::InvalidDescriptor {
1757 message: "reachable byte offset overflows".to_owned(),
1758 }
1759 })?)
1760 .ok_or_else(|| GroupError::InvalidDescriptor {
1761 message: "reachable byte offset overflows".to_owned(),
1762 })?;
1763 let byte_len = maximum
1764 .checked_sub(minimum)
1765 .and_then(|length| length.checked_add(1))
1766 .and_then(|length| length.checked_mul(element_size))
1767 .ok_or_else(|| GroupError::InvalidDescriptor {
1768 message: "reachable byte length overflows".to_owned(),
1769 })?;
1770 let range = ByteRange::new(byte_offset, byte_len);
1771 range
1772 .checked_end()
1773 .map_err(|error| GroupError::InvalidDescriptor {
1774 message: error.to_string(),
1775 })?;
1776 Ok(Some(range))
1777}
1778
1779#[cfg(test)]
1780pub(crate) fn test_logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1781 logical_element_count(shape)
1782}
1783
1784#[cfg(test)]
1785pub(crate) fn test_reachable_envelope(
1786 span: RootBoundSpan,
1787 layout: TensorLayout<DynRank>,
1788 element_size: usize,
1789) -> Result<Option<ByteRange>, GroupError> {
1790 reachable_envelope(&span, &layout, element_size)
1791}