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::prepared::{
11 prepare_read, prepare_write, validate_descriptor, AccessError, AccessTarget, CheckedDescriptor,
12 CheckedRead, CheckedWrite, PreparedRead, PreparedWrite, ProviderReadMapping,
13 ProviderWriteMapping, WriteInjectivityProof,
14};
15use super::root::{BackendAllocation, OwnedStorage, ProviderKind};
16use super::span::{ByteRange, RootBoundSpan};
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20pub(crate) struct AllocationSlot(u32);
21
22impl AllocationSlot {
23 pub(crate) const fn index(self) -> usize {
24 self.0 as usize
25 }
26
27 #[cfg(test)]
28 pub(crate) const fn test_raw(raw: u32) -> Self {
29 Self(raw)
30 }
31}
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
44pub struct DescriptorSlot(u32);
45
46impl DescriptorSlot {
47 pub const fn index(self) -> usize {
58 self.0 as usize
59 }
60
61 pub fn from_index(index: usize) -> Option<Self> {
72 match u32::try_from(index) {
73 Ok(index) => Some(Self(index)),
74 Err(_) => None,
75 }
76 }
77
78 #[cfg(test)]
79 pub(crate) const fn test_raw(raw: u32) -> Self {
80 Self(raw)
81 }
82}
83
84#[derive(Clone, Debug, PartialEq, Eq)]
86pub(crate) struct DescriptorInput<R: TensorRank> {
87 relative: ByteRange,
88 shape: R::Shape,
89 strides: R::Strides,
90 offset: isize,
91 require_injective: bool,
92}
93
94impl<R: TensorRank> DescriptorInput<R> {
95 pub(crate) fn new(
96 relative: ByteRange,
97 shape: R::Shape,
98 strides: R::Strides,
99 offset: isize,
100 require_injective: bool,
101 ) -> Self {
102 Self {
103 relative,
104 shape,
105 strides,
106 offset,
107 require_injective,
108 }
109 }
110}
111
112#[derive(Clone, Debug, PartialEq, Eq)]
114pub(crate) struct DescriptorRecord {
115 allocation: AllocationSlot,
116 root: super::identity::RootResourceIdentity,
117 span: RootBoundSpan,
118 layout: TensorLayout<DynRank>,
119 dtype: DType,
120 element_size: usize,
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.span
136 }
137
138 pub(crate) const fn dtype(&self) -> DType {
139 self.dtype
140 }
141
142 pub(crate) const fn element_size(&self) -> usize {
143 self.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.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
171#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
182pub enum GroupError {
183 #[error("group index overflows u32")]
184 IndexOverflow,
185 #[error("descriptor slot {slot} is outside the group")]
186 DescriptorSlotOutOfBounds { slot: usize },
187 #[error("descriptor slot {slot} is vacant")]
188 DescriptorSlotVacant { slot: usize },
189 #[error("allocation slot {slot} is outside the group")]
190 AllocationSlotOutOfBounds { slot: usize },
191 #[error("allocation slot {slot} is vacant")]
192 AllocationSlotVacant { slot: usize },
193 #[error("descriptor validation failed: {message}")]
194 InvalidDescriptor { message: String },
195 #[error("descriptor dtype mismatch: expected {expected:?}, actual {actual:?}")]
196 DTypeMismatch { expected: DType, actual: DType },
197 #[error("descriptor rank mismatch: expected {expected}, actual {actual}")]
198 RankMismatch { expected: usize, actual: usize },
199 #[error("allocation slot {allocation} has more than one live descriptor")]
200 AliasedAllocation { allocation: usize },
201}
202
203#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
205pub(crate) enum DisjointViewError {
206 #[error(transparent)]
207 Group(#[from] GroupError),
208 #[error("descriptor slot {slot} appears more than once")]
209 DuplicateSlot { slot: usize },
210 #[error("descriptor slot {slot} has a non-injective mutable layout")]
211 NonInjective { slot: usize },
212 #[error("requested mutable descriptor envelopes overlap")]
213 PairwiseOverlap,
214 #[error("requested mutable descriptors are not provably disjoint")]
215 NotProvablyDisjoint,
216}
217
218#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
220pub(crate) enum ExtractError {
221 #[error(transparent)]
222 Group(#[from] GroupError),
223 #[error("allocation slot {allocation} still has another descriptor")]
224 AliasedAllocation { allocation: usize },
225}
226
227#[derive(Default)]
240pub struct AllocationGroup {
241 allocations: SmallVec<[Option<OwnedStorage>; 1]>,
246 descriptors: SmallVec<[Option<DescriptorRecord>; 1]>,
247}
248
249pub(crate) struct GroupReadView<'a, T, R: TensorRank> {
251 owner: NonNull<OwnedStorage>,
252 descriptor: DescriptorRecord,
253 _borrow: PhantomData<(&'a OwnedStorage, T, R)>,
254}
255
256unsafe impl<'a, T: Send, R: TensorRank> Send for GroupReadView<'a, T, R> {}
259unsafe impl<'a, T: Sync, R: TensorRank> Sync for GroupReadView<'a, T, R> {}
260
261impl<'a, T, R: TensorRank> Clone for GroupReadView<'a, T, R> {
262 fn clone(&self) -> Self {
263 Self {
264 owner: self.owner,
265 descriptor: self.descriptor.clone(),
266 _borrow: PhantomData,
267 }
268 }
269}
270
271impl<'a, T, R: TensorRank> GroupReadView<'a, T, R> {
272 pub(crate) fn clone_dyn(&self) -> GroupReadView<'a, T, crate::DynRank> {
273 GroupReadView {
274 owner: self.owner,
275 descriptor: self.descriptor.clone(),
276 _borrow: PhantomData,
277 }
278 }
279}
280
281impl<T, R: TensorRank> std::fmt::Debug for GroupReadView<'_, T, R> {
282 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
283 formatter
284 .debug_struct("GroupReadView")
285 .field("descriptor", &self.descriptor)
286 .finish_non_exhaustive()
287 }
288}
289
290impl<'a, T: TensorScalar, R: TensorRank> GroupReadView<'a, T, R> {
291 pub(crate) fn descriptor(&self) -> &DescriptorRecord {
292 &self.descriptor
293 }
294
295 pub(crate) fn provider_kind(&self) -> ProviderKind {
296 self.descriptor.provider
297 }
298
299 pub(crate) fn backend_identity(
300 &self,
301 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
302 if self.descriptor.provider == ProviderKind::Cpu {
303 return None;
304 }
305 let key = unsafe { self.owner.as_ref().root_identity().extent().key() };
306 Some((key.domain(), key.local()))
307 }
308
309 pub(crate) fn storage_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
310 let buffer = unsafe {
313 self.owner
314 .as_ref()
315 .host_buffer::<T>()
316 .or_else(|| self.owner.as_ref().backend_buffer::<T>())
317 }?;
318 Some(unsafe {
319 std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
320 })
321 }
322
323 pub(crate) fn map_read(&self) -> Result<ProviderReadMapping<'_>, AccessError> {
324 unsafe {
327 self.owner
328 .as_ref()
329 .as_ref()
330 .map_read(self.descriptor.span, self.descriptor.dtype)
331 }
332 }
333
334 pub(crate) fn backend_allocation(&self) -> Option<&'a dyn BackendAllocation> {
335 let allocation = unsafe { self.owner.as_ref().backend_allocation() }?;
336 Some(unsafe {
337 std::mem::transmute::<&dyn BackendAllocation, &'a dyn BackendAllocation>(allocation)
338 })
339 }
340
341 pub(crate) fn prepare_device_read_for_layout(
342 &self,
343 layout: &TensorLayout<R>,
344 ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
345 let owner: crate::storage::root::StorageRef<'a> = unsafe { self.owner.as_ref().as_ref() };
346 let checked: CheckedRead<'a, R> = CheckedRead::new::<T>(
347 owner,
349 self.descriptor.span,
350 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
351 AccessError::InvalidLayout {
352 message: error.to_string(),
353 }
354 })?,
355 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
356 AccessError::InvalidLayout {
357 message: error.to_string(),
358 }
359 })?,
360 layout.offset(),
361 )?;
362 prepare_read::<T, R>(checked, AccessTarget::Device)
363 .map_err(|failure| failure.1)?
364 .into_device_state()
365 }
366
367 pub(crate) fn prepare_host_read(&self) -> Result<PreparedRead<'_, T, DynRank>, AccessError> {
368 prepare_read(
369 CheckedRead::from_validated(
370 unsafe { self.owner.as_ref().as_ref() },
372 self.descriptor.checked.clone(),
373 ),
374 AccessTarget::Host,
375 )
376 .map_err(|failure| failure.1)
377 }
378
379 pub(crate) fn host_slice(&self) -> Result<&'a [T], AccessError> {
380 unsafe {
382 self.owner
383 .as_ref()
384 .as_ref()
385 .host_slice(self.descriptor.span, self.descriptor.dtype)
386 }
387 }
388}
389
390impl<'a, T: 'static, R: TensorRank> GroupReadView<'a, T, R> {
391 pub(crate) fn backend_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
392 let buffer = unsafe { self.owner.as_ref().backend_buffer::<T>() }?;
395 Some(unsafe {
396 std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
397 })
398 }
399}
400
401pub(crate) struct GroupWriteView<'a, T, R: TensorRank> {
405 owner: NonNull<OwnedStorage>,
406 descriptor: DescriptorRecord,
407 _borrow: PhantomData<(&'a mut [u8], T, R)>,
408}
409
410unsafe impl<'a, T: Send, R: TensorRank> Send for GroupWriteView<'a, T, R> {}
413
414impl<T, R: TensorRank> std::fmt::Debug for GroupWriteView<'_, T, R> {
415 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
416 formatter
417 .debug_struct("GroupWriteView")
418 .field("descriptor", &self.descriptor)
419 .finish_non_exhaustive()
420 }
421}
422
423impl<'a, T: TensorScalar, R: TensorRank> GroupWriteView<'a, T, R> {
424 pub(crate) fn descriptor(&self) -> &DescriptorRecord {
425 &self.descriptor
426 }
427
428 pub(crate) fn map_write(&mut self) -> Result<ProviderWriteMapping<'_>, AccessError> {
429 unsafe {
433 self.owner
434 .as_mut()
435 .as_mut()
436 .map_write(self.descriptor.span, self.descriptor.dtype)
437 }
438 }
439
440 pub(crate) fn backend_buffer_mut(&mut self) -> Option<&'a mut crate::StorageBuffer<T>> {
441 let owner = unsafe { &mut *self.owner.as_ptr() };
444 let buffer = owner.backend_buffer_mut::<T>()?;
445 Some(unsafe {
446 std::mem::transmute::<&mut crate::StorageBuffer<T>, &'a mut crate::StorageBuffer<T>>(
447 buffer,
448 )
449 })
450 }
451
452 pub(crate) fn prepare_device_write_for_layout(
453 &mut self,
454 layout: &TensorLayout<R>,
455 ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
456 let owner: crate::storage::root::StorageMut<'a> =
457 unsafe { (&mut *self.owner.as_ptr()).as_mut() };
458 let checked: CheckedWrite<'a, R> = CheckedWrite::new::<T>(
459 owner,
461 self.descriptor.span,
462 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
463 AccessError::InvalidLayout {
464 message: error.to_string(),
465 }
466 })?,
467 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
468 AccessError::InvalidLayout {
469 message: error.to_string(),
470 }
471 })?,
472 layout.offset(),
473 )?;
474 prepare_write::<T, R>(checked, AccessTarget::Device)
475 .map_err(|failure| failure.1)?
476 .into_device_state()
477 }
478
479 pub(crate) fn prepare_host_write(
480 &mut self,
481 ) -> Result<PreparedWrite<'_, T, DynRank>, AccessError> {
482 let checked = CheckedWrite::from_validated(
483 unsafe { self.owner.as_mut().as_mut() },
485 self.descriptor.checked.clone(),
486 WriteInjectivityProof,
487 );
488 prepare_write(checked, AccessTarget::Host).map_err(|failure| failure.1)
489 }
490
491 pub(crate) fn host_slice_mut(&mut self) -> Result<&'a mut [T], AccessError> {
492 unsafe {
495 self.owner
496 .as_mut()
497 .as_mut()
498 .host_slice_mut(self.descriptor.span, self.descriptor.dtype)
499 }
500 }
501}
502
503impl<'a, T: 'static, R: TensorRank> GroupWriteView<'a, T, R> {
504 pub(crate) fn backend_buffer(&self) -> Option<&crate::StorageBuffer<T>> {
505 unsafe { self.owner.as_ref().backend_buffer::<T>() }
508 }
509}
510
511impl fmt::Debug for AllocationGroup {
512 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
513 formatter
514 .debug_struct("AllocationGroup")
515 .field("allocation_count", &self.allocations.len())
516 .field("descriptor_count", &self.descriptors.len())
517 .finish()
518 }
519}
520
521impl AllocationGroup {
522 pub(crate) fn new() -> Self {
523 Self::default()
524 }
525
526 pub fn from_tensors(
544 tensors: Vec<crate::Tensor>,
545 ) -> Result<(Self, Box<[DescriptorSlot]>), GroupError> {
546 let mut group = Self::new();
547 let mut bindings = Vec::with_capacity(tensors.len());
548 for tensor in tensors {
549 let (source, source_slot) = tensor.into_group_parts();
550 bindings.push(group.append_group(source, source_slot)?);
551 }
552 Ok((group, bindings.into_boxed_slice()))
553 }
554
555 pub fn read_views<'a>(
564 &'a self,
565 bindings: &[DescriptorSlot],
566 ) -> Result<Vec<TensorRead<'a>>, GroupError> {
567 bindings
568 .iter()
569 .map(|&slot| self.tensor_read(slot))
570 .collect()
571 }
572
573 pub fn read_view<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
582 self.tensor_read(slot)
583 }
584
585 fn tensor_read<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
586 let dtype = self.resolve_descriptor(slot)?.1.dtype();
587 let view = match dtype {
588 DType::F32 => tensor_view_from_group(self.view::<f32, DynRank>(slot)?),
589 DType::F64 => tensor_view_from_group(self.view::<f64, DynRank>(slot)?),
590 DType::I32 => tensor_view_from_group(self.view::<i32, DynRank>(slot)?),
591 DType::I64 => tensor_view_from_group(self.view::<i64, DynRank>(slot)?),
592 DType::Bool => tensor_view_from_group(self.view::<bool, DynRank>(slot)?),
593 DType::C32 => {
594 tensor_view_from_group(self.view::<num_complex::Complex32, DynRank>(slot)?)
595 }
596 DType::C64 => {
597 tensor_view_from_group(self.view::<num_complex::Complex64, DynRank>(slot)?)
598 }
599 }
600 .map_err(|error| GroupError::InvalidDescriptor {
601 message: error.to_string(),
602 })?;
603 Ok(TensorRead::from_view(view))
604 }
605
606 pub fn append_tensor(&mut self, tensor: crate::Tensor) -> Result<DescriptorSlot, GroupError> {
614 let (source, source_slot) = tensor.into_group_parts();
615 self.append_group(source, source_slot)
616 }
617
618 pub fn append_group(
625 &mut self,
626 mut source: AllocationGroup,
627 source_slot: DescriptorSlot,
628 ) -> Result<DescriptorSlot, GroupError> {
629 let allocation_offset =
630 u32::try_from(self.allocations.len()).map_err(|_| GroupError::IndexOverflow)?;
631 let descriptor_offset =
632 u32::try_from(self.descriptors.len()).map_err(|_| GroupError::IndexOverflow)?;
633 source.resolve_descriptor(source_slot)?;
634
635 for owner in source.allocations.drain(..) {
636 self.allocations.push(owner);
637 }
638 for descriptor in source.descriptors.drain(..) {
639 let descriptor = match descriptor {
640 Some(mut descriptor) => {
641 let allocation = descriptor
642 .allocation
643 .0
644 .checked_add(allocation_offset)
645 .ok_or(GroupError::IndexOverflow)?;
646 descriptor.allocation = AllocationSlot(allocation);
647 Some(descriptor)
648 }
649 None => None,
650 };
651 self.descriptors.push(descriptor);
652 }
653
654 let source_descriptor_index =
655 u32::try_from(source_slot.index()).map_err(|_| GroupError::IndexOverflow)?;
656 Ok(DescriptorSlot(
657 source_descriptor_index
658 .checked_add(descriptor_offset)
659 .ok_or(GroupError::IndexOverflow)?,
660 ))
661 }
662
663 pub(crate) fn set_descriptor_placement(
664 &mut self,
665 slot: DescriptorSlot,
666 placement: Placement,
667 ) -> Result<(), GroupError> {
668 let descriptor = self
669 .descriptors
670 .get_mut(slot.index())
671 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: slot.index() })?
672 .as_mut()
673 .ok_or(GroupError::DescriptorSlotVacant { slot: slot.index() })?;
674 descriptor.placement = placement;
675 Ok(())
676 }
677
678 pub(crate) fn from_host_vec<T: TensorScalar, R: TensorRank>(
679 shape: R::Shape,
680 data: Vec<T>,
681 ) -> Result<(Self, DescriptorSlot), GroupError> {
682 let owner =
683 super::root::import_host_vec(data).map_err(|error| GroupError::InvalidDescriptor {
684 message: error.to_string(),
685 })?;
686 let span = owner.root_span();
687 let mut group = Self::new();
688 let allocation = group.insert_owner(owner)?;
689 let layout =
690 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
691 message: error.to_string(),
692 })?;
693 let input = DescriptorInput::new(
694 ByteRange::new(0, span.byte_len()),
695 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
696 GroupError::InvalidDescriptor {
697 message: error.to_string(),
698 }
699 })?,
700 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
701 GroupError::InvalidDescriptor {
702 message: error.to_string(),
703 }
704 })?,
705 layout.offset(),
706 true,
707 );
708 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
709 Ok((group, slot))
710 }
711
712 #[doc(hidden)]
714 pub fn from_backend_allocation<T: TensorScalar, R: TensorRank>(
715 shape: R::Shape,
716 allocation: Box<dyn BackendAllocation>,
717 ) -> Result<(Self, DescriptorSlot), GroupError> {
718 let owner = super::root::import_unique_root(allocation).map_err(|error| {
719 GroupError::InvalidDescriptor {
720 message: error.to_string(),
721 }
722 })?;
723 let span = owner.root_span();
724 let mut group = Self::new();
725 let allocation = group.insert_owner(owner)?;
726 let layout =
727 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
728 message: error.to_string(),
729 })?;
730 let input = DescriptorInput::new(
731 ByteRange::new(0, span.byte_len()),
732 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
733 GroupError::InvalidDescriptor {
734 message: error.to_string(),
735 }
736 })?,
737 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
738 GroupError::InvalidDescriptor {
739 message: error.to_string(),
740 }
741 })?,
742 layout.offset(),
743 true,
744 );
745 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
746 Ok((group, slot))
747 }
748
749 pub(crate) fn from_backend_buffer<T: TensorScalar, R: TensorRank>(
750 shape: R::Shape,
751 buffer: crate::StorageBuffer<T>,
752 ) -> Result<(Self, DescriptorSlot), GroupError> {
753 let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
754 GroupError::InvalidDescriptor {
755 message: error.to_string(),
756 }
757 })?;
758 let span = owner.root_span();
759 let mut group = Self::new();
760 let allocation = group.insert_owner(owner)?;
761 let layout =
762 TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
763 message: error.to_string(),
764 })?;
765 let input = DescriptorInput::new(
766 ByteRange::new(0, span.byte_len()),
767 R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
768 GroupError::InvalidDescriptor {
769 message: error.to_string(),
770 }
771 })?,
772 R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
773 GroupError::InvalidDescriptor {
774 message: error.to_string(),
775 }
776 })?,
777 layout.offset(),
778 true,
779 );
780 let slot = group.insert_descriptor::<T, R>(allocation, input)?;
781 Ok((group, slot))
782 }
783
784 pub(crate) fn from_backend_root<T: Send + Sync + 'static>(
785 buffer: crate::StorageBuffer<T>,
786 ) -> Result<Self, GroupError> {
787 let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
788 GroupError::InvalidDescriptor {
789 message: error.to_string(),
790 }
791 })?;
792 let mut group = Self::new();
793 group.insert_owner(owner)?;
794 Ok(group)
795 }
796
797 pub(crate) fn insert_owner(
798 &mut self,
799 owner: OwnedStorage,
800 ) -> Result<AllocationSlot, GroupError> {
801 let index = self.allocations.len();
802 let slot = u32::try_from(index).map_err(|_| GroupError::IndexOverflow)?;
803 self.allocations.push(Some(owner));
804 Ok(AllocationSlot(slot))
805 }
806
807 pub(crate) fn insert_descriptor<T: TensorScalar, R: TensorRank>(
808 &mut self,
809 allocation: AllocationSlot,
810 input: DescriptorInput<R>,
811 ) -> Result<DescriptorSlot, GroupError> {
812 let allocation_index = allocation.index();
813 let owner = self
814 .allocations
815 .get(allocation_index)
816 .ok_or(GroupError::AllocationSlotOutOfBounds {
817 slot: allocation_index,
818 })?
819 .as_ref()
820 .ok_or(GroupError::AllocationSlotVacant {
821 slot: allocation_index,
822 })?;
823
824 let root = owner.as_ref().root_identity();
825 let span = root.bind_relative_range(input.relative).map_err(|error| {
826 GroupError::InvalidDescriptor {
827 message: error.to_string(),
828 }
829 })?;
830 let element_size = size_of::<T>();
831 if element_size == 0 || !span.byte_len().is_multiple_of(element_size) {
832 return Err(GroupError::InvalidDescriptor {
833 message: format!(
834 "byte span {} is not divisible by element size {}",
835 span.byte_len(),
836 element_size
837 ),
838 });
839 }
840 if !span
841 .guaranteed_alignment()
842 .get()
843 .is_multiple_of(align_of::<T>())
844 {
845 return Err(GroupError::InvalidDescriptor {
846 message: format!(
847 "span alignment {} is insufficient for {}-byte alignment",
848 span.guaranteed_alignment().get(),
849 align_of::<T>()
850 ),
851 });
852 }
853
854 let shape = R::shape_into_vec(input.shape);
855 let strides = R::strides_into_vec(input.strides);
856 let layout = TensorLayout::<DynRank>::from_parts(
857 shape,
858 strides,
859 input.offset,
860 span.byte_len() / element_size,
861 )
862 .map_err(|error| GroupError::InvalidDescriptor {
863 message: error.to_string(),
864 })?;
865 let element_count = logical_element_count(layout.shape())?;
866 let write_injective = if input.require_injective {
867 layout.validate_mutable_no_overlap().map_err(|error| {
868 GroupError::InvalidDescriptor {
869 message: error.to_string(),
870 }
871 })?;
872 true
873 } else {
874 false
875 };
876 let envelope = reachable_envelope(&span, &layout, element_size)?;
877 let (checked, _) = validate_descriptor::<T, DynRank>(
878 &root,
879 span,
880 layout.shape().iter().copied().collect(),
881 layout.strides().iter().copied().collect(),
882 layout.offset(),
883 false,
884 )
885 .map_err(|error| GroupError::InvalidDescriptor {
886 message: error.to_string(),
887 })?;
888 let record = DescriptorRecord {
889 allocation,
890 root,
891 span,
892 layout,
893 dtype: T::dtype(),
894 element_size,
895 element_count,
896 provider: owner.as_ref().provider_kind(),
897 placement: Placement::default(),
898 envelope,
899 write_injective,
900 checked,
901 };
902
903 let descriptor_index = self.descriptors.len();
904 let slot = u32::try_from(descriptor_index).map_err(|_| GroupError::IndexOverflow)?;
905 self.descriptors.push(Some(record));
906 Ok(DescriptorSlot(slot))
907 }
908
909 #[allow(clippy::result_large_err)]
914 pub(crate) fn update_descriptor_layout(
915 mut self,
916 slot: DescriptorSlot,
917 shape: Vec<usize>,
918 strides: Vec<isize>,
919 offset: isize,
920 ) -> Result<Self, (Self, GroupError)> {
921 let dtype = match self.resolve_descriptor(slot) {
922 Ok((_, descriptor)) => descriptor.dtype,
923 Err(error) => return Err((self, error)),
924 };
925 if let Some(Some(descriptor)) = self.descriptors.get_mut(slot.index()) {
926 descriptor.write_injective = false;
929 }
930 match dtype {
931 DType::F32 => self.reinterpret_descriptor::<f32, f32>(slot, shape, strides, offset),
932 DType::F64 => self.reinterpret_descriptor::<f64, f64>(slot, shape, strides, offset),
933 DType::I32 => self.reinterpret_descriptor::<i32, i32>(slot, shape, strides, offset),
934 DType::I64 => self.reinterpret_descriptor::<i64, i64>(slot, shape, strides, offset),
935 DType::Bool => self.reinterpret_descriptor::<bool, bool>(slot, shape, strides, offset),
936 DType::C32 => self
937 .reinterpret_descriptor::<num_complex::Complex32, num_complex::Complex32>(
938 slot, shape, strides, offset,
939 ),
940 DType::C64 => self
941 .reinterpret_descriptor::<num_complex::Complex64, num_complex::Complex64>(
942 slot, shape, strides, offset,
943 ),
944 }
945 }
946
947 #[allow(clippy::result_large_err)]
955 pub(crate) fn reinterpret_descriptor<T: TensorScalar, U: TensorScalar>(
956 mut self,
957 slot: DescriptorSlot,
958 shape: Vec<usize>,
959 strides: Vec<isize>,
960 offset: isize,
961 ) -> Result<Self, (Self, GroupError)> {
962 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
963 Ok((index, descriptor)) => (index, descriptor.clone()),
964 Err(error) => return Err((self, error)),
965 };
966 if descriptor.dtype != T::dtype() {
967 return Err((
968 self,
969 GroupError::DTypeMismatch {
970 expected: descriptor.dtype,
971 actual: T::dtype(),
972 },
973 ));
974 }
975 let allocation = descriptor.allocation;
976 let references = self
977 .descriptors
978 .iter()
979 .flatten()
980 .filter(|candidate| candidate.allocation == allocation)
981 .count();
982 if references != 1 {
983 return Err((
984 self,
985 GroupError::AliasedAllocation {
986 allocation: allocation.index(),
987 },
988 ));
989 }
990
991 let root = descriptor.root;
992 let span = descriptor.span;
993 let target_element_size = std::mem::size_of::<U>();
994 if target_element_size == 0 || !span.byte_len().is_multiple_of(target_element_size) {
995 return Err((
996 self,
997 GroupError::InvalidDescriptor {
998 message: format!(
999 "byte span {} is not divisible by target element size {}",
1000 span.byte_len(),
1001 target_element_size
1002 ),
1003 },
1004 ));
1005 }
1006 let layout = match TensorLayout::<DynRank>::from_parts(
1007 shape.into(),
1008 strides.into(),
1009 offset,
1010 span.byte_len() / target_element_size,
1011 ) {
1012 Ok(layout) => layout,
1013 Err(error) => {
1014 return Err((
1015 self,
1016 GroupError::InvalidDescriptor {
1017 message: error.to_string(),
1018 },
1019 ))
1020 }
1021 };
1022 let write_injective = descriptor.write_injective;
1023 let envelope = match reachable_envelope(&span, &layout, target_element_size) {
1024 Ok(envelope) => envelope,
1025 Err(error) => return Err((self, error)),
1026 };
1027 let (checked, _) = match validate_descriptor::<U, DynRank>(
1028 &root,
1029 span,
1030 layout.shape().iter().copied().collect(),
1031 layout.strides().iter().copied().collect(),
1032 layout.offset(),
1033 write_injective,
1034 ) {
1035 Ok(value) => value,
1036 Err(error) => {
1037 return Err((
1038 self,
1039 GroupError::InvalidDescriptor {
1040 message: error.to_string(),
1041 },
1042 ))
1043 }
1044 };
1045 let element_count = match logical_element_count(layout.shape()) {
1046 Ok(count) => count,
1047 Err(error) => return Err((self, error)),
1048 };
1049 self.descriptors[descriptor_index] = Some(DescriptorRecord {
1050 allocation,
1051 root,
1052 span,
1053 layout,
1054 dtype: U::dtype(),
1055 element_size: target_element_size,
1056 element_count,
1057 provider: descriptor.provider,
1058 placement: descriptor.placement.clone(),
1059 envelope,
1060 write_injective,
1061 checked,
1062 });
1063 Ok(self)
1064 }
1065
1066 pub(crate) fn view<T: TensorScalar, R: TensorRank>(
1067 &self,
1068 slot: DescriptorSlot,
1069 ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1070 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1071 check_typed::<T, R>(descriptor)?;
1072 let owner = self
1073 .allocations
1074 .get(descriptor.allocation.index())
1075 .ok_or(GroupError::AllocationSlotOutOfBounds {
1076 slot: descriptor.allocation.index(),
1077 })?
1078 .as_ref()
1079 .ok_or(GroupError::AllocationSlotVacant {
1080 slot: descriptor.allocation.index(),
1081 })?;
1082 let _ = descriptor_index;
1083 Ok(GroupReadView {
1084 owner: NonNull::from(owner),
1085 descriptor: descriptor.clone(),
1086 _borrow: PhantomData,
1087 })
1088 }
1089
1090 pub(crate) fn view_raw<T: 'static, R: TensorRank>(
1091 &self,
1092 slot: DescriptorSlot,
1093 ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1094 let (_, descriptor) = self.resolve_descriptor(slot)?;
1095 let owner = self
1096 .allocations
1097 .get(descriptor.allocation.index())
1098 .ok_or(GroupError::AllocationSlotOutOfBounds {
1099 slot: descriptor.allocation.index(),
1100 })?
1101 .as_ref()
1102 .ok_or(GroupError::AllocationSlotVacant {
1103 slot: descriptor.allocation.index(),
1104 })?;
1105 Ok(GroupReadView {
1106 owner: NonNull::from(owner),
1107 descriptor: descriptor.clone(),
1108 _borrow: PhantomData,
1109 })
1110 }
1111
1112 pub(crate) fn prepare_device_read_for_layout<T: TensorScalar, R: TensorRank>(
1113 &self,
1114 slot: DescriptorSlot,
1115 layout: &TensorLayout<R>,
1116 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1117 self.view_raw::<T, R>(slot)
1118 .map_err(|error| AccessError::InvalidLayout {
1119 message: error.to_string(),
1120 })?
1121 .prepare_device_read_for_layout(layout)
1122 }
1123
1124 pub(crate) fn allocation_index(&self, slot: DescriptorSlot) -> Result<usize, GroupError> {
1125 Ok(self.resolve_descriptor(slot)?.1.allocation.index())
1126 }
1127
1128 pub(crate) fn host_buffer_at<T: 'static>(
1129 &self,
1130 allocation_index: usize,
1131 ) -> Option<&crate::StorageBuffer<T>> {
1132 self.allocations
1133 .get(allocation_index)?
1134 .as_ref()?
1135 .host_buffer::<T>()
1136 }
1137
1138 pub(crate) fn host_root_metadata<T: 'static>(
1139 &self,
1140 slot: DescriptorSlot,
1141 ) -> Option<(usize, usize)> {
1142 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1143 let owner = self
1144 .allocations
1145 .get(descriptor.allocation.index())?
1146 .as_ref()?;
1147 if descriptor.span != owner.root_span() {
1148 return None;
1149 }
1150 let crate::StorageBuffer::Host(data) = owner.host_buffer::<T>()? else {
1151 return None;
1152 };
1153 let pointer = data.as_ptr() as usize;
1154 let byte_len = data.len().checked_mul(size_of::<T>())?;
1155 Some((pointer, byte_len))
1156 }
1157
1158 pub(crate) fn backend_buffer<T: 'static>(
1159 &self,
1160 slot: DescriptorSlot,
1161 ) -> Option<&crate::StorageBuffer<T>> {
1162 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1163 self.allocations
1164 .get(descriptor.allocation.index())?
1165 .as_ref()?
1166 .backend_buffer::<T>()
1167 }
1168
1169 pub(crate) fn descriptor_len(&self, slot: DescriptorSlot) -> Option<usize> {
1170 self.resolve_descriptor(slot)
1171 .ok()
1172 .map(|(_, descriptor)| descriptor.element_count)
1173 }
1174
1175 pub(crate) fn backend_identity(
1176 &self,
1177 slot: DescriptorSlot,
1178 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
1179 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1180 if descriptor.provider == ProviderKind::Cpu {
1181 return None;
1182 }
1183 let owner = self
1184 .allocations
1185 .get(descriptor.allocation.index())?
1186 .as_ref()?;
1187 let key = owner.root_identity().extent().key();
1188 Some((key.domain(), key.local()))
1189 }
1190
1191 pub(crate) fn provider_kind(&self, slot: DescriptorSlot) -> Option<ProviderKind> {
1192 self.resolve_descriptor(slot)
1193 .ok()
1194 .map(|(_, descriptor)| descriptor.provider)
1195 }
1196
1197 pub(crate) fn backend_root_buffer<T: 'static>(&self) -> Option<&crate::StorageBuffer<T>> {
1198 self.allocations
1199 .first()
1200 .and_then(|owner| owner.as_ref())
1201 .and_then(|owner| owner.backend_buffer::<T>())
1202 }
1203
1204 pub(crate) fn backend_buffer_mut<T: 'static>(
1205 &mut self,
1206 slot: DescriptorSlot,
1207 ) -> Option<&mut crate::StorageBuffer<T>> {
1208 let allocation = self.resolve_descriptor(slot).ok()?.1.allocation;
1209 self.allocations
1210 .get_mut(allocation.index())?
1211 .as_mut()?
1212 .backend_buffer_mut::<T>()
1213 }
1214
1215 pub(crate) fn backend_root_buffer_mut<T: 'static>(
1216 &mut self,
1217 ) -> Option<&mut crate::StorageBuffer<T>> {
1218 self.allocations
1219 .first_mut()
1220 .and_then(|owner| owner.as_mut())
1221 .and_then(|owner| owner.backend_buffer_mut::<T>())
1222 }
1223
1224 pub(crate) fn view_mut_raw<T: 'static, R: TensorRank>(
1225 &mut self,
1226 slot: DescriptorSlot,
1227 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1228 let (_, descriptor) = self.resolve_descriptor(slot)?;
1229 let descriptor = descriptor.clone();
1230 let owner = self
1231 .allocations
1232 .get_mut(descriptor.allocation.index())
1233 .ok_or(GroupError::AllocationSlotOutOfBounds {
1234 slot: descriptor.allocation.index(),
1235 })?
1236 .as_mut()
1237 .ok_or(GroupError::AllocationSlotVacant {
1238 slot: descriptor.allocation.index(),
1239 })?;
1240 Ok(GroupWriteView {
1241 owner: NonNull::from(owner),
1242 descriptor: descriptor.clone(),
1243 _borrow: PhantomData,
1244 })
1245 }
1246
1247 pub(crate) fn prepare_device_write_for_layout<T: TensorScalar, R: TensorRank>(
1248 &mut self,
1249 slot: DescriptorSlot,
1250 layout: &TensorLayout<R>,
1251 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1252 self.view_mut_raw::<T, R>(slot)
1253 .map_err(|error| AccessError::InvalidLayout {
1254 message: error.to_string(),
1255 })?
1256 .prepare_device_write_for_layout(layout)
1257 }
1258
1259 pub(crate) fn view_mut<T: TensorScalar, R: TensorRank>(
1260 &mut self,
1261 slot: DescriptorSlot,
1262 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1263 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1264 let mut descriptor = descriptor.clone();
1265 check_typed::<T, R>(&descriptor)?;
1266 if !descriptor.write_injective {
1267 descriptor
1268 .layout
1269 .validate_mutable_no_overlap()
1270 .map_err(|error| GroupError::InvalidDescriptor {
1271 message: error.to_string(),
1272 })?;
1273 if let Some(Some(retained)) = self.descriptors.get_mut(descriptor_index) {
1274 retained.write_injective = true;
1275 }
1276 descriptor.write_injective = true;
1277 }
1278 let owner = self
1279 .allocations
1280 .get_mut(descriptor.allocation.index())
1281 .ok_or(GroupError::AllocationSlotOutOfBounds {
1282 slot: descriptor.allocation.index(),
1283 })?
1284 .as_mut()
1285 .ok_or(GroupError::AllocationSlotVacant {
1286 slot: descriptor.allocation.index(),
1287 })?;
1288 Ok(GroupWriteView {
1289 owner: NonNull::from(owner),
1290 descriptor,
1291 _borrow: PhantomData,
1292 })
1293 }
1294
1295 pub(crate) fn split_mut<T: TensorScalar, R: TensorRank>(
1296 &mut self,
1297 slots: &[DescriptorSlot],
1298 ) -> Result<Vec<GroupWriteView<'_, T, R>>, DisjointViewError> {
1299 let mut selected = Vec::with_capacity(slots.len());
1300 for &slot in slots {
1301 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1302 if selected
1303 .iter()
1304 .any(|(seen, _, _): &(DescriptorSlot, usize, DescriptorRecord)| *seen == slot)
1305 {
1306 return Err(DisjointViewError::DuplicateSlot { slot: slot.index() });
1307 }
1308 check_typed::<T, R>(descriptor)?;
1309 if !descriptor.write_injective {
1310 descriptor
1311 .layout
1312 .validate_mutable_no_overlap()
1313 .map_err(|_| DisjointViewError::NonInjective { slot: slot.index() })?;
1314 }
1315 selected.push((slot, descriptor_index, descriptor.clone()));
1316 }
1317
1318 for left in 0..selected.len() {
1319 for right in (left + 1)..selected.len() {
1320 let first = &selected[left].2;
1321 let second = &selected[right].2;
1322 if first.root.root_resource() != second.root.root_resource() {
1323 continue;
1324 }
1325 match (first.envelope, second.envelope) {
1326 (None, _) | (_, None) => {}
1327 (Some(first), Some(second)) => {
1328 if first
1329 .overlaps(second)
1330 .map_err(|_| DisjointViewError::NotProvablyDisjoint)?
1331 {
1332 return Err(DisjointViewError::PairwiseOverlap);
1333 }
1334 }
1335 }
1336 }
1337 }
1338
1339 for (_, descriptor_index, _) in &selected {
1340 if let Some(Some(descriptor)) = self.descriptors.get_mut(*descriptor_index) {
1341 descriptor.write_injective = true;
1342 }
1343 }
1344
1345 let mut children = Vec::with_capacity(selected.len());
1346 for (_, _, mut descriptor) in selected {
1347 descriptor.write_injective = true;
1348 let owner = self
1349 .allocations
1350 .get_mut(descriptor.allocation.index())
1351 .ok_or(GroupError::AllocationSlotOutOfBounds {
1352 slot: descriptor.allocation.index(),
1353 })?
1354 .as_mut()
1355 .ok_or(GroupError::AllocationSlotVacant {
1356 slot: descriptor.allocation.index(),
1357 })?;
1358 children.push(GroupWriteView {
1359 owner: NonNull::from(owner),
1360 descriptor,
1361 _borrow: PhantomData,
1362 });
1363 }
1364 Ok(children)
1365 }
1366
1367 pub(crate) fn try_extract(
1368 &mut self,
1369 slot: DescriptorSlot,
1370 ) -> Result<OwnedStorage, ExtractError> {
1371 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1372 let allocation = descriptor.allocation;
1373 let references = self
1374 .descriptors
1375 .iter()
1376 .flatten()
1377 .filter(|candidate| candidate.allocation == allocation)
1378 .count();
1379 if references != 1 {
1380 return Err(ExtractError::AliasedAllocation {
1381 allocation: allocation.index(),
1382 });
1383 }
1384 let _ = self.descriptors[descriptor_index].take();
1385 self.allocations
1386 .get_mut(allocation.index())
1387 .ok_or(GroupError::AllocationSlotOutOfBounds {
1388 slot: allocation.index(),
1389 })?
1390 .take()
1391 .ok_or(GroupError::AllocationSlotVacant {
1392 slot: allocation.index(),
1393 })
1394 .map_err(ExtractError::from)
1395 }
1396
1397 #[allow(clippy::result_large_err)]
1409 pub fn take_tensor(&mut self, slot: DescriptorSlot) -> Result<crate::Tensor, GroupError> {
1410 let (_, descriptor) = self.resolve_descriptor(slot)?;
1411 let descriptor = descriptor.clone();
1412 let dtype = descriptor.dtype;
1413 let layout = descriptor.layout.clone();
1414 let placement = descriptor.placement.clone();
1415 let owner = self
1416 .try_extract(slot)
1417 .map_err(|error| GroupError::InvalidDescriptor {
1418 message: error.to_string(),
1419 })?;
1420 let mut extracted = Self::new();
1421 extracted.allocations.push(Some(owner));
1422 let mut descriptor = descriptor.clone();
1423 descriptor.allocation = AllocationSlot(0);
1424 extracted.descriptors.push(Some(descriptor));
1425 Ok(tensor_from_group(
1426 extracted,
1427 DescriptorSlot(0),
1428 0,
1429 dtype,
1430 layout,
1431 placement,
1432 ))
1433 }
1434
1435 #[allow(clippy::result_large_err)]
1438 pub(crate) fn into_owner(
1439 mut self,
1440 slot: DescriptorSlot,
1441 ) -> Result<OwnedStorage, (Self, ExtractError)> {
1442 let result = self.try_extract(slot);
1443 match result {
1444 Ok(owner) => Ok(owner),
1445 Err(error) => Err((self, error)),
1446 }
1447 }
1448
1449 #[allow(clippy::result_large_err)]
1463 pub fn into_tensor(self, slot: DescriptorSlot) -> Result<crate::Tensor, (Self, GroupError)> {
1464 let (_, descriptor) = match self.resolve_descriptor(slot) {
1465 Ok(value) => value,
1466 Err(error) => return Err((self, error)),
1467 };
1468 let allocation = descriptor.allocation;
1469 let references = self
1470 .descriptors
1471 .iter()
1472 .flatten()
1473 .filter(|candidate| candidate.allocation == allocation)
1474 .count();
1475 if references != 1 {
1476 return Err((
1477 self,
1478 GroupError::AliasedAllocation {
1479 allocation: allocation.index(),
1480 },
1481 ));
1482 }
1483 let dtype = descriptor.dtype;
1484 let layout = descriptor.layout.clone();
1485 let placement = descriptor.placement.clone();
1486 let (group, slot) = match self.into_single_descriptor(slot) {
1487 Ok(value) => value,
1488 Err((group, error)) => return Err((group, error)),
1489 };
1490 Ok(tensor_from_group(group, slot, 0, dtype, layout, placement))
1491 }
1492
1493 #[allow(clippy::result_large_err)]
1496 fn into_single_descriptor(
1497 mut self,
1498 slot: DescriptorSlot,
1499 ) -> Result<(Self, DescriptorSlot), (Self, GroupError)> {
1500 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1501 Ok(value) => value,
1502 Err(error) => return Err((self, error)),
1503 };
1504 let allocation = descriptor.allocation;
1505 let mut descriptor = match self.descriptors[descriptor_index].take() {
1506 Some(descriptor) => descriptor,
1507 None => {
1508 return Err((
1509 self,
1510 GroupError::DescriptorSlotVacant { slot: slot.index() },
1511 ))
1512 }
1513 };
1514 let owner = match self.allocations.get_mut(allocation.index()) {
1515 Some(owner) => match owner.take() {
1516 Some(owner) => owner,
1517 None => {
1518 self.descriptors[descriptor_index] = Some(descriptor);
1519 return Err((
1520 self,
1521 GroupError::AllocationSlotVacant {
1522 slot: allocation.index(),
1523 },
1524 ));
1525 }
1526 },
1527 None => {
1528 self.descriptors[descriptor_index] = Some(descriptor);
1529 return Err((
1530 self,
1531 GroupError::AllocationSlotOutOfBounds {
1532 slot: allocation.index(),
1533 },
1534 ));
1535 }
1536 };
1537 let mut group = Self::new();
1538 descriptor.allocation = AllocationSlot(0);
1539 group.allocations.push(Some(owner));
1540 group.descriptors.push(Some(descriptor));
1541 Ok((group, DescriptorSlot(0)))
1542 }
1543
1544 pub(crate) fn into_host_vec<T: TensorScalar>(
1545 self,
1546 slot: DescriptorSlot,
1547 ) -> Result<Vec<T>, String> {
1548 let owner = self
1549 .into_owner(slot)
1550 .map_err(|(_, error)| error.to_string())?;
1551 owner
1552 .into_host_vec::<T>()
1553 .map_err(|error| error.to_string())
1554 }
1555
1556 fn resolve_descriptor(
1557 &self,
1558 slot: DescriptorSlot,
1559 ) -> Result<(usize, &DescriptorRecord), GroupError> {
1560 let index = slot.index();
1561 let descriptor = self
1562 .descriptors
1563 .get(index)
1564 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: index })?
1565 .as_ref()
1566 .ok_or(GroupError::DescriptorSlotVacant { slot: index })?;
1567 Ok((index, descriptor))
1568 }
1569
1570 #[cfg(test)]
1571 pub(crate) fn test_vacate_allocation(&mut self, slot: AllocationSlot) {
1572 if let Some(entry) = self.allocations.get_mut(slot.index()) {
1573 *entry = None;
1574 }
1575 }
1576}
1577
1578fn check_typed<T: TensorScalar, R: TensorRank>(
1579 descriptor: &DescriptorRecord,
1580) -> Result<(), GroupError> {
1581 if descriptor.dtype != T::dtype() {
1582 return Err(GroupError::DTypeMismatch {
1583 expected: descriptor.dtype,
1584 actual: T::dtype(),
1585 });
1586 }
1587 if let Some(expected) = R::RANK {
1588 let actual = descriptor.layout.shape().len();
1589 if expected != actual {
1590 return Err(GroupError::RankMismatch { expected, actual });
1591 }
1592 }
1593 Ok(())
1594}
1595
1596fn logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1597 shape.iter().try_fold(1usize, |count, &extent| {
1598 count
1599 .checked_mul(extent)
1600 .ok_or_else(|| GroupError::InvalidDescriptor {
1601 message: "logical element count overflows".to_owned(),
1602 })
1603 })
1604}
1605
1606fn reachable_envelope(
1607 span: &RootBoundSpan,
1608 layout: &TensorLayout<DynRank>,
1609 element_size: usize,
1610) -> Result<Option<ByteRange>, GroupError> {
1611 if layout.shape().contains(&0) {
1612 return Ok(None);
1613 }
1614 let mut minimum = layout.offset() as i128;
1615 let mut maximum = minimum;
1616 for (&extent, &stride) in layout.shape().iter().zip(layout.strides()) {
1617 let steps = i128::try_from(extent - 1).map_err(|_| GroupError::InvalidDescriptor {
1618 message: "layout extent does not fit i128".to_owned(),
1619 })?;
1620 let contribution =
1621 (stride as i128)
1622 .checked_mul(steps)
1623 .ok_or_else(|| GroupError::InvalidDescriptor {
1624 message: "reachable layout arithmetic overflows".to_owned(),
1625 })?;
1626 if contribution < 0 {
1627 minimum =
1628 minimum
1629 .checked_add(contribution)
1630 .ok_or_else(|| GroupError::InvalidDescriptor {
1631 message: "reachable layout minimum overflows".to_owned(),
1632 })?;
1633 } else {
1634 maximum =
1635 maximum
1636 .checked_add(contribution)
1637 .ok_or_else(|| GroupError::InvalidDescriptor {
1638 message: "reachable layout maximum overflows".to_owned(),
1639 })?;
1640 }
1641 }
1642 let minimum = usize::try_from(minimum).map_err(|_| GroupError::InvalidDescriptor {
1643 message: "reachable layout minimum is negative".to_owned(),
1644 })?;
1645 let maximum = usize::try_from(maximum).map_err(|_| GroupError::InvalidDescriptor {
1646 message: "reachable layout maximum is negative or too large".to_owned(),
1647 })?;
1648 let byte_offset = span
1649 .byte_offset()
1650 .checked_add(minimum.checked_mul(element_size).ok_or_else(|| {
1651 GroupError::InvalidDescriptor {
1652 message: "reachable byte offset overflows".to_owned(),
1653 }
1654 })?)
1655 .ok_or_else(|| GroupError::InvalidDescriptor {
1656 message: "reachable byte offset overflows".to_owned(),
1657 })?;
1658 let byte_len = maximum
1659 .checked_sub(minimum)
1660 .and_then(|length| length.checked_add(1))
1661 .and_then(|length| length.checked_mul(element_size))
1662 .ok_or_else(|| GroupError::InvalidDescriptor {
1663 message: "reachable byte length overflows".to_owned(),
1664 })?;
1665 let range = ByteRange::new(byte_offset, byte_len);
1666 range
1667 .checked_end()
1668 .map_err(|error| GroupError::InvalidDescriptor {
1669 message: error.to_string(),
1670 })?;
1671 Ok(Some(range))
1672}
1673
1674#[cfg(test)]
1675pub(crate) fn test_logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1676 logical_element_count(shape)
1677}
1678
1679#[cfg(test)]
1680pub(crate) fn test_reachable_envelope(
1681 span: RootBoundSpan,
1682 layout: TensorLayout<DynRank>,
1683 element_size: usize,
1684) -> Result<Option<ByteRange>, GroupError> {
1685 reachable_envelope(&span, &layout, element_size)
1686}