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 set_host_recycler<T: TensorScalar>(
1129 &mut self,
1130 allocation_index: usize,
1131 recycler: std::sync::Weak<dyn super::root::HostBufferRecycler<T>>,
1132 ) -> Result<(), AccessError> {
1133 let owner = self
1134 .allocations
1135 .get_mut(allocation_index)
1136 .and_then(Option::as_mut)
1137 .ok_or(AccessError::Unsupported {
1138 backend: "missing host allocation",
1139 })?;
1140 owner.set_host_recycler(recycler)
1141 }
1142
1143 pub(crate) fn host_buffer_at<T: 'static>(
1144 &self,
1145 allocation_index: usize,
1146 ) -> Option<&crate::StorageBuffer<T>> {
1147 self.allocations
1148 .get(allocation_index)?
1149 .as_ref()?
1150 .host_buffer::<T>()
1151 }
1152
1153 pub(crate) fn host_root_metadata<T: 'static>(
1154 &self,
1155 slot: DescriptorSlot,
1156 ) -> Option<(usize, usize)> {
1157 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1158 let owner = self
1159 .allocations
1160 .get(descriptor.allocation.index())?
1161 .as_ref()?;
1162 if descriptor.span != owner.root_span() {
1163 return None;
1164 }
1165 let crate::StorageBuffer::Host(data) = owner.host_buffer::<T>()? else {
1166 return None;
1167 };
1168 let pointer = data.as_ptr() as usize;
1169 let byte_len = data.len().checked_mul(size_of::<T>())?;
1170 Some((pointer, byte_len))
1171 }
1172
1173 pub(crate) fn backend_buffer<T: 'static>(
1174 &self,
1175 slot: DescriptorSlot,
1176 ) -> Option<&crate::StorageBuffer<T>> {
1177 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1178 self.allocations
1179 .get(descriptor.allocation.index())?
1180 .as_ref()?
1181 .backend_buffer::<T>()
1182 }
1183
1184 pub(crate) fn descriptor_len(&self, slot: DescriptorSlot) -> Option<usize> {
1185 self.resolve_descriptor(slot)
1186 .ok()
1187 .map(|(_, descriptor)| descriptor.element_count)
1188 }
1189
1190 pub(crate) fn backend_identity(
1191 &self,
1192 slot: DescriptorSlot,
1193 ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
1194 let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1195 if descriptor.provider == ProviderKind::Cpu {
1196 return None;
1197 }
1198 let owner = self
1199 .allocations
1200 .get(descriptor.allocation.index())?
1201 .as_ref()?;
1202 let key = owner.root_identity().extent().key();
1203 Some((key.domain(), key.local()))
1204 }
1205
1206 pub(crate) fn provider_kind(&self, slot: DescriptorSlot) -> Option<ProviderKind> {
1207 self.resolve_descriptor(slot)
1208 .ok()
1209 .map(|(_, descriptor)| descriptor.provider)
1210 }
1211
1212 pub(crate) fn backend_root_buffer<T: 'static>(&self) -> Option<&crate::StorageBuffer<T>> {
1213 self.allocations
1214 .first()
1215 .and_then(|owner| owner.as_ref())
1216 .and_then(|owner| owner.backend_buffer::<T>())
1217 }
1218
1219 pub(crate) fn backend_buffer_mut<T: 'static>(
1220 &mut self,
1221 slot: DescriptorSlot,
1222 ) -> Option<&mut crate::StorageBuffer<T>> {
1223 let allocation = self.resolve_descriptor(slot).ok()?.1.allocation;
1224 self.allocations
1225 .get_mut(allocation.index())?
1226 .as_mut()?
1227 .backend_buffer_mut::<T>()
1228 }
1229
1230 pub(crate) fn backend_root_buffer_mut<T: 'static>(
1231 &mut self,
1232 ) -> Option<&mut crate::StorageBuffer<T>> {
1233 self.allocations
1234 .first_mut()
1235 .and_then(|owner| owner.as_mut())
1236 .and_then(|owner| owner.backend_buffer_mut::<T>())
1237 }
1238
1239 pub(crate) fn view_mut_raw<T: 'static, R: TensorRank>(
1240 &mut self,
1241 slot: DescriptorSlot,
1242 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1243 let (_, descriptor) = self.resolve_descriptor(slot)?;
1244 let descriptor = descriptor.clone();
1245 let owner = self
1246 .allocations
1247 .get_mut(descriptor.allocation.index())
1248 .ok_or(GroupError::AllocationSlotOutOfBounds {
1249 slot: descriptor.allocation.index(),
1250 })?
1251 .as_mut()
1252 .ok_or(GroupError::AllocationSlotVacant {
1253 slot: descriptor.allocation.index(),
1254 })?;
1255 Ok(GroupWriteView {
1256 owner: NonNull::from(owner),
1257 descriptor: descriptor.clone(),
1258 _borrow: PhantomData,
1259 })
1260 }
1261
1262 pub(crate) fn prepare_device_write_for_layout<T: TensorScalar, R: TensorRank>(
1263 &mut self,
1264 slot: DescriptorSlot,
1265 layout: &TensorLayout<R>,
1266 ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1267 self.view_mut_raw::<T, R>(slot)
1268 .map_err(|error| AccessError::InvalidLayout {
1269 message: error.to_string(),
1270 })?
1271 .prepare_device_write_for_layout(layout)
1272 }
1273
1274 pub(crate) fn view_mut<T: TensorScalar, R: TensorRank>(
1275 &mut self,
1276 slot: DescriptorSlot,
1277 ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1278 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1279 let mut descriptor = descriptor.clone();
1280 check_typed::<T, R>(&descriptor)?;
1281 if !descriptor.write_injective {
1282 descriptor
1283 .layout
1284 .validate_mutable_no_overlap()
1285 .map_err(|error| GroupError::InvalidDescriptor {
1286 message: error.to_string(),
1287 })?;
1288 if let Some(Some(retained)) = self.descriptors.get_mut(descriptor_index) {
1289 retained.write_injective = true;
1290 }
1291 descriptor.write_injective = true;
1292 }
1293 let owner = self
1294 .allocations
1295 .get_mut(descriptor.allocation.index())
1296 .ok_or(GroupError::AllocationSlotOutOfBounds {
1297 slot: descriptor.allocation.index(),
1298 })?
1299 .as_mut()
1300 .ok_or(GroupError::AllocationSlotVacant {
1301 slot: descriptor.allocation.index(),
1302 })?;
1303 Ok(GroupWriteView {
1304 owner: NonNull::from(owner),
1305 descriptor,
1306 _borrow: PhantomData,
1307 })
1308 }
1309
1310 pub(crate) fn split_mut<T: TensorScalar, R: TensorRank>(
1311 &mut self,
1312 slots: &[DescriptorSlot],
1313 ) -> Result<Vec<GroupWriteView<'_, T, R>>, DisjointViewError> {
1314 let mut selected = Vec::with_capacity(slots.len());
1315 for &slot in slots {
1316 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1317 if selected
1318 .iter()
1319 .any(|(seen, _, _): &(DescriptorSlot, usize, DescriptorRecord)| *seen == slot)
1320 {
1321 return Err(DisjointViewError::DuplicateSlot { slot: slot.index() });
1322 }
1323 check_typed::<T, R>(descriptor)?;
1324 if !descriptor.write_injective {
1325 descriptor
1326 .layout
1327 .validate_mutable_no_overlap()
1328 .map_err(|_| DisjointViewError::NonInjective { slot: slot.index() })?;
1329 }
1330 selected.push((slot, descriptor_index, descriptor.clone()));
1331 }
1332
1333 for left in 0..selected.len() {
1334 for right in (left + 1)..selected.len() {
1335 let first = &selected[left].2;
1336 let second = &selected[right].2;
1337 if first.root.root_resource() != second.root.root_resource() {
1338 continue;
1339 }
1340 match (first.envelope, second.envelope) {
1341 (None, _) | (_, None) => {}
1342 (Some(first), Some(second)) => {
1343 if first
1344 .overlaps(second)
1345 .map_err(|_| DisjointViewError::NotProvablyDisjoint)?
1346 {
1347 return Err(DisjointViewError::PairwiseOverlap);
1348 }
1349 }
1350 }
1351 }
1352 }
1353
1354 for (_, descriptor_index, _) in &selected {
1355 if let Some(Some(descriptor)) = self.descriptors.get_mut(*descriptor_index) {
1356 descriptor.write_injective = true;
1357 }
1358 }
1359
1360 let mut children = Vec::with_capacity(selected.len());
1361 for (_, _, mut descriptor) in selected {
1362 descriptor.write_injective = true;
1363 let owner = self
1364 .allocations
1365 .get_mut(descriptor.allocation.index())
1366 .ok_or(GroupError::AllocationSlotOutOfBounds {
1367 slot: descriptor.allocation.index(),
1368 })?
1369 .as_mut()
1370 .ok_or(GroupError::AllocationSlotVacant {
1371 slot: descriptor.allocation.index(),
1372 })?;
1373 children.push(GroupWriteView {
1374 owner: NonNull::from(owner),
1375 descriptor,
1376 _borrow: PhantomData,
1377 });
1378 }
1379 Ok(children)
1380 }
1381
1382 pub(crate) fn try_extract(
1383 &mut self,
1384 slot: DescriptorSlot,
1385 ) -> Result<OwnedStorage, ExtractError> {
1386 let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1387 let allocation = descriptor.allocation;
1388 let references = self
1389 .descriptors
1390 .iter()
1391 .flatten()
1392 .filter(|candidate| candidate.allocation == allocation)
1393 .count();
1394 if references != 1 {
1395 return Err(ExtractError::AliasedAllocation {
1396 allocation: allocation.index(),
1397 });
1398 }
1399 let _ = self.descriptors[descriptor_index].take();
1400 self.allocations
1401 .get_mut(allocation.index())
1402 .ok_or(GroupError::AllocationSlotOutOfBounds {
1403 slot: allocation.index(),
1404 })?
1405 .take()
1406 .ok_or(GroupError::AllocationSlotVacant {
1407 slot: allocation.index(),
1408 })
1409 .map_err(ExtractError::from)
1410 }
1411
1412 #[allow(clippy::result_large_err)]
1424 pub fn take_tensor(&mut self, slot: DescriptorSlot) -> Result<crate::Tensor, GroupError> {
1425 let (_, descriptor) = self.resolve_descriptor(slot)?;
1426 let descriptor = descriptor.clone();
1427 let dtype = descriptor.dtype;
1428 let layout = descriptor.layout.clone();
1429 let placement = descriptor.placement.clone();
1430 let owner = self
1431 .try_extract(slot)
1432 .map_err(|error| GroupError::InvalidDescriptor {
1433 message: error.to_string(),
1434 })?;
1435 let mut extracted = Self::new();
1436 extracted.allocations.push(Some(owner));
1437 let mut descriptor = descriptor.clone();
1438 descriptor.allocation = AllocationSlot(0);
1439 extracted.descriptors.push(Some(descriptor));
1440 Ok(tensor_from_group(
1441 extracted,
1442 DescriptorSlot(0),
1443 0,
1444 dtype,
1445 layout,
1446 placement,
1447 ))
1448 }
1449
1450 #[allow(clippy::result_large_err)]
1453 pub(crate) fn into_owner(
1454 mut self,
1455 slot: DescriptorSlot,
1456 ) -> Result<OwnedStorage, (Self, ExtractError)> {
1457 let result = self.try_extract(slot);
1458 match result {
1459 Ok(owner) => Ok(owner),
1460 Err(error) => Err((self, error)),
1461 }
1462 }
1463
1464 #[allow(clippy::result_large_err)]
1478 pub fn into_tensor(self, slot: DescriptorSlot) -> Result<crate::Tensor, (Self, GroupError)> {
1479 let (_, descriptor) = match self.resolve_descriptor(slot) {
1480 Ok(value) => value,
1481 Err(error) => return Err((self, error)),
1482 };
1483 let allocation = descriptor.allocation;
1484 let references = self
1485 .descriptors
1486 .iter()
1487 .flatten()
1488 .filter(|candidate| candidate.allocation == allocation)
1489 .count();
1490 if references != 1 {
1491 return Err((
1492 self,
1493 GroupError::AliasedAllocation {
1494 allocation: allocation.index(),
1495 },
1496 ));
1497 }
1498 let dtype = descriptor.dtype;
1499 let layout = descriptor.layout.clone();
1500 let placement = descriptor.placement.clone();
1501 let (group, slot) = match self.into_single_descriptor(slot) {
1502 Ok(value) => value,
1503 Err((group, error)) => return Err((group, error)),
1504 };
1505 Ok(tensor_from_group(group, slot, 0, dtype, layout, placement))
1506 }
1507
1508 #[allow(clippy::result_large_err)]
1511 fn into_single_descriptor(
1512 mut self,
1513 slot: DescriptorSlot,
1514 ) -> Result<(Self, DescriptorSlot), (Self, GroupError)> {
1515 let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1516 Ok(value) => value,
1517 Err(error) => return Err((self, error)),
1518 };
1519 let allocation = descriptor.allocation;
1520 let mut descriptor = match self.descriptors[descriptor_index].take() {
1521 Some(descriptor) => descriptor,
1522 None => {
1523 return Err((
1524 self,
1525 GroupError::DescriptorSlotVacant { slot: slot.index() },
1526 ))
1527 }
1528 };
1529 let owner = match self.allocations.get_mut(allocation.index()) {
1530 Some(owner) => match owner.take() {
1531 Some(owner) => owner,
1532 None => {
1533 self.descriptors[descriptor_index] = Some(descriptor);
1534 return Err((
1535 self,
1536 GroupError::AllocationSlotVacant {
1537 slot: allocation.index(),
1538 },
1539 ));
1540 }
1541 },
1542 None => {
1543 self.descriptors[descriptor_index] = Some(descriptor);
1544 return Err((
1545 self,
1546 GroupError::AllocationSlotOutOfBounds {
1547 slot: allocation.index(),
1548 },
1549 ));
1550 }
1551 };
1552 let mut group = Self::new();
1553 descriptor.allocation = AllocationSlot(0);
1554 group.allocations.push(Some(owner));
1555 group.descriptors.push(Some(descriptor));
1556 Ok((group, DescriptorSlot(0)))
1557 }
1558
1559 pub(crate) fn into_host_vec<T: TensorScalar>(
1560 self,
1561 slot: DescriptorSlot,
1562 ) -> Result<Vec<T>, String> {
1563 let owner = self
1564 .into_owner(slot)
1565 .map_err(|(_, error)| error.to_string())?;
1566 owner
1567 .into_host_vec::<T>()
1568 .map_err(|error| error.to_string())
1569 }
1570
1571 fn resolve_descriptor(
1572 &self,
1573 slot: DescriptorSlot,
1574 ) -> Result<(usize, &DescriptorRecord), GroupError> {
1575 let index = slot.index();
1576 let descriptor = self
1577 .descriptors
1578 .get(index)
1579 .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: index })?
1580 .as_ref()
1581 .ok_or(GroupError::DescriptorSlotVacant { slot: index })?;
1582 Ok((index, descriptor))
1583 }
1584
1585 #[cfg(test)]
1586 pub(crate) fn test_vacate_allocation(&mut self, slot: AllocationSlot) {
1587 if let Some(entry) = self.allocations.get_mut(slot.index()) {
1588 *entry = None;
1589 }
1590 }
1591}
1592
1593fn check_typed<T: TensorScalar, R: TensorRank>(
1594 descriptor: &DescriptorRecord,
1595) -> Result<(), GroupError> {
1596 if descriptor.dtype != T::dtype() {
1597 return Err(GroupError::DTypeMismatch {
1598 expected: descriptor.dtype,
1599 actual: T::dtype(),
1600 });
1601 }
1602 if let Some(expected) = R::RANK {
1603 let actual = descriptor.layout.shape().len();
1604 if expected != actual {
1605 return Err(GroupError::RankMismatch { expected, actual });
1606 }
1607 }
1608 Ok(())
1609}
1610
1611fn logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1612 shape.iter().try_fold(1usize, |count, &extent| {
1613 count
1614 .checked_mul(extent)
1615 .ok_or_else(|| GroupError::InvalidDescriptor {
1616 message: "logical element count overflows".to_owned(),
1617 })
1618 })
1619}
1620
1621fn reachable_envelope(
1622 span: &RootBoundSpan,
1623 layout: &TensorLayout<DynRank>,
1624 element_size: usize,
1625) -> Result<Option<ByteRange>, GroupError> {
1626 if layout.shape().contains(&0) {
1627 return Ok(None);
1628 }
1629 let mut minimum = layout.offset() as i128;
1630 let mut maximum = minimum;
1631 for (&extent, &stride) in layout.shape().iter().zip(layout.strides()) {
1632 let steps = i128::try_from(extent - 1).map_err(|_| GroupError::InvalidDescriptor {
1633 message: "layout extent does not fit i128".to_owned(),
1634 })?;
1635 let contribution =
1636 (stride as i128)
1637 .checked_mul(steps)
1638 .ok_or_else(|| GroupError::InvalidDescriptor {
1639 message: "reachable layout arithmetic overflows".to_owned(),
1640 })?;
1641 if contribution < 0 {
1642 minimum =
1643 minimum
1644 .checked_add(contribution)
1645 .ok_or_else(|| GroupError::InvalidDescriptor {
1646 message: "reachable layout minimum overflows".to_owned(),
1647 })?;
1648 } else {
1649 maximum =
1650 maximum
1651 .checked_add(contribution)
1652 .ok_or_else(|| GroupError::InvalidDescriptor {
1653 message: "reachable layout maximum overflows".to_owned(),
1654 })?;
1655 }
1656 }
1657 let minimum = usize::try_from(minimum).map_err(|_| GroupError::InvalidDescriptor {
1658 message: "reachable layout minimum is negative".to_owned(),
1659 })?;
1660 let maximum = usize::try_from(maximum).map_err(|_| GroupError::InvalidDescriptor {
1661 message: "reachable layout maximum is negative or too large".to_owned(),
1662 })?;
1663 let byte_offset = span
1664 .byte_offset()
1665 .checked_add(minimum.checked_mul(element_size).ok_or_else(|| {
1666 GroupError::InvalidDescriptor {
1667 message: "reachable byte offset overflows".to_owned(),
1668 }
1669 })?)
1670 .ok_or_else(|| GroupError::InvalidDescriptor {
1671 message: "reachable byte offset overflows".to_owned(),
1672 })?;
1673 let byte_len = maximum
1674 .checked_sub(minimum)
1675 .and_then(|length| length.checked_add(1))
1676 .and_then(|length| length.checked_mul(element_size))
1677 .ok_or_else(|| GroupError::InvalidDescriptor {
1678 message: "reachable byte length overflows".to_owned(),
1679 })?;
1680 let range = ByteRange::new(byte_offset, byte_len);
1681 range
1682 .checked_end()
1683 .map_err(|error| GroupError::InvalidDescriptor {
1684 message: error.to_string(),
1685 })?;
1686 Ok(Some(range))
1687}
1688
1689#[cfg(test)]
1690pub(crate) fn test_logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1691 logical_element_count(shape)
1692}
1693
1694#[cfg(test)]
1695pub(crate) fn test_reachable_envelope(
1696 span: RootBoundSpan,
1697 layout: TensorLayout<DynRank>,
1698 element_size: usize,
1699) -> Result<Option<ByteRange>, GroupError> {
1700 reachable_envelope(&span, &layout, element_size)
1701}