use num_complex::{Complex, Complex32, Complex64};
use num_traits::{One, Zero};
use std::any::Any;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem::{align_of, needs_drop, offset_of, size_of};
use std::num::NonZeroUsize;
use std::ops::Deref;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::config::SliceConfig;
use crate::error::ReinterpretError;
pub use tenferro_tensor_core::{DType, DynRank, Rank, TensorLayout, TensorRank};
use tenferro_tensor_core::{ShapeVec, StrideVec};
use tenferro_tensor_core::{SliceSpec as CoreSliceSpec, ValidationError};
use tinyvec::TinyVec;
use crate::storage::{
AllocationGroup, AllocationSlot, BackendAllocation, DescriptorSlot, GroupError, GroupReadView,
GroupWriteView,
};
mod accessors;
mod col_major;
mod shape_packing;
mod strided_view;
#[cfg(test)]
mod tests;
pub use col_major::{ColMajorView, ColMajorViewMut};
pub use strided_view::StridedSliceSpec;
pub(crate) fn shape_vec(shape: &[usize]) -> ShapeVec {
shape.iter().copied().collect()
}
pub(crate) fn stride_vec(strides: &[isize]) -> StrideVec {
strides.iter().copied().collect()
}
pub(crate) fn shape_vec_zeros(len: usize) -> ShapeVec {
std::iter::repeat_n(0, len).collect()
}
fn representation_pair_error(
op: &'static str,
from: DType,
to: DType,
message: impl Into<String>,
) -> crate::Error {
crate::Error::unsupported_dtype_conversion(op, from, to, message)
}
fn validate_representation_pair(op: &'static str, from: DType, to: DType) -> crate::Result<()> {
let valid = match (from, to) {
(DType::C32, DType::F32) | (DType::F32, DType::C32) => {
size_of::<Complex32>() == 2 * size_of::<f32>()
&& align_of::<Complex32>() == align_of::<f32>()
&& offset_of!(Complex32, re) == 0
&& offset_of!(Complex32, im) == size_of::<f32>()
&& !needs_drop::<Complex32>()
&& !needs_drop::<f32>()
}
(DType::C64, DType::F64) | (DType::F64, DType::C64) => {
size_of::<Complex64>() == 2 * size_of::<f64>()
&& align_of::<Complex64>() == align_of::<f64>()
&& offset_of!(Complex64, re) == 0
&& offset_of!(Complex64, im) == size_of::<f64>()
&& !needs_drop::<Complex64>()
&& !needs_drop::<f64>()
}
_ => false,
};
if valid {
Ok(())
} else {
Err(representation_pair_error(
op,
from,
to,
"only the sealed Complex<f32><->f32 and Complex<f64><->f64 representations are supported",
))
}
}
fn reinterpret_complex_to_real_layout(
shape: &[usize],
strides: &[isize],
offset: isize,
complex_buffer_len: usize,
op: &'static str,
) -> crate::Result<TensorLayout<DynRank>> {
let real_buffer_len = complex_buffer_len
.checked_mul(2)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
let rank = shape
.len()
.checked_add(1)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
let mut real_shape = ShapeVec::with_capacity(rank);
real_shape.push(2);
real_shape.extend_from_slice(shape);
let mut real_strides = StrideVec::with_capacity(rank);
real_strides.push(1);
for &stride in strides {
real_strides.push(
stride
.checked_mul(2)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?,
);
}
let real_offset = offset
.checked_mul(2)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
TensorLayout::from_parts(real_shape, real_strides, real_offset, real_buffer_len)
.map_err(|err| tensor_layout_error(op, err))
}
fn reinterpret_real_to_complex_layout(
shape: &[usize],
strides: &[isize],
offset: isize,
real_buffer_len: usize,
op: &'static str,
) -> crate::Result<TensorLayout<DynRank>> {
if shape.first().copied() != Some(2) {
return Err(crate::Error::invalid_argument(
op,
"shape",
"the leading extent must be 2 for a complex reinterpretation",
));
}
if strides.first().copied() != Some(1) {
return Err(crate::Error::invalid_argument(
op,
"strides",
"the leading stride must be 1 for a complex reinterpretation",
));
}
if offset % 2 != 0 {
return Err(crate::Error::invalid_argument(
op,
"offset",
"the offset must be divisible by 2 for a complex reinterpretation",
));
}
let mut complex_strides = StrideVec::with_capacity(strides.len() - 1);
for &stride in &strides[1..] {
if stride % 2 != 0 {
return Err(crate::Error::invalid_argument(
op,
"strides",
"all non-leading strides must be divisible by 2",
));
}
complex_strides.push(stride / 2);
}
let complex_buffer_len = real_buffer_len / 2;
TensorLayout::from_parts(
shape[1..].iter().copied().collect(),
complex_strides,
offset / 2,
complex_buffer_len,
)
.map_err(|err| tensor_layout_error(op, err))
}
fn reinterpret_host_slice<'a, T: TensorScalar, U: TensorScalar>(
data: &'a [T],
op: &'static str,
) -> crate::Result<&'a [U]> {
let byte_len = data
.len()
.checked_mul(size_of::<T>())
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
if !byte_len.is_multiple_of(size_of::<U>()) {
return Err(crate::Error::validation(
op,
ValidationError::ViewOutOfBounds,
));
}
if data.as_ptr().align_offset(align_of::<U>()) != 0 {
return Err(crate::Error::invalid_argument(
op,
"alignment",
"the source allocation is not aligned for the target representation",
));
}
Ok(unsafe { std::slice::from_raw_parts(data.as_ptr().cast::<U>(), byte_len / size_of::<U>()) })
}
fn reinterpret_host_slice_mut<'a, T: TensorScalar, U: TensorScalar>(
data: &'a mut [T],
op: &'static str,
) -> crate::Result<&'a mut [U]> {
let byte_len = data
.len()
.checked_mul(size_of::<T>())
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
if !byte_len.is_multiple_of(size_of::<U>()) {
return Err(crate::Error::validation(
op,
ValidationError::ViewOutOfBounds,
));
}
if data.as_mut_ptr().align_offset(align_of::<U>()) != 0 {
return Err(crate::Error::invalid_argument(
op,
"alignment",
"the source allocation is not aligned for the target representation",
));
}
Ok(unsafe {
std::slice::from_raw_parts_mut(data.as_mut_ptr().cast::<U>(), byte_len / size_of::<U>())
})
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum MemoryKind {
Device,
PinnedHost,
UnpinnedHost,
Managed,
Other(String),
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum DeviceKind {
Cpu,
Gpu(GpuBackendKind),
Other(String),
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum GpuBackendKind {
Cuda,
WebGpu,
Rocm,
Other(String),
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct DeviceId {
pub kind: DeviceKind,
pub ordinal: usize,
}
#[repr(transparent)]
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct CpuDomainId(u64);
impl CpuDomainId {
pub const fn new(id: u64) -> Self {
Self(id)
}
pub const fn as_u64(self) -> u64 {
self.0
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Placement {
pub memory_kind: MemoryKind,
pub device: Option<DeviceId>,
pub cpu_affinity: Option<CpuDomainId>,
}
impl Default for Placement {
fn default() -> Self {
default_placement()
}
}
pub struct BackendStorageHandle<T> {
id: u64,
len: usize,
allocation_domain: AllocationDomainId,
_phantom: std::marker::PhantomData<T>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct AllocationDomainId(u64);
impl AllocationDomainId {
pub fn fresh() -> Self {
static NEXT_DOMAIN_ID: AtomicU64 = AtomicU64::new(1);
Self(NEXT_DOMAIN_ID.fetch_add(1, Ordering::Relaxed))
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct AllocationId(u64);
impl AllocationId {
pub const fn from_backend_id(id: u64) -> Self {
Self(id)
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[non_exhaustive]
pub enum HostAccessError {
#[error("backend `{backend}` does not support guarded host access")]
Unsupported { backend: &'static str },
#[error("allocation belongs to domain {actual:?}, expected {expected:?}")]
ForeignDomain {
expected: AllocationDomainId,
actual: AllocationDomainId,
},
#[error("the allocation already has an active host mapping")]
OverlappingHostMapping,
#[error("GPU access is in progress for the allocation")]
GpuAccessInProgress,
#[error("the allocation is mapped for host access")]
MappedForHost,
#[error("backend host mapping failed: {message}")]
BackendFailure { message: String },
#[error("host write length mismatch: expected {expected}, got {actual}")]
LengthMismatch { expected: usize, actual: usize },
}
#[doc(hidden)]
#[derive(Clone, Copy, Debug)]
pub struct DeviceAccessRequest<'a> {
allocation_domain: AllocationDomainId,
allocation_id: AllocationId,
byte_len: usize,
element_size: usize,
shape: &'a [usize],
strides: &'a [isize],
offset: isize,
}
impl<'a> DeviceAccessRequest<'a> {
pub(crate) fn new(
allocation_domain: AllocationDomainId,
allocation_id: AllocationId,
byte_len: usize,
element_size: usize,
shape: &'a [usize],
strides: &'a [isize],
offset: isize,
) -> Self {
Self {
allocation_domain,
allocation_id,
byte_len,
element_size,
shape,
strides,
offset,
}
}
pub fn allocation_domain(&self) -> AllocationDomainId {
self.allocation_domain
}
pub fn allocation_id(&self) -> AllocationId {
self.allocation_id
}
pub fn byte_len(&self) -> usize {
self.byte_len
}
pub fn element_size(&self) -> usize {
self.element_size
}
pub fn shape(&self) -> &[usize] {
self.shape
}
pub fn strides(&self) -> &[isize] {
self.strides
}
pub fn offset(&self) -> isize {
self.offset
}
}
#[doc(hidden)]
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum DeviceAccessError {
#[error("backend `{backend}` does not support prepared device access")]
Unsupported { backend: &'static str },
#[error("prepared device access request is invalid: {message}")]
InvalidRequest { message: String },
#[error("provider device preparation failed: {message}")]
ProviderFailure { message: String },
}
#[doc(hidden)]
pub trait PreparedDeviceAccess: Debug {
fn as_any(&self) -> &dyn Any;
fn into_any(self: Box<Self>) -> Box<dyn Any>;
}
trait ReadGuardAccess<T> {
fn as_slice(&self) -> &[T];
}
impl<T, G> ReadGuardAccess<T> for G
where
G: Deref,
G::Target: AsRef<[T]>,
{
fn as_slice(&self) -> &[T] {
self.deref().as_ref()
}
}
pub struct HostReadGuard<'a, T> {
access: Box<dyn ReadGuardAccess<T> + 'a>,
}
impl<T> Debug for HostReadGuard<'_, T> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("HostReadGuard")
.field("len", &self.len())
.finish_non_exhaustive()
}
}
impl<'a, T> HostReadGuard<'a, T> {
pub fn new<G>(guard: G) -> Self
where
G: Deref + 'a,
G::Target: AsRef<[T]>,
T: 'a,
{
Self {
access: Box::new(guard),
}
}
pub fn as_view(&self) -> crate::Result<TypedTensorView<'_, T, DynRank>>
where
T: 'static,
{
let data: &[T] = self;
TypedTensorView::from_slice(vec![data.len()], vec![1], 0, data)
}
}
impl<T> Deref for HostReadGuard<'_, T> {
type Target = [T];
fn deref(&self) -> &Self::Target {
self.access.as_slice()
}
}
pub trait SharedTensorAllocationDomain: Debug + Send + Sync + 'static {
fn id(&self) -> AllocationDomainId;
fn allocate(&self, dtype: DType, shape: &[usize]) -> crate::Result<Tensor>;
}
type HostWriteCopy<'a, T> = dyn FnMut(&[T]) -> Result<(), HostAccessError> + 'a;
pub struct HostWriteGuard<'a, T> {
len: usize,
copy: Box<HostWriteCopy<'a, T>>,
}
impl<T> Debug for HostWriteGuard<'_, T> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("HostWriteGuard")
.field("len", &self.len)
.finish_non_exhaustive()
}
}
impl<'a, T> HostWriteGuard<'a, T> {
pub fn new<F>(len: usize, copy: F) -> Self
where
F: FnMut(&[T]) -> Result<(), HostAccessError> + 'a,
T: 'a,
{
Self {
len,
copy: Box::new(copy),
}
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn copy_from_slice(&mut self, source: &[T]) -> Result<(), HostAccessError> {
if source.len() != self.len {
return Err(HostAccessError::LengthMismatch {
expected: self.len,
actual: source.len(),
});
}
(self.copy)(source)
}
}
impl<T> Debug for BackendStorageHandle<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BackendStorageHandle")
.field("id", &self.id)
.finish()
}
}
impl<T> BackendStorageHandle<T> {
pub fn new(id: u64) -> Self {
Self::new_with_len(id, 0)
}
pub fn new_with_len(id: u64, len: usize) -> Self {
Self {
id,
len,
allocation_domain: AllocationDomainId::fresh(),
_phantom: std::marker::PhantomData,
}
}
}
pub trait BackendStorage<T>: Debug + Send + Sync + 'static {
fn backend_family(&self) -> &'static str;
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn allocation_domain(&self) -> Option<AllocationDomainId> {
None
}
fn allocation_id(&self) -> Option<AllocationId> {
None
}
#[doc(hidden)]
fn prepare_device_access(
&self,
_request: DeviceAccessRequest<'_>,
) -> Result<Box<dyn PreparedDeviceAccess>, DeviceAccessError> {
Err(DeviceAccessError::Unsupported {
backend: self.backend_family(),
})
}
fn map_read(&self) -> Result<HostReadGuard<'_, T>, HostAccessError> {
Err(HostAccessError::Unsupported {
backend: self.backend_family(),
})
}
fn map_write(&mut self) -> Result<HostWriteGuard<'_, T>, HostAccessError> {
Err(HostAccessError::Unsupported {
backend: self.backend_family(),
})
}
fn as_any(&self) -> &dyn Any;
}
impl<T: Send + Sync + 'static> BackendStorage<T> for BackendStorageHandle<T> {
fn backend_family(&self) -> &'static str {
"opaque"
}
fn len(&self) -> usize {
self.len
}
fn allocation_domain(&self) -> Option<AllocationDomainId> {
Some(self.allocation_domain)
}
fn allocation_id(&self) -> Option<AllocationId> {
Some(AllocationId::from_backend_id(self.id))
}
fn as_any(&self) -> &dyn Any {
self
}
}
#[derive(Debug)]
pub enum StorageBuffer<T> {
Host(Vec<T>),
Backend(Box<dyn BackendStorage<T>>),
}
impl<T: 'static> StorageBuffer<T> {
pub fn len(&self) -> usize {
match self {
Self::Host(data) => data.len(),
Self::Backend(buffer) => buffer.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn is_backend(&self) -> bool {
matches!(self, Self::Backend(_))
}
}
#[derive(Debug)]
pub struct Host;
#[derive(Debug)]
pub struct Gpu;
#[derive(Debug)]
pub struct Dynamic;
mod representation_sealed {
pub trait Sealed {}
impl Sealed for super::Host {}
impl Sealed for super::Gpu {}
impl Sealed for super::Dynamic {}
}
pub trait Representation: representation_sealed::Sealed + 'static {
#[doc(hidden)]
type Storage<T, R: TensorRank>;
}
impl Representation for Host {
#[doc(hidden)]
type Storage<T, R: TensorRank> = HostStorage<T>;
}
impl Representation for Gpu {
#[doc(hidden)]
type Storage<T, R: TensorRank> = GroupStorage<R>;
}
impl Representation for Dynamic {
#[doc(hidden)]
type Storage<T, R: TensorRank> = DynamicStorage<T, R>;
}
#[doc(hidden)]
pub struct HostStorage<T> {
data: HostData<T>,
}
#[doc(hidden)]
pub struct GroupStorage<R: TensorRank> {
group: Box<OwnedTensorGroup<R>>,
}
#[doc(hidden)]
pub enum DynamicStorage<T, R: TensorRank> {
Host(HostStorage<T>),
Group(GroupStorage<R>),
}
impl<T> std::fmt::Debug for HostStorage<T> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.data.fmt(formatter)
}
}
impl<R: TensorRank> std::fmt::Debug for GroupStorage<R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("GroupStorage")
.field("group", &self.group)
.finish()
}
}
impl<T, R: TensorRank> std::fmt::Debug for DynamicStorage<T, R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Host(host) => formatter.debug_tuple("Host").field(host).finish(),
Self::Group(group) => formatter.debug_tuple("Group").field(group).finish(),
}
}
}
pub struct TypedTensor<T, R: TensorRank = DynRank, D: Representation = Dynamic> {
shape: R::Shape,
placement: Placement,
layout: StrideVec,
storage: D::Storage<T, R>,
}
fn typed_host_tensor_from_vec_col_major<T, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
op: &'static str,
) -> crate::Result<TypedTensor<T, R, Host>> {
let shape = shape
.into_rank_shape()
.map_err(|err| tensor_layout_error(op, err))?;
let layout = tenferro_tensor_core::col_major_strides(shape.as_ref())
.map_err(|err| tensor_layout_error(op, err))?;
try_checked_shape_len(shape.as_ref(), data.len(), op)?;
Ok(TypedTensor {
shape,
placement: default_placement(),
layout,
storage: HostStorage {
data: HostData::new(data),
},
})
}
fn row_major_reorder<T: Clone, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
op: &'static str,
) -> crate::Result<(R::Shape, Vec<T>)> {
let shape = shape
.into_rank_shape()
.map_err(|err| tensor_layout_error(op, err))?;
tenferro_tensor_core::col_major_strides(shape.as_ref())
.map_err(|err| tensor_layout_error(op, err))?;
try_checked_shape_len(shape.as_ref(), data.len(), op)?;
let mut row_strides = shape_vec_zeros(shape.as_ref().len());
let mut stride = 1usize;
for axis in (0..row_strides.len()).rev() {
row_strides[axis] = stride;
stride = stride
.checked_mul(shape.as_ref()[axis])
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
}
let mut coordinates = shape_vec_zeros(row_strides.len());
let mut source_offset = 0usize;
let mut reordered = Vec::with_capacity(data.len());
for _ in 0..data.len() {
reordered.push(data[source_offset].clone());
for axis in 0..coordinates.len() {
coordinates[axis] += 1;
if coordinates[axis] < shape.as_ref()[axis] {
source_offset += row_strides[axis];
break;
}
coordinates[axis] = 0;
source_offset -= row_strides[axis] * (shape.as_ref()[axis] - 1);
}
}
Ok((shape, reordered))
}
impl<T, R: TensorRank> TypedTensor<T, R, Host> {
pub fn from_host_vec_col_major(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
) -> crate::Result<Self> {
typed_host_tensor_from_vec_col_major(shape, data, "from_host_vec_col_major")
}
pub fn from_host_vec_row_major(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
) -> crate::Result<Self>
where
T: Clone,
{
let (shape, data) = row_major_reorder(shape, data, "from_host_vec_row_major")?;
typed_host_tensor_from_vec_col_major(shape, data, "from_host_vec_row_major")
}
pub fn as_slice(&self) -> &[T] {
self.storage.data.as_slice()
}
pub fn host_data(&self) -> &[T] {
self.as_slice()
}
pub fn host_data_mut(&mut self) -> &mut [T] {
self.storage.data.as_mut_slice()
}
pub fn get(&self, indices: &[usize]) -> crate::Result<&T> {
let offset = self.linear_offset(indices)?;
self.as_slice().get(offset).ok_or_else(|| {
crate::Error::validation("TypedTensor::get", ValidationError::ViewOutOfBounds)
})
}
pub fn get_mut(&mut self, indices: &[usize]) -> crate::Result<&mut T> {
let offset = self.linear_offset(indices)?;
self.host_data_mut().get_mut(offset).ok_or_else(|| {
crate::Error::validation("TypedTensor::get_mut", ValidationError::ViewOutOfBounds)
})
}
pub fn into_host_vec(self) -> Vec<T> {
self.storage.data.into_vec()
}
pub fn into_vec_col_major(self) -> (Vec<usize>, Vec<T>) {
let shape = self.shape().to_vec();
(shape, self.into_host_vec())
}
pub fn duplicate(&self) -> Self
where
T: Clone,
{
let mut copy = typed_host_tensor_from_vec_col_major::<T, R>(
self.shape.clone(),
self.as_slice().to_vec(),
"TypedTensor::duplicate",
)
.unwrap_or_else(|err| unreachable!("a validated host owner re-copies: {err}"));
copy.placement = self.placement.clone();
copy
}
pub fn into_dynamic(self) -> TypedTensor<T, R, Dynamic> {
TypedTensor {
shape: self.shape,
placement: self.placement,
layout: self.layout,
storage: DynamicStorage::Host(self.storage),
}
}
pub fn as_view(&self) -> TypedTensorView<'_, T, R, Host> {
TypedTensorView {
buffer: TensorStorageRef::Host(self.as_slice()),
root: None,
layout: self.layout(),
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
}
}
pub fn as_view_mut(&mut self) -> TypedTensorViewMut<'_, T, R, Host> {
let layout = self.layout();
let placement = self.placement.clone();
TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(self.storage.data.as_mut_slice()),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
}
}
}
impl<T: Clone, R: TensorRank> Clone for TypedTensor<T, R, Host> {
fn clone(&self) -> Self {
self.duplicate()
}
}
impl<T, R: TensorRank> std::ops::Index<&[usize]> for TypedTensor<T, R, Host> {
type Output = T;
fn index(&self, indices: &[usize]) -> &T {
match self.get(indices) {
Ok(value) => value,
Err(err) => panic!("TypedTensor host index {indices:?} is invalid: {err}"),
}
}
}
impl<T, R: TensorRank> std::ops::IndexMut<&[usize]> for TypedTensor<T, R, Host> {
fn index_mut(&mut self, indices: &[usize]) -> &mut T {
match self.get_mut(indices) {
Ok(value) => value,
Err(err) => panic!("TypedTensor host index {indices:?} is invalid: {err}"),
}
}
}
impl<T, const N: usize> std::ops::Index<[usize; N]> for TypedTensor<T, Rank<N>, Host> {
type Output = T;
fn index(&self, indices: [usize; N]) -> &T {
match self.get(&indices) {
Ok(value) => value,
Err(err) => panic!("TypedTensor host index {indices:?} is invalid: {err}"),
}
}
}
impl<T, const N: usize> std::ops::Index<&[usize; N]> for TypedTensor<T, Rank<N>, Host> {
type Output = T;
fn index(&self, indices: &[usize; N]) -> &T {
std::ops::Index::index(self, *indices)
}
}
impl<T, const N: usize> std::ops::IndexMut<[usize; N]> for TypedTensor<T, Rank<N>, Host> {
fn index_mut(&mut self, indices: [usize; N]) -> &mut T {
match self.get_mut(&indices) {
Ok(value) => value,
Err(err) => panic!("TypedTensor host index {indices:?} is invalid: {err}"),
}
}
}
impl<T, const N: usize> std::ops::IndexMut<&[usize; N]> for TypedTensor<T, Rank<N>, Host> {
fn index_mut(&mut self, indices: &[usize; N]) -> &mut T {
std::ops::IndexMut::index_mut(self, *indices)
}
}
impl<T, const N: usize> std::ops::Index<[usize; N]> for TypedTensor<T, Rank<N>> {
type Output = T;
fn index(&self, indices: [usize; N]) -> &T {
match self.get(&indices) {
Ok(value) => value,
Err(err) => panic!("TypedTensor index {indices:?} is invalid: {err}"),
}
}
}
impl<T, const N: usize> std::ops::IndexMut<[usize; N]> for TypedTensor<T, Rank<N>> {
fn index_mut(&mut self, indices: [usize; N]) -> &mut T {
match self.get_mut(&indices) {
Ok(value) => value,
Err(err) => panic!("TypedTensor index {indices:?} is invalid: {err}"),
}
}
}
fn host_read_guard<T>(data: &[T]) -> HostReadGuard<'_, T> {
HostReadGuard::new(data)
}
fn host_write_guard<T: Clone>(data: &mut [T]) -> HostWriteGuard<'_, T> {
let len = data.len();
HostWriteGuard::new(len, move |source| {
data.clone_from_slice(source);
Ok(())
})
}
fn promote_host_group<T: TensorScalar, R: TensorRank>(
shape: R::Shape,
host: HostData<T>,
) -> crate::Result<OwnedTensorGroup<R>> {
let recycler = host.recycler.clone();
let mut group = OwnedTensorGroup::from_host_vec(shape, host.into_vec())?;
if let Some(recycler) = recycler {
group
.group
.set_host_recycler(group.allocation_index.index(), recycler)
.map_err(|error| crate::Error::runtime_state_source("TypedTensor::promote", error))?;
}
Ok(group)
}
impl<T, R: TensorRank> TypedTensor<T, R, Host> {
pub fn map_read(&self) -> HostReadGuard<'_, T> {
host_read_guard(self.as_slice())
}
pub fn map_write(&mut self) -> HostWriteGuard<'_, T>
where
T: Clone,
{
host_write_guard(self.host_data_mut())
}
}
impl<T: TensorScalar, R: TensorRank> TypedTensor<T, R, Host> {
pub fn promote(self) -> crate::Result<TypedTensor<T, R, Gpu>> {
let TypedTensor {
shape,
placement,
layout,
storage,
} = self;
let group = promote_host_group(shape.clone(), storage.data)?;
Ok(TypedTensor {
shape,
placement,
layout,
storage: GroupStorage {
group: Box::new(group),
},
})
}
}
impl<T, R: TensorRank> TypedTensor<T, R, Dynamic> {
pub fn map_read(&self) -> crate::Result<HostReadGuard<'_, T>> {
match &self.storage {
DynamicStorage::Host(host) => Ok(host_read_guard(host.data.as_slice())),
DynamicStorage::Group(core) => {
let range = group_logical_range(&core.group, "TypedTensor::map_read")?;
let data = core.group.host_slice::<T>()?;
data.get(range).map(host_read_guard).ok_or_else(|| {
crate::Error::validation(
"TypedTensor::map_read",
ValidationError::ViewOutOfBounds,
)
})
}
}
}
pub fn map_write(&mut self) -> crate::Result<HostWriteGuard<'_, T>>
where
T: Clone,
{
match &mut self.storage {
DynamicStorage::Host(host) => Ok(host_write_guard(host.data.as_mut_slice())),
DynamicStorage::Group(core) => {
let range = group_logical_range(&core.group, "TypedTensor::map_write")?;
let data = core.group.host_slice_mut::<T>()?;
data.get_mut(range).map(host_write_guard).ok_or_else(|| {
crate::Error::validation(
"TypedTensor::map_write",
ValidationError::ViewOutOfBounds,
)
})
}
}
}
}
impl<T, R: TensorRank> TypedTensor<T, R, Gpu> {
pub fn from_backend_buffer_col_major(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
buffer: StorageBuffer<T>,
placement: Placement,
) -> crate::Result<Self>
where
T: TensorScalar + Send + Sync + 'static,
{
let op = "TypedTensor::from_backend_buffer_col_major";
let layout = try_compact_layout(shape, op)?;
try_checked_shape_len(layout.shape(), buffer.len(), op)?;
let group_shape = R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error(op, err))?;
match buffer {
StorageBuffer::Host(_) => Err(crate::Error::runtime_state(
op,
"a host buffer is plain host storage; use TypedTensor::<T, R, Host>::from_vec_col_major",
)),
StorageBuffer::Backend(buffer) => {
let group = OwnedTensorGroup::from_backend_buffer(
group_shape.clone(),
StorageBuffer::Backend(buffer),
placement.clone(),
)?;
Ok(TypedTensor {
shape: group_shape,
placement,
layout: stride_vec(layout.strides()),
storage: GroupStorage {
group: Box::new(group),
},
})
}
}
}
pub fn into_dynamic(self) -> TypedTensor<T, R, Dynamic> {
TypedTensor {
shape: self.shape,
placement: self.placement,
layout: self.layout,
storage: DynamicStorage::Group(self.storage),
}
}
pub fn is_backend_buffer(&self) -> bool {
self.storage.group.is_backend_buffer()
}
pub fn host_data(&self) -> crate::Result<&[T]> {
let data = self.storage.group.host_slice::<T>()?;
let range = group_logical_range(&self.storage.group, "TypedTensor::host_data")?;
data.get(range).ok_or_else(|| {
crate::Error::validation("TypedTensor::host_data", ValidationError::ViewOutOfBounds)
})
}
pub fn host_data_mut(&mut self) -> crate::Result<&mut [T]> {
let range = group_logical_range(&self.storage.group, "TypedTensor::host_data_mut")?;
let data = self.storage.group.host_slice_mut::<T>()?;
data.get_mut(range).ok_or_else(|| {
crate::Error::validation(
"TypedTensor::host_data_mut",
ValidationError::ViewOutOfBounds,
)
})
}
#[doc(hidden)]
pub fn prepare_device_read(
&self,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>>
where
T: TensorScalar + 'static,
{
let layout = group_physical_layout(&self.storage.group);
self.storage
.group
.prepare_device_read_for_layout::<T>(&layout)
.map_err(|error| crate::Error::runtime_state_source(op, error))
}
#[doc(hidden)]
pub fn prepare_device_write(
&mut self,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>>
where
T: TensorScalar + 'static,
{
let layout = group_physical_layout(&self.storage.group);
self.storage
.group
.prepare_device_write_for_layout::<T>(&layout)
.map_err(|error| crate::Error::runtime_state_source(op, error))
}
}
impl<T, R: TensorRank> TypedTensor<T, R, Dynamic> {
#[allow(clippy::result_large_err)]
pub fn into_host(self) -> Result<TypedTensor<T, R, Host>, ReinterpretError<Self>> {
let TypedTensor {
shape,
placement,
layout,
storage,
} = self;
match storage {
DynamicStorage::Host(host) => Ok(TypedTensor {
shape,
placement,
layout,
storage: host,
}),
DynamicStorage::Group(group) => Err(ReinterpretError::new(
TypedTensor {
shape,
placement,
layout,
storage: DynamicStorage::Group(group),
},
crate::Error::runtime_state(
"TypedTensor::into_host",
"the tensor is group-backed rather than plain host storage",
),
)),
}
}
#[allow(clippy::result_large_err)]
pub fn into_gpu(self) -> Result<TypedTensor<T, R, Gpu>, ReinterpretError<Self>> {
let TypedTensor {
shape,
placement,
layout,
storage,
} = self;
match storage {
DynamicStorage::Group(group) => Ok(TypedTensor {
shape,
placement,
layout,
storage: group,
}),
DynamicStorage::Host(host) => Err(ReinterpretError::new(
TypedTensor {
shape,
placement,
layout,
storage: DynamicStorage::Host(host),
},
crate::Error::runtime_state(
"TypedTensor::into_gpu",
"the tensor is a plain host owner rather than group-backed storage",
),
)),
}
}
}
impl<T, R: TensorRank> TypedTensor<T, R, Dynamic> {
pub(crate) fn physical_layout(&self) -> TensorLayout<R> {
match &self.storage {
DynamicStorage::Host(_) => TensorLayout::<R>::compact(self.shape.clone())
.unwrap_or_else(|err| unreachable!("validated owned shape: {err}")),
DynamicStorage::Group(core) => group_physical_layout(&core.group),
}
}
}
fn group_physical_layout<R: TensorRank>(core: &OwnedTensorGroup<R>) -> TensorLayout<R> {
let descriptor = core
.group
.descriptor_layout(core.slot)
.unwrap_or_else(|| unreachable!("a live group carries its descriptor layout"));
let shape = R::shape_from_vec(shape_vec(descriptor.shape()))
.unwrap_or_else(|err| unreachable!("validated descriptor shape: {err:?}"));
let strides = R::strides_from_vec(stride_vec(descriptor.strides()))
.unwrap_or_else(|err| unreachable!("validated descriptor strides: {err:?}"));
let buffer_len = core
.group
.descriptor_backing_len(core.slot)
.unwrap_or_else(|| unreachable!("a live group carries its root length"));
TensorLayout::<R>::from_parts(shape, strides, descriptor.offset(), buffer_len)
.unwrap_or_else(|err| unreachable!("validated descriptor layout: {err:?}"))
}
fn group_logical_range<R: TensorRank>(
core: &OwnedTensorGroup<R>,
op: &'static str,
) -> crate::Result<std::ops::Range<usize>> {
core.group
.descriptor_logical_range(core.slot)
.ok_or_else(|| crate::Error::validation(op, ValidationError::NonContiguousViewAsSlice))
}
impl<T, R: TensorRank, D: Representation> std::fmt::Debug for TypedTensor<T, R, D> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("TypedTensor")
.field("shape", &self.shape.as_ref())
.field("placement", &self.placement)
.finish_non_exhaustive()
}
}
impl<T, R: TensorRank, D: Representation> TypedTensor<T, R, D> {
pub fn n_elements(&self) -> usize {
match try_shape_product(self.shape(), "TypedTensor::n_elements") {
Ok(n) => n,
Err(err) => {
unreachable!("TypedTensor compact shape is validated at construction: {err}")
}
}
}
pub fn shape(&self) -> &[usize] {
self.shape.as_ref()
}
pub fn strides(&self) -> &[isize] {
&self.layout
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn layout(&self) -> TensorLayout<R> {
TensorLayout::compact(self.shape.clone())
.unwrap_or_else(|err| unreachable!("validated owned shape: {err}"))
}
pub fn placement(&self) -> &Placement {
&self.placement
}
pub fn set_placement(&mut self, placement: Placement) {
self.placement = placement;
}
pub fn set_cpu_affinity(&mut self, cpu_affinity: Option<CpuDomainId>) {
self.placement.cpu_affinity = cpu_affinity;
}
pub fn linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
try_linear_offset_for_shape(self.shape(), indices, "TypedTensor::linear_offset")
}
pub fn into_layout(self) -> TensorLayout<R> {
self.layout()
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
try_linear_offset_for_shape(self.shape(), indices, "TypedTensor::layout_linear_offset")
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
self.layout()
.is_compact_col_major()
.map_err(|err| tensor_layout_error("TypedTensor::is_col_major_contiguous", err))
}
pub fn layout_summary(&self) -> String {
let layout = self.layout();
layout_summary(self.shape(), layout.strides(), layout.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
let layout = self.layout();
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
layout.strides(),
layout.offset(),
"TypedTensor::assert_col_major_contiguous",
)
}
}
struct HostData<T> {
buffer: StorageBuffer<T>,
recycler: Option<std::sync::Weak<dyn crate::HostBufferRecycler<T>>>,
}
impl<T> std::fmt::Debug for HostData<T> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("HostData")
.field("pooled", &self.recycler.is_some())
.finish_non_exhaustive()
}
}
impl<T> HostData<T> {
fn new(data: Vec<T>) -> Self {
Self {
buffer: StorageBuffer::Host(data),
recycler: None,
}
}
fn into_vec(mut self) -> Vec<T> {
self.recycler = None;
match std::mem::replace(&mut self.buffer, StorageBuffer::Host(Vec::new())) {
StorageBuffer::Host(data) => data,
StorageBuffer::Backend(_) => unreachable!("plain host storage is host-owned"),
}
}
fn as_slice(&self) -> &[T] {
match &self.buffer {
StorageBuffer::Host(data) => data,
StorageBuffer::Backend(_) => unreachable!("plain host storage is host-owned"),
}
}
fn as_mut_slice(&mut self) -> &mut [T] {
match &mut self.buffer {
StorageBuffer::Host(data) => data,
StorageBuffer::Backend(_) => unreachable!("plain host storage is host-owned"),
}
}
}
impl<T> Drop for HostData<T> {
fn drop(&mut self) {
let Some(recycler) = self.recycler.take().and_then(|weak| weak.upgrade()) else {
return;
};
if let StorageBuffer::Host(data) = &mut self.buffer {
recycler.recycle(std::mem::take(data));
}
}
}
pub(crate) struct OwnedTensorGroup<R: TensorRank> {
group: AllocationGroup,
slot: DescriptorSlot,
allocation_index: AllocationSlot,
host_ptr: Option<NonZeroUsize>,
host_byte_len: usize,
_rank: PhantomData<R>,
}
impl<R: TensorRank> Debug for OwnedTensorGroup<R> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("OwnedTensorGroup")
.field("slot", &self.slot)
.finish_non_exhaustive()
}
}
impl<R: TensorRank> OwnedTensorGroup<R> {
fn from_host_vec<T: TensorScalar>(shape: R::Shape, data: Vec<T>) -> crate::Result<Self> {
let (group, slot) = AllocationGroup::from_host_vec::<T, R>(shape, data)
.map_err(|error| group_error("TypedTensor::from_host_vec", error))?;
let allocation_index = group
.allocation_index(slot)
.map_err(|error| group_error("TypedTensor::from_host_vec", error))?;
let (host_ptr, host_byte_len) = host_metadata::<T>(&group, slot);
Ok(Self {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
})
}
fn from_backend_buffer<T: TensorScalar + Send + Sync + 'static>(
shape: R::Shape,
buffer: StorageBuffer<T>,
placement: Placement,
) -> crate::Result<Self> {
let (mut group, slot) = AllocationGroup::from_backend_buffer::<T, R>(shape, buffer)
.map_err(|error| group_error("TypedTensor::from_backend_buffer", error))?;
group
.set_descriptor_placement(slot, placement)
.map_err(|error| group_error("TypedTensor::from_backend_buffer", error))?;
let allocation_index = group
.allocation_index(slot)
.map_err(|error| group_error("TypedTensor::from_backend_buffer", error))?;
let (host_ptr, host_byte_len) = host_metadata::<T>(&group, slot);
Ok(Self {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
})
}
fn view<T: TensorScalar>(&self) -> crate::Result<GroupReadView<'_, T, R>> {
self.group
.view(self.slot)
.map_err(|error| group_error("TypedTensor::group_view", error))
}
fn view_dyn<T: TensorScalar>(&self) -> crate::Result<GroupReadView<'_, T, DynRank>> {
self.group
.view(self.slot)
.map_err(|error| group_error("TypedTensor::group_view", error))
}
fn view_mut<T: TensorScalar>(&mut self) -> crate::Result<GroupWriteView<'_, T, R>> {
self.group
.view_mut(self.slot)
.map_err(|error| group_error("TypedTensor::group_view_mut", error))
}
fn view_mut_dyn<T: TensorScalar>(&mut self) -> crate::Result<GroupWriteView<'_, T, DynRank>> {
self.group
.view_mut(self.slot)
.map_err(|error| group_error("TypedTensor::group_view_mut", error))
}
fn prepare_device_read_for_layout<T: TensorScalar>(
&self,
layout: &TensorLayout<R>,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>> {
self.group
.prepare_device_read_for_layout::<T, R>(self.slot, layout)
.map_err(|error| {
crate::Error::runtime_state_source("TypedTensor::prepare_device_read", error)
})
}
fn prepare_device_write_for_layout<T: TensorScalar>(
&mut self,
layout: &TensorLayout<R>,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>> {
self.group
.prepare_device_write_for_layout::<T, R>(self.slot, layout)
.map_err(|error| {
crate::Error::runtime_state_source("TypedTensor::prepare_device_write", error)
})
}
fn host_buffer<T: 'static>(&self) -> Option<&StorageBuffer<T>> {
self.group.host_buffer_at::<T>(self.allocation_index)
}
fn host_slice<T>(&self) -> crate::Result<&[T]> {
let Some(pointer) = self.host_ptr else {
return Err(crate::Error::runtime_state(
"TypedTensor::host_data",
"backend storage cannot be borrowed as host data; download explicitly first",
));
};
let element_size = size_of::<T>();
let Some(element_count) = self.host_byte_len.checked_div(element_size) else {
return Err(crate::Error::runtime_state(
"TypedTensor::host_data",
"host allocation byte length is not aligned to the requested dtype",
));
};
Ok(unsafe { std::slice::from_raw_parts(pointer.get() as *const T, element_count) })
}
fn host_slice_mut<T>(&mut self) -> crate::Result<&mut [T]> {
let Some(pointer) = self.host_ptr else {
return Err(crate::Error::runtime_state(
"TypedTensor::host_data_mut",
"backend storage cannot be borrowed as host data; download explicitly first",
));
};
let element_size = size_of::<T>();
let Some(element_count) = self.host_byte_len.checked_div(element_size) else {
return Err(crate::Error::runtime_state(
"TypedTensor::host_data_mut",
"host allocation byte length is not aligned to the requested dtype",
));
};
Ok(unsafe { std::slice::from_raw_parts_mut(pointer.get() as *mut T, element_count) })
}
fn backend_buffer<T: 'static>(&self) -> Option<&StorageBuffer<T>> {
self.group.backend_buffer::<T>(self.slot)
}
fn backend_buffer_mut<T: 'static>(&mut self) -> Option<&mut StorageBuffer<T>> {
self.group.backend_buffer_mut::<T>(self.slot)
}
#[allow(clippy::result_large_err)]
fn into_host_vec<T: 'static>(self) -> std::result::Result<Vec<T>, (Self, crate::Error)> {
if self.backend_buffer::<T>().is_some() {
return Err((
self,
crate::Error::runtime_state(
"TypedTensor::into_host_vec",
"backend buffers cannot be exported as host Vec; download the tensor first",
),
));
}
let OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
..
} = self;
match group.into_host_vec::<T>(slot) {
Ok(data) => Ok(data),
Err((group, error)) => Err((
OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
},
crate::Error::runtime_state("TypedTensor::into_host_vec", error),
)),
}
}
fn into_parts(self) -> (AllocationGroup, DescriptorSlot) {
(self.group, self.slot)
}
#[allow(clippy::result_large_err)]
fn reinterpret<T: TensorScalar, U: TensorScalar>(
self,
shape: Vec<usize>,
strides: Vec<isize>,
offset: isize,
) -> Result<OwnedTensorGroup<DynRank>, (Self, crate::Error)> {
let OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: _,
} = self;
match group.reinterpret_descriptor::<T, U>(slot, shape, strides, offset) {
Ok(group) => Ok(OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
}),
Err((group, error)) => Err((
OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
},
group_error("TypedTensor::reinterpret", error),
)),
}
}
}
fn host_metadata<T: TensorScalar>(
group: &AllocationGroup,
slot: DescriptorSlot,
) -> (Option<NonZeroUsize>, usize) {
group
.host_root_metadata::<T>(slot)
.map_or((None, 0), |(pointer, byte_len)| {
(NonZeroUsize::new(pointer), byte_len)
})
}
fn group_error(op: &'static str, error: GroupError) -> crate::Error {
crate::Error::runtime_state(op, error.to_string())
}
#[derive(Debug)]
pub enum TensorStorageRef<'a, T> {
Host(&'a [T]),
Backend(&'a dyn BackendStorage<T>),
#[doc(hidden)]
Root(&'a dyn BackendAllocation),
}
impl<T> Clone for TensorStorageRef<'_, T> {
fn clone(&self) -> Self {
match self {
Self::Host(data) => Self::Host(data),
Self::Backend(buffer) => Self::Backend(*buffer),
Self::Root(allocation) => Self::Root(*allocation),
}
}
}
impl<T: 'static> TensorStorageRef<'_, T> {
pub fn len(&self) -> usize {
match self {
Self::Host(data) => data.len(),
Self::Backend(buffer) => buffer.len(),
Self::Root(allocation) => allocation
.root_extent()
.byte_len()
.checked_div(std::mem::size_of::<T>())
.unwrap_or(0),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[derive(Debug)]
pub enum TensorStorageRefMut<'a, T> {
Host(&'a mut [T]),
Backend(&'a mut dyn BackendStorage<T>),
}
impl<T: 'static> TensorStorageRefMut<'_, T> {
pub fn len(&self) -> usize {
match self {
Self::Host(data) => data.len(),
Self::Backend(buffer) => buffer.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
pub struct TypedTensorView<'a, T, R: TensorRank = DynRank, D: Representation = Dynamic> {
buffer: TensorStorageRef<'a, T>,
root: Option<GroupReadView<'a, T, R>>,
layout: TensorLayout<R>,
placement: Placement,
_representation: std::marker::PhantomData<D>,
}
impl<'a, T, R: TensorRank, D: Representation> Clone for TypedTensorView<'a, T, R, D> {
fn clone(&self) -> Self {
Self {
buffer: self.buffer.clone(),
root: self.root.clone(),
layout: self.layout.clone(),
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
}
}
}
impl<'a, T, R: TensorRank, D: Representation> std::fmt::Debug for TypedTensorView<'a, T, R, D> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("TypedTensorView")
.field("shape", &self.layout.shape())
.field("placement", &self.placement)
.finish_non_exhaustive()
}
}
impl<'a, T: 'static> TypedTensorView<'a, T, DynRank> {
pub fn from_col_major(shape: &[usize], data: &'a [T]) -> crate::Result<Self> {
let layout = TensorLayout::<DynRank>::compact(shape_vec(shape))
.map_err(|err| tensor_layout_error("TypedTensorView::from_col_major", err))?;
Self::from_buffer_ref(
shape_vec(layout.shape()),
stride_vec(layout.strides()),
layout.offset(),
TensorStorageRef::Host(data),
default_placement(),
"TypedTensorView::from_col_major",
)
}
pub fn from_slice(
shape: impl AsRef<[usize]>,
strides: impl AsRef<[isize]>,
offset: isize,
data: &'a [T],
) -> crate::Result<Self> {
Self::from_buffer_ref(
shape_vec(shape.as_ref()),
stride_vec(strides.as_ref()),
offset,
TensorStorageRef::Host(data),
default_placement(),
"TypedTensorView::from_slice",
)
}
}
impl<'a, T: 'static, R: TensorRank, D: Representation> TypedTensorView<'a, T, R, D> {
pub fn from_slice_ranked(
shape: impl Into<R::Shape>,
strides: impl Into<R::Strides>,
offset: isize,
data: &'a [T],
) -> crate::Result<Self> {
Self::from_buffer_ref(
shape,
strides,
offset,
TensorStorageRef::Host(data),
default_placement(),
"TypedTensorView::from_slice_ranked",
)
}
}
impl<'a, T: 'static, R: TensorRank, D: Representation> TypedTensorView<'a, T, R, D> {
fn from_buffer_ref(
shape: impl Into<R::Shape>,
strides: impl Into<R::Strides>,
offset: isize,
buffer: TensorStorageRef<'a, T>,
placement: Placement,
op: &'static str,
) -> crate::Result<Self> {
let layout = TensorLayout::from_parts(shape.into(), strides.into(), offset, buffer.len())
.map_err(|err| tensor_layout_error(op, err))?;
Ok(Self {
buffer,
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
})
}
}
impl<'a, T: 'static, R: TensorRank> TypedTensorView<'a, T, R, Host> {
pub fn from_host_slice(
shape: impl Into<R::Shape>,
strides: impl Into<R::Strides>,
offset: isize,
data: &'a [T],
) -> crate::Result<Self> {
let buffer = TensorStorageRef::Host(data);
let layout = TensorLayout::from_parts(shape.into(), strides.into(), offset, buffer.len())
.map_err(|err| tensor_layout_error("TypedTensorView::from_host_slice", err))?;
Ok(Self {
buffer,
root: None,
layout,
placement: default_placement(),
_representation: std::marker::PhantomData,
})
}
pub fn as_host_slice(&self) -> &'a [T] {
match &self.buffer {
TensorStorageRef::Host(data) => data,
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
unreachable!("a Host-marked view always borrows host storage")
}
}
}
}
impl<'a, T: 'static, R: TensorRank, D: Representation> TypedTensorView<'a, T, R, D> {
pub fn into_tensor_read(self) -> crate::Result<TensorRead<'a>>
where
T: TensorScalar,
{
let layout = TensorLayout::<DynRank>::from_parts(
shape_vec(self.layout.shape()),
stride_vec(self.layout.strides()),
self.layout.offset(),
self.buffer.len(),
)
.map_err(|err| tensor_layout_error("TypedTensorView::into_tensor_read", err))?;
let view = TypedTensorView {
buffer: self.buffer,
root: self.root.map(GroupReadView::into_dyn),
layout,
placement: self.placement,
_representation: std::marker::PhantomData,
};
Ok(TensorRead::from_view(T::tensor_view(view)))
}
pub fn shape(&self) -> &[usize] {
self.layout.shape()
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn strides(&self) -> &[isize] {
self.layout.strides()
}
pub fn offset(&self) -> isize {
self.layout.offset()
}
pub fn host_storage(&self) -> crate::Result<&'a [T]> {
match &self.buffer {
TensorStorageRef::Host(data) => Ok(data),
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
Err(crate::Error::runtime_state(
"TypedTensorView::host_storage",
"backend buffers cannot expose host storage; download explicitly first",
))
}
}
}
pub fn n_elements(&self) -> usize {
match checked_view_element_count(self.shape(), "TypedTensorView::n_elements") {
Ok(n) => n,
Err(err) => {
unreachable!("TypedTensorView layout shape is validated at construction: {err}")
}
}
}
pub fn layout(&self) -> &TensorLayout<R> {
&self.layout
}
pub fn placement(&self) -> &Placement {
&self.placement
}
#[doc(hidden)]
pub fn backing_len(&self) -> usize {
self.buffer.len()
}
#[doc(hidden)]
pub fn backend_buffer(&self) -> Option<&dyn BackendStorage<T>> {
match &self.buffer {
TensorStorageRef::Host(_) => None,
TensorStorageRef::Backend(buffer) => Some(*buffer),
TensorStorageRef::Root(_) => {
self.root
.as_ref()?
.backend_buffer()
.and_then(|buffer| match buffer {
StorageBuffer::Host(_) => None,
StorageBuffer::Backend(buffer) => Some(buffer.as_ref()),
})
}
}
}
#[doc(hidden)]
pub fn backend_family(&self) -> Option<&'static str>
where
T: TensorScalar + 'static,
{
self.root
.as_ref()
.and_then(|root| {
root.backend_allocation()
.map(|_| root.provider_kind().as_str())
})
.or_else(|| self.backend_buffer().map(|buffer| buffer.backend_family()))
}
#[doc(hidden)]
pub fn allocation_domain(&self) -> Option<AllocationDomainId>
where
T: TensorScalar + 'static,
{
self.root
.as_ref()
.and_then(|root| root.backend_identity().map(|(domain, _)| domain))
.or_else(|| {
self.backend_buffer()
.and_then(|buffer| buffer.allocation_domain())
})
}
#[doc(hidden)]
pub fn allocation_id(&self) -> Option<AllocationId>
where
T: TensorScalar + 'static,
{
self.root
.as_ref()
.and_then(|root| root.backend_identity().map(|(_, id)| id))
.or_else(|| {
self.backend_buffer()
.and_then(|buffer| buffer.allocation_id())
})
}
#[doc(hidden)]
pub fn prepare_device_read(
&self,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>>
where
T: TensorScalar + 'static,
{
if let Some(root) = &self.root {
return root
.prepare_device_read_for_layout(&self.layout)
.map_err(|error| crate::Error::runtime_state_source(op, error));
}
let buffer = self.backend_buffer().ok_or_else(|| {
crate::Error::runtime_state_source(
op,
crate::AccessError::Unsupported { backend: "host" },
)
})?;
prepare_backend_access(buffer, &self.layout, op)
}
pub fn linear_offset(&self, indices: &[usize]) -> Option<usize> {
checked_view_offset(self.shape(), self.strides(), self.offset(), indices)
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
checked_view_offset_result(
self.shape(),
self.strides(),
self.offset(),
indices,
"TypedTensorView::layout_linear_offset",
)
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
self.layout
.is_compact_col_major()
.map_err(|err| tensor_layout_error("TypedTensorView::is_col_major_contiguous", err))
}
pub fn layout_summary(&self) -> String {
layout_summary(self.shape(), self.strides(), self.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
self.strides(),
self.offset(),
"TypedTensorView::assert_col_major_contiguous",
)
}
pub fn get(&self, indices: &[usize]) -> Option<&T> {
let offset = self.linear_offset(indices)?;
match &self.buffer {
TensorStorageRef::Host(data) => data.get(offset),
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => None,
}
}
pub fn as_slice(&self) -> crate::Result<&'a [T]> {
let data = match &self.buffer {
TensorStorageRef::Host(data) => data,
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
return Err(crate::Error::runtime_state(
"TypedTensorView::as_slice",
"backend buffers cannot be inspected as host slices; download explicitly first",
))
}
};
contiguous_layout_slice(self.layout(), data, "TypedTensorView::as_slice")
}
#[doc(hidden)]
pub fn with_host_read<U>(&self, f: impl FnOnce(&[T]) -> U) -> crate::Result<U>
where
T: TensorScalar + 'static,
{
const OP: &str = "TypedTensorView::with_host_read";
if !self.is_col_major_contiguous()? {
return Err(crate::Error::unsupported(
OP,
"host guard access requires a compact descriptor",
));
}
if let Some(buffer) = self.backend_buffer() {
let guard = buffer
.map_read()
.map_err(|error| crate::Error::host_access(OP, error))?;
return Ok(f(contiguous_layout_slice(self.layout(), &guard, OP)?));
}
if let Some(root) = &self.root {
let prepared = root
.prepare_host_read_for_layout(self.layout())
.map_err(|error| crate::Error::runtime_state_source(OP, error))?;
let slice = prepared.as_slice().ok_or_else(|| {
crate::Error::unsupported(OP, "host guard access requires a compact descriptor")
})?;
return Ok(f(slice));
}
Ok(f(self.as_slice()?))
}
pub fn duplicate(&self) -> crate::Result<TypedTensor<T, R>>
where
T: Clone,
{
let data = self.as_slice()?.to_vec();
let shape = R::shape_from_vec(shape_vec(self.shape()))
.map_err(|err| tensor_layout_error("TypedTensorView::duplicate", err))?;
let mut tensor = TypedTensor::from_vec_col_major(shape, data)?;
tensor.set_placement(self.placement.clone());
Ok(tensor)
}
pub fn to_col_major(&self) -> crate::Result<TypedTensor<T, R>>
where
T: Clone,
{
let data = self.host_storage()?;
let mut coordinates = shape_vec_zeros(self.shape().len());
let element_count = self.n_elements();
let mut copied = Vec::with_capacity(element_count);
for _ in 0..element_count {
let offset = self.layout_linear_offset(&coordinates)?;
let value = data.get(offset).ok_or_else(|| {
crate::Error::validation(
"TypedTensorView::to_col_major",
ValidationError::ViewOutOfBounds,
)
})?;
copied.push(value.clone());
for axis in 0..coordinates.len() {
coordinates[axis] += 1;
if coordinates[axis] < self.shape()[axis] {
break;
}
coordinates[axis] = 0;
}
}
let shape = R::shape_from_vec(shape_vec(self.shape()))
.map_err(|err| tensor_layout_error("TypedTensorView::to_col_major", err))?;
let mut tensor = TypedTensor::from_vec_col_major(shape, copied)?;
tensor.set_placement(self.placement.clone());
Ok(tensor)
}
pub fn transpose_view(&self, axes: impl AsRef<[usize]>) -> crate::Result<Self> {
let layout = self
.layout
.transpose_view(axes)
.map_err(|err| tensor_layout_error("TypedTensorView::transpose_view", err))?;
Ok(Self {
buffer: self.buffer.clone(),
root: self.root.clone(),
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
pub fn slice_view(&self, slices: &[StridedSliceSpec]) -> crate::Result<Self> {
let specs = core_slice_specs(slices, self.shape(), "TypedTensorView::slice_view")?;
let layout = self
.layout
.slice_view(specs, self.buffer.len())
.map_err(|err| tensor_layout_error("TypedTensorView::slice_view", err))?;
Ok(Self {
buffer: self.buffer.clone(),
root: self.root.clone(),
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
pub fn slice_axis_view(&self, axis: usize, slice: StridedSliceSpec) -> crate::Result<Self> {
let slices = slice_axis_specs(
self.shape().len(),
axis,
slice,
"TypedTensorView::slice_axis_view",
)?;
self.slice_view(&slices)
}
pub fn reshape_view(
&self,
shape: &[usize],
) -> crate::Result<TypedTensorView<'a, T, DynRank, D>> {
let layout = reshape_layout_dyn(
&self.layout,
shape,
self.buffer.len(),
"TypedTensorView::reshape_view",
)?;
Ok(TypedTensorView {
buffer: self.buffer.clone(),
root: self.root.as_ref().map(GroupReadView::clone_dyn),
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorView<'a, Complex32, R> {
pub fn as_real_view(&self) -> crate::Result<TypedTensorView<'a, f32, DynRank>> {
let op = "TypedTensorView::as_real_view";
validate_representation_pair(op, DType::C32, DType::F32)?;
let layout = reinterpret_complex_to_real_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
let buffer = match &self.buffer {
TensorStorageRef::Host(data) => {
TensorStorageRef::Host(reinterpret_host_slice::<Complex32, f32>(data, op)?)
}
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorView {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorView<'a, Complex64, R> {
pub fn as_real_view(&self) -> crate::Result<TypedTensorView<'a, f64, DynRank>> {
let op = "TypedTensorView::as_real_view";
validate_representation_pair(op, DType::C64, DType::F64)?;
let layout = reinterpret_complex_to_real_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
let buffer = match &self.buffer {
TensorStorageRef::Host(data) => {
TensorStorageRef::Host(reinterpret_host_slice::<Complex64, f64>(data, op)?)
}
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorView {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorView<'a, f32, R> {
pub fn as_complex_view(&self) -> crate::Result<TypedTensorView<'a, Complex32, DynRank>> {
let op = "TypedTensorView::as_complex_view";
validate_representation_pair(op, DType::F32, DType::C32)?;
let layout = reinterpret_real_to_complex_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
let buffer = match &self.buffer {
TensorStorageRef::Host(data) => {
TensorStorageRef::Host(reinterpret_host_slice::<f32, Complex32>(data, op)?)
}
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorView {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorView<'a, f64, R> {
pub fn as_complex_view(&self) -> crate::Result<TypedTensorView<'a, Complex64, DynRank>> {
let op = "TypedTensorView::as_complex_view";
validate_representation_pair(op, DType::F64, DType::C64)?;
let layout = reinterpret_real_to_complex_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
let buffer = match &self.buffer {
TensorStorageRef::Host(data) => {
TensorStorageRef::Host(reinterpret_host_slice::<f64, Complex64>(data, op)?)
}
TensorStorageRef::Backend(_) | TensorStorageRef::Root(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorView {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
pub struct TypedTensorViewMut<'a, T, R: TensorRank = DynRank, D: Representation = Dynamic> {
buffer: TensorStorageRefMut<'a, T>,
root: Option<GroupWriteView<'a, T, R>>,
layout: TensorLayout<R>,
placement: Placement,
_representation: std::marker::PhantomData<D>,
}
impl<'a, T, R: TensorRank, D: Representation> std::fmt::Debug for TypedTensorViewMut<'a, T, R, D> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("TypedTensorViewMut")
.field("shape", &self.layout.shape())
.field("placement", &self.placement)
.finish_non_exhaustive()
}
}
pub type TypedTensorViewMutSplit<'a, T, R = DynRank> =
(TypedTensorViewMut<'a, T, R>, TypedTensorViewMut<'a, T, R>);
impl<'a, T: 'static> TypedTensorViewMut<'a, T, DynRank> {
pub fn from_col_major(shape: &[usize], data: &'a mut [T]) -> crate::Result<Self> {
let layout = TensorLayout::<DynRank>::compact(shape_vec(shape))
.map_err(|err| tensor_layout_error("TypedTensorViewMut::from_col_major", err))?;
Self::from_buffer_ref_mut(
shape_vec(layout.shape()),
stride_vec(layout.strides()),
layout.offset(),
TensorStorageRefMut::Host(data),
default_placement(),
"TypedTensorViewMut::from_col_major",
)
}
pub fn from_slice(
shape: impl AsRef<[usize]>,
strides: impl AsRef<[isize]>,
offset: isize,
data: &'a mut [T],
) -> crate::Result<Self> {
Self::from_buffer_ref_mut(
shape_vec(shape.as_ref()),
stride_vec(strides.as_ref()),
offset,
TensorStorageRefMut::Host(data),
default_placement(),
"TypedTensorViewMut::from_slice",
)
}
}
impl<'a, T: 'static, R: TensorRank, D: Representation> TypedTensorViewMut<'a, T, R, D> {
pub fn from_slice_ranked(
shape: impl Into<R::Shape>,
strides: impl Into<R::Strides>,
offset: isize,
data: &'a mut [T],
) -> crate::Result<Self> {
Self::from_buffer_ref_mut(
shape,
strides,
offset,
TensorStorageRefMut::Host(data),
default_placement(),
"TypedTensorViewMut::from_slice_ranked",
)
}
}
impl<'a, T: 'static, R: TensorRank, D: Representation> TypedTensorViewMut<'a, T, R, D> {
fn from_buffer_ref_mut(
shape: impl Into<R::Shape>,
strides: impl Into<R::Strides>,
offset: isize,
buffer: TensorStorageRefMut<'a, T>,
placement: Placement,
op: &'static str,
) -> crate::Result<Self> {
let layout = TensorLayout::from_parts(shape.into(), strides.into(), offset, buffer.len())
.map_err(|err| tensor_layout_error(op, err))?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
Ok(Self {
buffer,
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
})
}
}
impl<'a, T: 'static, R: TensorRank> TypedTensorViewMut<'a, T, R, Host> {
pub fn from_host_slice(
shape: impl Into<R::Shape>,
strides: impl Into<R::Strides>,
offset: isize,
data: &'a mut [T],
) -> crate::Result<Self> {
let buffer = TensorStorageRefMut::Host(data);
let layout =
TensorLayout::from_parts(shape.into(), strides.into(), offset, buffer.len())
.map_err(|err| tensor_layout_error("TypedTensorViewMut::from_host_slice", err))?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::from_host_slice", err))?;
Ok(Self {
buffer,
root: None,
layout,
placement: default_placement(),
_representation: std::marker::PhantomData,
})
}
pub fn as_host_slice_mut(&mut self) -> &mut [T] {
match &mut self.buffer {
TensorStorageRefMut::Host(data) => data,
TensorStorageRefMut::Backend(_) => {
unreachable!("a Host-marked view always borrows host storage")
}
}
}
}
impl<'a, T: 'static, R: TensorRank, D: Representation> TypedTensorViewMut<'a, T, R, D> {
pub fn shape(&self) -> &[usize] {
self.layout.shape()
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn strides(&self) -> &[isize] {
self.layout.strides()
}
pub fn offset(&self) -> isize {
self.layout.offset()
}
pub fn host_storage(&self) -> crate::Result<&[T]> {
match &self.buffer {
TensorStorageRefMut::Host(data) => Ok(data),
TensorStorageRefMut::Backend(_) => Err(crate::Error::runtime_state(
"TypedTensorViewMut::host_storage",
"backend buffers cannot expose host storage; download explicitly first",
)),
}
}
pub fn host_storage_mut(&mut self) -> crate::Result<&mut [T]> {
match &mut self.buffer {
TensorStorageRefMut::Host(data) => Ok(data),
TensorStorageRefMut::Backend(_) => Err(crate::Error::runtime_state(
"TypedTensorViewMut::host_storage_mut",
"backend buffers cannot expose mutable host storage; download explicitly first",
)),
}
}
pub fn n_elements(&self) -> usize {
match checked_view_element_count(self.shape(), "TypedTensorViewMut::n_elements") {
Ok(n) => n,
Err(err) => {
unreachable!("TypedTensorViewMut layout shape is validated at construction: {err}")
}
}
}
pub fn layout(&self) -> &TensorLayout<R> {
&self.layout
}
pub fn placement(&self) -> &Placement {
&self.placement
}
#[doc(hidden)]
pub fn backend_buffer(&self) -> Option<&dyn BackendStorage<T>> {
match &self.buffer {
TensorStorageRefMut::Host(_) => None,
TensorStorageRefMut::Backend(buffer) => Some(&**buffer),
}
}
#[doc(hidden)]
pub fn prepare_device_write(
&mut self,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>>
where
T: TensorScalar + 'static,
{
let layout = self.layout.clone();
if self.root.is_none() {
let buffer = self.backend_buffer().ok_or_else(|| {
crate::Error::runtime_state_source(
op,
crate::AccessError::Unsupported { backend: "host" },
)
})?;
return prepare_backend_access(buffer, &self.layout, op);
}
let root = self
.root
.as_mut()
.ok_or_else(|| crate::Error::runtime_state(op, "expected a root-backed tensor view"))?;
root.prepare_device_write_for_layout(&layout)
.map_err(|error| crate::Error::runtime_state_source(op, error))
}
pub fn linear_offset(&self, indices: &[usize]) -> Option<usize> {
checked_view_offset(self.shape(), self.strides(), self.offset(), indices)
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
checked_view_offset_result(
self.shape(),
self.strides(),
self.offset(),
indices,
"TypedTensorViewMut::layout_linear_offset",
)
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
self.layout
.is_compact_col_major()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::is_col_major_contiguous", err))
}
pub fn layout_summary(&self) -> String {
layout_summary(self.shape(), self.strides(), self.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
self.strides(),
self.offset(),
"TypedTensorViewMut::assert_col_major_contiguous",
)
}
pub fn get(&self, indices: &[usize]) -> Option<&T> {
let offset = self.linear_offset(indices)?;
match &self.buffer {
TensorStorageRefMut::Host(data) => data.get(offset),
TensorStorageRefMut::Backend(_) => None,
}
}
pub fn get_mut(&mut self, indices: &[usize]) -> Option<&mut T> {
let offset = self.linear_offset(indices)?;
match &mut self.buffer {
TensorStorageRefMut::Host(data) => data.get_mut(offset),
TensorStorageRefMut::Backend(_) => None,
}
}
pub fn duplicate(&self) -> crate::Result<TypedTensor<T, R>>
where
T: Clone,
{
self.as_read_only().duplicate()
}
pub fn to_col_major(&self) -> crate::Result<TypedTensor<T, R>>
where
T: Clone,
{
self.as_read_only().to_col_major()
}
pub fn as_read_only(&self) -> TypedTensorView<'_, T, R, D> {
let buffer = match &self.buffer {
TensorStorageRefMut::Host(data) => TensorStorageRef::Host(data),
TensorStorageRefMut::Backend(buffer) => TensorStorageRef::Backend(&**buffer),
};
TypedTensorView {
buffer,
root: None,
layout: self.layout.clone(),
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
}
}
pub fn into_read_only(self) -> TypedTensorView<'a, T, R, D> {
let buffer = match self.buffer {
TensorStorageRefMut::Host(data) => TensorStorageRef::Host(data),
TensorStorageRefMut::Backend(buffer) => TensorStorageRef::Backend(buffer),
};
TypedTensorView {
buffer,
root: None,
layout: self.layout,
placement: self.placement,
_representation: std::marker::PhantomData,
}
}
pub fn transpose_view(
self,
axes: impl AsRef<[usize]>,
) -> crate::Result<TypedTensorViewMut<'a, T, R>> {
let Self {
buffer,
root,
layout,
placement,
_representation: std::marker::PhantomData,
} = self;
let layout = layout
.transpose_view(axes)
.map_err(|err| tensor_layout_error("TypedTensorViewMut::transpose_view", err))?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::transpose_view", err))?;
match buffer {
TensorStorageRefMut::Host(data) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(data),
root,
layout,
placement,
_representation: std::marker::PhantomData,
}),
TensorStorageRefMut::Backend(buffer) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(buffer),
root,
layout,
placement,
_representation: std::marker::PhantomData,
}),
}
}
pub fn slice_view(
&mut self,
slices: &[StridedSliceSpec],
) -> crate::Result<TypedTensorViewMut<'_, T, R>> {
let specs = core_slice_specs(slices, self.shape(), "TypedTensorViewMut::slice_view")?;
let layout = self
.layout
.slice_view(specs, self.buffer.len())
.map_err(|err| tensor_layout_error("TypedTensorViewMut::slice_view", err))?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::slice_view", err))?;
let placement = self.placement.clone();
match &mut self.buffer {
TensorStorageRefMut::Host(data) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(data),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
}),
TensorStorageRefMut::Backend(buffer) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(*buffer),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
}),
}
}
pub fn slice_axis_view(
&mut self,
axis: usize,
slice: StridedSliceSpec,
) -> crate::Result<TypedTensorViewMut<'_, T, R>> {
let slices = slice_axis_specs(
self.shape().len(),
axis,
slice,
"TypedTensorViewMut::slice_axis_view",
)?;
self.slice_view(&slices)
}
pub fn try_multi_slice_mut(
&mut self,
first: &[StridedSliceSpec],
second: &[StridedSliceSpec],
) -> crate::Result<Option<TypedTensorViewMutSplit<'_, T, R>>> {
let op = "TypedTensorViewMut::try_multi_slice_mut";
let first_specs = core_slice_specs(first, self.shape(), op)?;
let second_specs = core_slice_specs(second, self.shape(), op)?;
let buffer_len = self.buffer.len();
let first_layout = self
.layout
.slice_view(first_specs, buffer_len)
.map_err(|err| tensor_layout_error(op, err))?;
let second_layout = self
.layout
.slice_view(second_specs, buffer_len)
.map_err(|err| tensor_layout_error(op, err))?;
first_layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
second_layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
match (
reachable_layout_span(
first_layout.shape(),
first_layout.strides(),
first_layout.offset(),
)?,
reachable_layout_span(
second_layout.shape(),
second_layout.strides(),
second_layout.offset(),
)?,
) {
(Some(first_span), Some(second_span)) => {
let first_offset = adjusted_view_offset(first_layout.offset(), first_span.0)?;
let second_offset = adjusted_view_offset(second_layout.offset(), second_span.0)?;
let (first_data, second_data) = match &mut self.buffer {
TensorStorageRefMut::Host(data) => {
match split_two_mut_ranges(data, first_span, second_span) {
Some(ranges) => ranges,
None => return Ok(None),
}
}
TensorStorageRefMut::Backend(_) => return Ok(None),
};
let first_view = view_mut_from_layout_and_slice(
&first_layout,
first_offset,
first_data,
self.placement.clone(),
)?;
let second_view = view_mut_from_layout_and_slice(
&second_layout,
second_offset,
second_data,
self.placement.clone(),
)?;
Ok(Some((first_view, second_view)))
}
(None, Some(second_span)) => {
let second_offset = adjusted_view_offset(second_layout.offset(), second_span.0)?;
let (_, after_start) = match &mut self.buffer {
TensorStorageRefMut::Host(data) => data.split_at_mut(second_span.0),
TensorStorageRefMut::Backend(_) => return Ok(None),
};
let (second_data, _) = after_start.split_at_mut(second_span.1 - second_span.0 + 1);
let first_view = view_mut_from_layout_and_slice(
&first_layout,
0,
&mut [],
self.placement.clone(),
)?;
let second_view = view_mut_from_layout_and_slice(
&second_layout,
second_offset,
second_data,
self.placement.clone(),
)?;
Ok(Some((first_view, second_view)))
}
(Some(first_span), None) => {
let first_offset = adjusted_view_offset(first_layout.offset(), first_span.0)?;
let (_, after_start) = match &mut self.buffer {
TensorStorageRefMut::Host(data) => data.split_at_mut(first_span.0),
TensorStorageRefMut::Backend(_) => return Ok(None),
};
let (first_data, _) = after_start.split_at_mut(first_span.1 - first_span.0 + 1);
let first_view = view_mut_from_layout_and_slice(
&first_layout,
first_offset,
first_data,
self.placement.clone(),
)?;
let second_view = view_mut_from_layout_and_slice(
&second_layout,
0,
&mut [],
self.placement.clone(),
)?;
Ok(Some((first_view, second_view)))
}
(None, None) => {
let first_view = view_mut_from_layout_and_slice(
&first_layout,
0,
&mut [],
self.placement.clone(),
)?;
let second_view = view_mut_from_layout_and_slice(
&second_layout,
0,
&mut [],
self.placement.clone(),
)?;
Ok(Some((first_view, second_view)))
}
}
}
pub fn reshape_view(
&mut self,
shape: &[usize],
) -> crate::Result<TypedTensorViewMut<'_, T, DynRank, D>> {
let layout = reshape_layout_dyn(
&self.layout,
shape,
self.buffer.len(),
"TypedTensorViewMut::reshape_view",
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::reshape_view", err))?;
let placement = self.placement.clone();
match &mut self.buffer {
TensorStorageRefMut::Host(data) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(data),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
}),
TensorStorageRefMut::Backend(buffer) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(*buffer),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
}),
}
}
}
impl<'a, R: TensorRank> TypedTensorViewMut<'a, Complex32, R> {
pub fn as_real_view_mut(&mut self) -> crate::Result<TypedTensorViewMut<'_, f32, DynRank>> {
let op = "TypedTensorViewMut::as_real_view_mut";
validate_representation_pair(op, DType::C32, DType::F32)?;
let layout = reinterpret_complex_to_real_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
let buffer = match &mut self.buffer {
TensorStorageRefMut::Host(data) => {
TensorStorageRefMut::Host(reinterpret_host_slice_mut::<Complex32, f32>(data, op)?)
}
TensorStorageRefMut::Backend(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorViewMut<'a, Complex64, R> {
pub fn as_real_view_mut(&mut self) -> crate::Result<TypedTensorViewMut<'_, f64, DynRank>> {
let op = "TypedTensorViewMut::as_real_view_mut";
validate_representation_pair(op, DType::C64, DType::F64)?;
let layout = reinterpret_complex_to_real_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
let buffer = match &mut self.buffer {
TensorStorageRefMut::Host(data) => {
TensorStorageRefMut::Host(reinterpret_host_slice_mut::<Complex64, f64>(data, op)?)
}
TensorStorageRefMut::Backend(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorViewMut<'a, f32, R> {
pub fn as_complex_view_mut(
&mut self,
) -> crate::Result<TypedTensorViewMut<'_, Complex32, DynRank>> {
let op = "TypedTensorViewMut::as_complex_view_mut";
validate_representation_pair(op, DType::F32, DType::C32)?;
let layout = reinterpret_real_to_complex_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
let buffer = match &mut self.buffer {
TensorStorageRefMut::Host(data) => {
TensorStorageRefMut::Host(reinterpret_host_slice_mut::<f32, Complex32>(data, op)?)
}
TensorStorageRefMut::Backend(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
impl<'a, R: TensorRank> TypedTensorViewMut<'a, f64, R> {
pub fn as_complex_view_mut(
&mut self,
) -> crate::Result<TypedTensorViewMut<'_, Complex64, DynRank>> {
let op = "TypedTensorViewMut::as_complex_view_mut";
validate_representation_pair(op, DType::F64, DType::C64)?;
let layout = reinterpret_real_to_complex_layout(
self.shape(),
self.strides(),
self.offset(),
self.buffer.len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
let buffer = match &mut self.buffer {
TensorStorageRefMut::Host(data) => {
TensorStorageRefMut::Host(reinterpret_host_slice_mut::<f64, Complex64>(data, op)?)
}
TensorStorageRefMut::Backend(_) => {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
))
}
};
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
}
pub trait TensorScalar: Copy + Clone + Send + Sync + 'static + private::Sealed {
type Real: TensorScalar;
fn dtype() -> DType;
fn into_default_scalars(
shape: Vec<usize>,
data: Vec<Self>,
) -> crate::Result<crate::DefaultScalars>;
fn default_scalars_slice(set: &crate::DefaultScalars) -> Option<&[Self]>;
fn default_scalars_slice_mut(set: &mut crate::DefaultScalars) -> Option<&mut [Self]>;
fn from_default_scalars(set: crate::DefaultScalars)
-> Option<TypedTensor<Self, DynRank, Host>>;
fn into_tensor(shape: Vec<usize>, data: Vec<Self>) -> crate::Result<Tensor>;
fn typed_tensor_into_tensor(tensor: TypedTensor<Self>) -> Tensor;
fn tensor_read(tensor: &TypedTensor<Self>) -> TensorRead<'_>;
fn tensor_view<'a>(view: TypedTensorView<'a, Self>) -> TensorView<'a>;
fn tensor_view_mut<'a>(view: TypedTensorViewMut<'a, Self>) -> TensorViewMut<'a>;
fn tensor_write(tensor: &mut TypedTensor<Self>) -> TensorWrite<'_>;
fn as_slice(tensor: &Tensor) -> crate::Result<&[Self]>;
fn as_slice_mut(tensor: &mut Tensor) -> crate::Result<&mut [Self]>;
fn into_typed(
tensor: Tensor,
) -> std::result::Result<TypedTensor<Self>, ReinterpretError<Tensor>>;
}
mod private {
pub trait Sealed {}
impl Sealed for f64 {}
impl Sealed for f32 {}
impl Sealed for i32 {}
impl Sealed for i64 {}
impl Sealed for bool {}
impl Sealed for num_complex::Complex64 {}
impl Sealed for num_complex::Complex32 {}
}
macro_rules! impl_tensor_scalar {
($ty:ty, $real:ty, $dtype:ident, $variant:ident) => {
impl TensorScalar for $ty {
type Real = $real;
#[inline]
fn dtype() -> DType {
DType::$dtype
}
fn into_default_scalars(
shape: Vec<usize>,
data: Vec<Self>,
) -> crate::Result<crate::DefaultScalars> {
TypedTensor::<Self, DynRank, Host>::from_host_vec_col_major(shape, data).map(
|tensor| {
crate::DefaultScalars::from_payload(
crate::default_scalars::DefaultScalarsValue::$variant(tensor),
)
},
)
}
fn default_scalars_slice(set: &crate::DefaultScalars) -> Option<&[Self]> {
match set.payload() {
crate::default_scalars::DefaultScalarsValue::$variant(tensor) => {
Some(tensor.as_slice())
}
_ => None,
}
}
fn default_scalars_slice_mut(set: &mut crate::DefaultScalars) -> Option<&mut [Self]> {
match set.payload_mut() {
crate::default_scalars::DefaultScalarsValue::$variant(tensor) => {
Some(tensor.host_data_mut())
}
_ => None,
}
}
fn from_default_scalars(
set: crate::DefaultScalars,
) -> Option<TypedTensor<Self, DynRank, Host>> {
match set.into_payload() {
crate::default_scalars::DefaultScalarsValue::$variant(tensor) => Some(tensor),
_ => None,
}
}
fn into_tensor(shape: Vec<usize>, data: Vec<Self>) -> crate::Result<Tensor> {
TypedTensor::from_vec_col_major(shape, data).map(Self::typed_tensor_into_tensor)
}
fn typed_tensor_into_tensor(tensor: TypedTensor<Self>) -> Tensor {
Tensor {
payload: TensorPayload::Native(PresetTensor::$variant(tensor)),
}
}
fn tensor_read(tensor: &TypedTensor<Self>) -> TensorRead<'_> {
TensorRead::from_view(TensorView::$variant(tensor.as_view()))
}
#[inline]
fn tensor_view<'a>(view: TypedTensorView<'a, Self>) -> TensorView<'a> {
TensorView::$variant(view)
}
#[inline]
fn tensor_view_mut<'a>(view: TypedTensorViewMut<'a, Self>) -> TensorViewMut<'a> {
TensorViewMut::$variant(view)
}
fn tensor_write(tensor: &mut TypedTensor<Self>) -> TensorWrite<'_> {
TensorWrite::from_view(TensorViewMut::$variant(tensor.as_view_mut()))
}
fn as_slice(tensor: &Tensor) -> crate::Result<&[Self]> {
tensor
.as_typed::<Self>()
.ok_or_else(|| {
crate::Error::validation(
"Tensor::as_slice",
ValidationError::DTypeMismatch {
expected: Self::dtype(),
actual: tensor.dtype(),
},
)
})?
.host_data()
}
fn as_slice_mut(tensor: &mut Tensor) -> crate::Result<&mut [Self]> {
let actual = tensor.dtype();
let typed = tensor.as_typed_mut::<Self>().ok_or_else(|| {
crate::Error::validation(
"Tensor::as_slice_mut",
ValidationError::DTypeMismatch {
expected: Self::dtype(),
actual,
},
)
})?;
typed.host_data_mut()
}
fn into_typed(
tensor: Tensor,
) -> std::result::Result<TypedTensor<Self>, ReinterpretError<Tensor>> {
let actual = tensor.dtype();
match tensor.payload {
TensorPayload::Native(PresetTensor::$variant(typed)) => Ok(typed),
payload => Err(ReinterpretError::new(
Tensor { payload },
crate::Error::validation(
"TensorScalar::into_typed",
ValidationError::DTypeMismatch {
expected: Self::dtype(),
actual,
},
),
)),
}
}
}
};
}
impl_tensor_scalar!(f64, f64, F64, F64);
impl_tensor_scalar!(f32, f32, F32, F32);
impl_tensor_scalar!(i64, i64, I64, I64);
impl_tensor_scalar!(i32, i32, I32, I32);
impl_tensor_scalar!(bool, bool, Bool, Bool);
impl_tensor_scalar!(Complex64, f64, C64, C64);
impl_tensor_scalar!(Complex32, f32, C32, C32);
#[derive(Debug)]
enum TensorPayload {
Native(PresetTensor),
External(crate::ErasedHostTensor, Placement),
}
#[derive(Debug)]
enum PresetTensor {
F32(TypedTensor<f32>),
F64(TypedTensor<f64>),
I32(TypedTensor<i32>),
I64(TypedTensor<i64>),
Bool(TypedTensor<bool>),
C32(TypedTensor<Complex32>),
C64(TypedTensor<Complex64>),
}
macro_rules! with_preset {
($preset:expr, |$value:ident| $body:expr) => {
match $preset {
PresetTensor::F32($value) => $body,
PresetTensor::F64($value) => $body,
PresetTensor::I32($value) => $body,
PresetTensor::I64($value) => $body,
PresetTensor::Bool($value) => $body,
PresetTensor::C32($value) => $body,
PresetTensor::C64($value) => $body,
}
};
}
#[derive(Debug)]
pub struct Tensor {
payload: TensorPayload,
}
impl<R: TensorRank> OwnedTensorGroup<R> {
fn is_backend_buffer(&self) -> bool {
!matches!(
self.group.provider_kind(self.slot),
None | Some(crate::storage::ProviderKind::Cpu)
)
}
}
impl Tensor {
#[must_use]
pub fn external(payload: crate::ErasedHostTensor) -> Self {
Self {
payload: TensorPayload::External(payload, Placement::default()),
}
}
#[must_use]
pub fn from_typed<T: TensorScalar>(typed: TypedTensor<T>) -> Self {
T::typed_tensor_into_tensor(typed)
}
#[must_use]
pub fn external_payload(&self) -> Option<&crate::ErasedHostTensor> {
match &self.payload {
TensorPayload::External(payload, _) => Some(payload),
TensorPayload::Native(_) => None,
}
}
#[must_use]
pub fn external_with_placement(payload: crate::ErasedHostTensor, placement: Placement) -> Self {
Self {
payload: TensorPayload::External(payload, placement),
}
}
#[must_use]
pub fn external_payload_mut(&mut self) -> Option<&mut crate::ErasedHostTensor> {
match &mut self.payload {
TensorPayload::External(payload, _) => Some(payload),
TensorPayload::Native(_) => None,
}
}
pub(crate) fn into_group_parts(self) -> (AllocationGroup, DescriptorSlot) {
match self.payload {
TensorPayload::Native(preset) => with_preset!(preset, |typed| typed.into_group_parts()),
TensorPayload::External(..) => {
unreachable!("an externally defined payload has no allocation group")
}
}
}
}
#[derive(Clone, Debug)]
pub enum TensorView<'a> {
F32(TypedTensorView<'a, f32>),
F64(TypedTensorView<'a, f64>),
I32(TypedTensorView<'a, i32>),
I64(TypedTensorView<'a, i64>),
Bool(TypedTensorView<'a, bool>),
C32(TypedTensorView<'a, Complex<f32>>),
C64(TypedTensorView<'a, Complex<f64>>),
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
pub enum TensorViewMut<'a> {
F32(TypedTensorViewMut<'a, f32>),
F64(TypedTensorViewMut<'a, f64>),
I32(TypedTensorViewMut<'a, i32>),
I64(TypedTensorViewMut<'a, i64>),
Bool(TypedTensorViewMut<'a, bool>),
C32(TypedTensorViewMut<'a, Complex<f32>>),
C64(TypedTensorViewMut<'a, Complex<f64>>),
}
#[allow(clippy::large_enum_variant)]
#[derive(Clone, Debug)]
pub enum TensorRead<'a> {
Tensor(&'a Tensor),
View(TensorView<'a>),
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
pub enum TypedTensorWrite<'a, T> {
Tensor(&'a mut TypedTensor<T>),
View(TypedTensorViewMut<'a, T>),
}
impl<'a, T> TypedTensorWrite<'a, T> {
pub fn from_tensor(tensor: &'a mut TypedTensor<T>) -> Self {
Self::Tensor(tensor)
}
pub fn from_view(view: TypedTensorViewMut<'a, T>) -> Self {
Self::View(view)
}
}
impl<'a, T: TensorScalar> TypedTensorWrite<'a, T> {
pub fn into_tensor_write(self) -> TensorWrite<'a> {
match self {
Self::Tensor(tensor) => T::tensor_write(tensor),
Self::View(view) => TensorWrite::from_view(T::tensor_view_mut(view)),
}
}
}
impl<'a, T> From<&'a mut TypedTensor<T>> for TypedTensorWrite<'a, T> {
fn from(tensor: &'a mut TypedTensor<T>) -> Self {
Self::from_tensor(tensor)
}
}
impl<'a, T> From<TypedTensorViewMut<'a, T>> for TypedTensorWrite<'a, T> {
fn from(view: TypedTensorViewMut<'a, T>) -> Self {
Self::from_view(view)
}
}
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
pub enum TensorWrite<'a> {
Tensor(&'a mut Tensor),
View(TensorViewMut<'a>),
}
#[derive(Debug)]
pub struct TensorValue {
owner: Tensor,
layout: TensorLayout<DynRank>,
}
#[derive(Debug)]
pub struct TensorValueViewError {
value: TensorValue,
source: crate::Error,
}
impl TensorValueViewError {
pub fn into_parts(self) -> (TensorValue, crate::Error) {
(self.value, self.source)
}
fn new(value: TensorValue, source: crate::Error) -> Self {
Self { value, source }
}
}
impl std::fmt::Display for TensorValueViewError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.source, formatter)
}
}
impl std::error::Error for TensorValueViewError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(&self.source)
}
}
impl TensorValue {
pub fn duplicate(&self) -> crate::Result<Self> {
let tensor = duplicate_backing_storage(&self.owner)?;
Self::from_parts(
tensor,
self.shape().to_vec(),
self.strides().to_vec(),
self.offset(),
)
}
pub fn from_tensor(tensor: Tensor) -> Self {
let layout = tensor_layout(&tensor);
Self {
owner: tensor,
layout,
}
}
pub fn from_parts(
tensor: Tensor,
shape: Vec<usize>,
strides: Vec<isize>,
offset: isize,
) -> crate::Result<Self> {
let layout = TensorLayout::from_parts(
shape_vec(&shape),
stride_vec(&strides),
offset,
tensor_buffer_len(&tensor),
)
.map_err(|err| tensor_layout_error("TensorValue::from_parts", err))?;
Ok(Self {
owner: tensor,
layout,
})
}
#[doc(hidden)]
#[allow(clippy::result_large_err)]
pub fn try_into_group_parts(
self,
) -> std::result::Result<(AllocationGroup, DescriptorSlot, DType, Vec<usize>), Self> {
let Self { owner, layout } = self;
if owner.external_payload().is_some() {
return Err(Self { owner, layout });
}
let dtype = owner.dtype();
let shape = layout.shape().to_vec();
let strides = layout.strides().to_vec();
let offset = layout.offset();
let (group, slot) = owner.into_group_parts();
match group.update_descriptor_layout(slot, shape, strides, offset) {
Ok(group) => Ok((group, slot, dtype, layout.shape().to_vec())),
Err((_group, _error)) => {
unreachable!("TensorValue layout was validated before group ownership transfer")
}
}
}
pub fn into_tensor(self) -> crate::Result<Tensor> {
if self.layout != tensor_layout(&self.owner) {
return Err(crate::Error::unsupported(
"TensorValue::into_tensor",
"a metadata-only view has no compact tensor owner",
));
}
Ok(self.owner)
}
pub fn as_tensor(&self) -> Option<&Tensor> {
(self.layout == tensor_layout(&self.owner)).then_some(&self.owner)
}
pub fn is_view(&self) -> bool {
self.as_tensor().is_none()
}
pub fn dtype(&self) -> DType {
self.owner.dtype()
}
pub fn shape(&self) -> &[usize] {
self.layout.shape()
}
pub fn strides(&self) -> &[isize] {
self.layout.strides()
}
pub fn offset(&self) -> isize {
self.layout.offset()
}
pub fn tensor_view(&self) -> TensorView<'_> {
tensor_view_with_layout(&self.owner, self.layout.clone())
}
pub fn tensor_read(&self) -> TensorRead<'_> {
self.as_tensor()
.map(TensorRead::from_tensor)
.unwrap_or_else(|| TensorRead::from_view(self.tensor_view()))
}
pub fn transpose_view(self, axes: impl AsRef<[usize]>) -> crate::Result<Self> {
let layout = self
.layout
.transpose_view(axes)
.map_err(|err| tensor_layout_error("TensorValue::transpose_view", err))?;
Ok(Self {
owner: self.owner,
layout,
})
}
#[allow(clippy::result_large_err)]
pub fn try_reshape_view(
self,
shape: impl tenferro_tensor_core::IntoShapeVec,
) -> std::result::Result<Self, TensorValueViewError> {
let shape = shape.into_shape_vec();
let layout = match reshape_layout_dyn(
&self.layout,
&shape,
tensor_buffer_len(&self.owner),
"TensorValue::reshape_view",
) {
Ok(layout) => layout,
Err(error) => return Err(TensorValueViewError::new(self, error)),
};
Ok(Self {
owner: self.owner,
layout,
})
}
pub fn reshape_view(
self,
shape: impl tenferro_tensor_core::IntoShapeVec,
) -> crate::Result<Self> {
self.try_reshape_view(shape).map_err(|error| error.source)
}
pub fn slice_view(self, config: &SliceConfig) -> crate::Result<Self> {
let op = "TensorValue::slice_view";
if config.starts.len() != self.shape().len()
|| config.limits.len() != self.shape().len()
|| config.strides.len() != self.shape().len()
{
return Err(crate::Error::validation(
op,
ValidationError::RankMismatch {
expected: self.shape().len(),
actual: config.starts.len(),
},
));
}
let mut specs = TinyVec::<[CoreSliceSpec; 8]>::with_capacity(self.shape().len());
for (((&start, &limit), &stride), &axis_len) in config
.starts
.iter()
.zip(config.limits.iter())
.zip(config.strides.iter())
.zip(self.shape())
{
let start = isize::try_from(start).map_err(|_| {
crate::Error::invalid_argument(
op,
"slice start",
"slice start does not fit in isize",
)
})?;
let limit = isize::try_from(limit).map_err(|_| {
crate::Error::invalid_argument(
op,
"slice limit",
"slice limit does not fit in isize",
)
})?;
let stride = isize::try_from(stride).map_err(|_| {
crate::Error::invalid_argument(
op,
"slice stride",
"slice stride does not fit in isize",
)
})?;
specs.push(core_slice_spec(
StridedSliceSpec::new(start, Some(limit), stride),
axis_len,
op,
)?);
}
let layout = self
.layout
.slice_view(&specs, tensor_buffer_len(&self.owner))
.map_err(|err| tensor_layout_error(op, err))?;
Ok(Self {
owner: self.owner,
layout,
})
}
pub fn broadcast_in_dim_view(
self,
shape: impl tenferro_tensor_core::IntoShapeVec,
dims: impl AsRef<[usize]>,
) -> crate::Result<Self> {
let shape = shape.into_shape_vec();
let layout = self
.layout
.broadcast_in_dim_view::<DynRank>(shape, dims, tensor_buffer_len(&self.owner))
.map_err(|err| tensor_layout_error("TensorValue::broadcast_in_dim_view", err))?;
Ok(Self {
owner: self.owner,
layout,
})
}
}
fn duplicate_backing_storage(tensor: &Tensor) -> crate::Result<Tensor> {
fn typed<T: TensorScalar>(tensor: &TypedTensor<T>) -> crate::Result<Tensor> {
let layout = tensor.physical_layout();
let root = tensor.as_view().host_storage()?.to_vec();
let op = "TensorValue::duplicate";
let (group, slot) =
AllocationGroup::from_host_vec::<T, DynRank>(shape_vec(&[root.len()]), root)
.map_err(|error| crate::Error::runtime_state_source(op, error))?;
let group = group
.update_descriptor_layout(
slot,
layout.shape().to_vec(),
layout.strides().to_vec(),
layout.offset(),
)
.map_err(|(_, error)| crate::Error::runtime_state_source(op, error))?;
let allocation = group
.allocation_index(slot)
.unwrap_or_else(|_| unreachable!("the duplicate group carries a live descriptor"));
let descriptor_layout = group
.descriptor_layout(slot)
.unwrap_or_else(|| unreachable!("the duplicate group carries a live descriptor"));
let placement = tensor.placement().clone();
Ok(tensor_from_group(
group,
slot,
allocation,
T::dtype(),
descriptor_layout,
placement,
))
}
match &tensor.payload {
TensorPayload::External(..) => tensor.duplicate(),
TensorPayload::Native(preset) => with_preset!(preset, |inner| typed(inner)),
}
}
fn tensor_layout(tensor: &Tensor) -> TensorLayout<DynRank> {
match &tensor.payload {
TensorPayload::External(payload, _) => TensorLayout::<DynRank>::from_parts(
shape_vec(payload.shape()),
stride_vec(payload.strides()),
payload.offset(),
payload.payload_element_count(),
)
.unwrap_or_else(|_| {
unreachable!("a validated payload yields a layout inside its own storage")
}),
TensorPayload::Native(preset) => with_preset!(preset, |typed| typed.physical_layout()),
}
}
fn tensor_buffer_len(tensor: &Tensor) -> usize {
match &tensor.payload {
TensorPayload::External(payload, _) => payload.element_count(),
TensorPayload::Native(preset) => with_preset!(preset, |typed| typed.buffer_len()),
}
}
fn prepare_backend_access<'a, T: 'static, R: TensorRank>(
buffer: &'a dyn BackendStorage<T>,
layout: &'a TensorLayout<R>,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + 'a>> {
let domain = buffer.allocation_domain().ok_or_else(|| {
crate::Error::runtime_state(op, "backend buffer is missing an allocation domain")
})?;
let allocation_id = buffer.allocation_id().ok_or_else(|| {
crate::Error::runtime_state(op, "backend buffer is missing an allocation identity")
})?;
let byte_len = buffer
.len()
.checked_mul(size_of::<T>())
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
let request = DeviceAccessRequest::new(
domain,
allocation_id,
byte_len,
size_of::<T>(),
layout.shape(),
layout.strides(),
layout.offset(),
);
buffer
.prepare_device_access(request)
.map_err(|error| crate::Error::runtime_state_source(op, error))
}
fn cast_view_slice<S: 'static, T: TensorScalar>(source: &[S]) -> crate::Result<&[T]> {
if size_of::<S>() != size_of::<T>() || align_of::<S>() != align_of::<T>() {
return Err(crate::Error::invalid_argument(
"TensorView::as_slice",
"dtype",
"matching dtypes must have identical scalar layout",
));
}
Ok(unsafe { std::slice::from_raw_parts(source.as_ptr().cast::<T>(), source.len()) })
}
fn broadcast_typed_view<'a, T: 'static, R: TensorRank>(
view: TypedTensorView<'a, T, R>,
shape: &[usize],
dims: &[usize],
) -> crate::Result<TypedTensorView<'a, T, R>> {
let shape = R::shape_from_vec(shape_vec(shape))
.map_err(|err| tensor_layout_error("TensorRead::broadcast_in_dim_view", err))?;
let layout = view
.layout()
.broadcast_in_dim_view::<R>(shape, dims, view.backing_len())
.map_err(|err| tensor_layout_error("TensorRead::broadcast_in_dim_view", err))?;
Ok(TypedTensorView {
buffer: view.buffer,
root: view.root,
layout,
placement: view.placement,
_representation: std::marker::PhantomData,
})
}
fn tensor_view_with_layout(tensor: &Tensor, layout: TensorLayout<DynRank>) -> TensorView<'_> {
match tensor.dtype() {
DType::F32 => TensorView::F32(typed_view_with_layout(
tensor
.as_typed::<f32>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::F64 => TensorView::F64(typed_view_with_layout(
tensor
.as_typed::<f64>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::I32 => TensorView::I32(typed_view_with_layout(
tensor
.as_typed::<i32>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::I64 => TensorView::I64(typed_view_with_layout(
tensor
.as_typed::<i64>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::Bool => TensorView::Bool(typed_view_with_layout(
tensor
.as_typed::<bool>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::C32 => TensorView::C32(typed_view_with_layout(
tensor
.as_typed::<Complex<f32>>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::C64 => TensorView::C64(typed_view_with_layout(
tensor
.as_typed::<Complex<f64>>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm")),
layout,
)),
DType::External(_) => unreachable!("views cover the preset scalars"),
}
}
fn typed_view_with_layout<T: TensorScalar + 'static>(
tensor: &TypedTensor<T>,
layout: TensorLayout<DynRank>,
) -> TypedTensorView<'_, T> {
match &tensor.storage {
DynamicStorage::Host(HostStorage { data }) => TypedTensorView {
buffer: TensorStorageRef::Host(data.as_slice()),
root: None,
layout,
placement: tensor.placement.clone(),
_representation: std::marker::PhantomData,
},
DynamicStorage::Group(core) => {
let root = core.group.view::<T>().unwrap_or_else(|error| {
unreachable!("typed tensor group descriptor mismatch: {error}")
});
let buffer = if let Some(allocation) = root.backend_allocation() {
TensorStorageRef::Root(allocation)
} else {
TensorStorageRef::Host(tensor.group_host_slice())
};
TypedTensorView {
buffer,
root: Some(root),
layout,
placement: tensor.placement.clone(),
_representation: std::marker::PhantomData,
}
}
}
}
pub(crate) fn tensor_view_from_group<'a, T: TensorScalar>(
view: GroupReadView<'a, T, DynRank>,
) -> crate::Result<TensorView<'a>> {
let buffer = if let Some(allocation) = view.backend_allocation() {
TensorStorageRef::Root(allocation)
} else {
let storage = view.storage_buffer();
match storage {
Some(StorageBuffer::Host(data)) => TensorStorageRef::Host(data),
Some(StorageBuffer::Backend(buffer)) => TensorStorageRef::Backend(buffer.as_ref()),
None => match view.host_root_slice() {
Some(data) => TensorStorageRef::Host(data),
None => {
return Err(crate::Error::runtime_state(
"AllocationGroup::tensor_read",
"group descriptor has no backing storage",
));
}
},
}
};
let layout = view.descriptor().layout().clone();
let placement = view.descriptor().placement().clone();
let typed = TypedTensorView {
buffer,
root: Some(view.clone()),
layout,
placement,
_representation: std::marker::PhantomData,
};
Ok(T::tensor_view(typed))
}
pub(crate) fn tensor_from_group(
group: AllocationGroup,
slot: DescriptorSlot,
allocation_index: AllocationSlot,
dtype: DType,
layout: TensorLayout<DynRank>,
placement: Placement,
) -> Tensor {
fn typed<T: TensorScalar>(
group: AllocationGroup,
slot: DescriptorSlot,
allocation_index: AllocationSlot,
layout: TensorLayout<DynRank>,
placement: Placement,
) -> TypedTensor<T> {
let descriptor_dtype = group
.descriptor_dtype(slot)
.unwrap_or_else(|| unreachable!("tensor_from_group requires a live typed descriptor"));
assert_eq!(
descriptor_dtype,
T::dtype(),
"managed typed owner must match its descriptor dtype",
);
let (host_ptr, host_byte_len) = host_metadata::<T>(&group, slot);
TypedTensor {
shape: shape_vec(layout.shape()),
placement,
layout: stride_vec(layout.strides()),
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
}),
}),
}
}
match dtype {
DType::F32 => <f32 as TensorScalar>::typed_tensor_into_tensor(typed::<f32>(
group,
slot,
allocation_index,
layout,
placement,
)),
DType::F64 => <f64 as TensorScalar>::typed_tensor_into_tensor(typed::<f64>(
group,
slot,
allocation_index,
layout,
placement,
)),
DType::I32 => <i32 as TensorScalar>::typed_tensor_into_tensor(typed::<i32>(
group,
slot,
allocation_index,
layout,
placement,
)),
DType::I64 => <i64 as TensorScalar>::typed_tensor_into_tensor(typed::<i64>(
group,
slot,
allocation_index,
layout,
placement,
)),
DType::Bool => <bool as TensorScalar>::typed_tensor_into_tensor(typed::<bool>(
group,
slot,
allocation_index,
layout,
placement,
)),
DType::C32 => <Complex<f32> as TensorScalar>::typed_tensor_into_tensor(
typed::<Complex<f32>>(group, slot, allocation_index, layout, placement),
),
DType::C64 => <Complex<f64> as TensorScalar>::typed_tensor_into_tensor(
typed::<Complex<f64>>(group, slot, allocation_index, layout, placement),
),
DType::External(_) => unreachable!("allocation groups are preset-typed"),
}
}
impl From<TypedTensor<f64>> for Tensor {
fn from(t: TypedTensor<f64>) -> Self {
Tensor::from_typed(t)
}
}
impl From<TypedTensor<f32>> for Tensor {
fn from(t: TypedTensor<f32>) -> Self {
Tensor::from_typed(t)
}
}
impl From<TypedTensor<i64>> for Tensor {
fn from(t: TypedTensor<i64>) -> Self {
Tensor::from_typed(t)
}
}
impl From<TypedTensor<i32>> for Tensor {
fn from(t: TypedTensor<i32>) -> Self {
Tensor::from_typed(t)
}
}
impl From<TypedTensor<bool>> for Tensor {
fn from(t: TypedTensor<bool>) -> Self {
Tensor::from_typed(t)
}
}
impl From<TypedTensor<Complex<f64>>> for Tensor {
fn from(t: TypedTensor<Complex<f64>>) -> Self {
Tensor::from_typed(t)
}
}
impl From<TypedTensor<Complex<f32>>> for Tensor {
fn from(t: TypedTensor<Complex<f32>>) -> Self {
Tensor::from_typed(t)
}
}
impl<'a> TensorView<'a> {
pub fn f32(shape: &'a [usize], data: &'a [f32]) -> crate::Result<Self> {
Ok(Self::F32(TypedTensorView::from_col_major(shape, data)?))
}
pub fn f64(shape: &'a [usize], data: &'a [f64]) -> crate::Result<Self> {
Ok(Self::F64(TypedTensorView::from_col_major(shape, data)?))
}
pub fn i64(shape: &'a [usize], data: &'a [i64]) -> crate::Result<Self> {
Ok(Self::I64(TypedTensorView::from_col_major(shape, data)?))
}
pub fn i32(shape: &'a [usize], data: &'a [i32]) -> crate::Result<Self> {
Ok(Self::I32(TypedTensorView::from_col_major(shape, data)?))
}
pub fn bool(shape: &'a [usize], data: &'a [bool]) -> crate::Result<Self> {
Ok(Self::Bool(TypedTensorView::from_col_major(shape, data)?))
}
pub fn c32(shape: &'a [usize], data: &'a [Complex32]) -> crate::Result<Self> {
Ok(Self::C32(TypedTensorView::from_col_major(shape, data)?))
}
pub fn c64(shape: &'a [usize], data: &'a [Complex64]) -> crate::Result<Self> {
Ok(Self::C64(TypedTensorView::from_col_major(shape, data)?))
}
pub fn dtype(&self) -> DType {
match self {
Self::F32(_) => DType::F32,
Self::F64(_) => DType::F64,
Self::I32(_) => DType::I32,
Self::I64(_) => DType::I64,
Self::Bool(_) => DType::Bool,
Self::C32(_) => DType::C32,
Self::C64(_) => DType::C64,
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::F32(t) => t.shape(),
Self::F64(t) => t.shape(),
Self::I32(t) => t.shape(),
Self::I64(t) => t.shape(),
Self::Bool(t) => t.shape(),
Self::C32(t) => t.shape(),
Self::C64(t) => t.shape(),
}
}
pub fn as_slice<T: TensorScalar>(&self) -> crate::Result<&'a [T]> {
if self.dtype() != T::dtype() {
return Err(crate::Error::validation(
"TensorView::as_slice",
ValidationError::DTypeMismatch {
expected: T::dtype(),
actual: self.dtype(),
},
));
}
match self {
Self::F32(view) => cast_view_slice(view.as_slice()?),
Self::F64(view) => cast_view_slice(view.as_slice()?),
Self::I32(view) => cast_view_slice(view.as_slice()?),
Self::I64(view) => cast_view_slice(view.as_slice()?),
Self::Bool(view) => cast_view_slice(view.as_slice()?),
Self::C32(view) => cast_view_slice(view.as_slice()?),
Self::C64(view) => cast_view_slice(view.as_slice()?),
}
}
pub fn as_real_view(&self) -> crate::Result<Self> {
match self {
Self::C32(t) => t.as_real_view().map(Self::F32),
Self::C64(t) => t.as_real_view().map(Self::F64),
_ => Err(crate::Error::unsupported(
"TensorView::as_real_view",
"only complex views have a sealed real representation",
)),
}
}
pub fn as_complex_view(&self) -> crate::Result<Self> {
match self {
Self::F32(t) => t.as_complex_view().map(Self::C32),
Self::F64(t) => t.as_complex_view().map(Self::C64),
_ => Err(crate::Error::unsupported(
"TensorView::as_complex_view",
"only real views have a sealed complex representation",
)),
}
}
pub fn placement(&self) -> &Placement {
match self {
Self::F32(t) => t.placement(),
Self::F64(t) => t.placement(),
Self::I32(t) => t.placement(),
Self::I64(t) => t.placement(),
Self::Bool(t) => t.placement(),
Self::C32(t) => t.placement(),
Self::C64(t) => t.placement(),
}
}
pub fn backend_family(&self) -> Option<&'static str> {
match self {
Self::F32(t) => t.backend_family(),
Self::F64(t) => t.backend_family(),
Self::I32(t) => t.backend_family(),
Self::I64(t) => t.backend_family(),
Self::Bool(t) => t.backend_family(),
Self::C32(t) => t.backend_family(),
Self::C64(t) => t.backend_family(),
}
}
pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
match self {
Self::F32(t) => t.allocation_domain(),
Self::F64(t) => t.allocation_domain(),
Self::I32(t) => t.allocation_domain(),
Self::I64(t) => t.allocation_domain(),
Self::Bool(t) => t.allocation_domain(),
Self::C32(t) => t.allocation_domain(),
Self::C64(t) => t.allocation_domain(),
}
}
pub fn strides(&self) -> &[isize] {
match self {
Self::F32(t) => t.strides(),
Self::F64(t) => t.strides(),
Self::I32(t) => t.strides(),
Self::I64(t) => t.strides(),
Self::Bool(t) => t.strides(),
Self::C32(t) => t.strides(),
Self::C64(t) => t.strides(),
}
}
pub fn offset(&self) -> isize {
match self {
Self::F32(t) => t.offset(),
Self::F64(t) => t.offset(),
Self::I32(t) => t.offset(),
Self::I64(t) => t.offset(),
Self::Bool(t) => t.offset(),
Self::C32(t) => t.offset(),
Self::C64(t) => t.offset(),
}
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
match self {
Self::F32(t) => t.layout_linear_offset(indices),
Self::F64(t) => t.layout_linear_offset(indices),
Self::I32(t) => t.layout_linear_offset(indices),
Self::I64(t) => t.layout_linear_offset(indices),
Self::Bool(t) => t.layout_linear_offset(indices),
Self::C32(t) => t.layout_linear_offset(indices),
Self::C64(t) => t.layout_linear_offset(indices),
}
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
match self {
Self::F32(t) => t.is_col_major_contiguous(),
Self::F64(t) => t.is_col_major_contiguous(),
Self::I32(t) => t.is_col_major_contiguous(),
Self::I64(t) => t.is_col_major_contiguous(),
Self::Bool(t) => t.is_col_major_contiguous(),
Self::C32(t) => t.is_col_major_contiguous(),
Self::C64(t) => t.is_col_major_contiguous(),
}
}
pub fn layout_summary(&self) -> String {
layout_summary(self.shape(), self.strides(), self.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
self.strides(),
self.offset(),
"TensorView::assert_col_major_contiguous",
)
}
pub fn duplicate(&self) -> crate::Result<Tensor> {
fn duplicate_typed<T: TensorScalar>(
view: &TypedTensorView<'_, T>,
) -> crate::Result<TypedTensor<T>> {
let mut tensor = TypedTensor::<T>::from_vec_col_major(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
)?;
tensor.set_placement(view.placement().clone());
Ok(tensor)
}
match self {
Self::F32(view) => duplicate_typed(view).map(Tensor::from_typed),
Self::F64(view) => duplicate_typed(view).map(Tensor::from_typed),
Self::I32(view) => duplicate_typed(view).map(Tensor::from_typed),
Self::I64(view) => duplicate_typed(view).map(Tensor::from_typed),
Self::Bool(view) => duplicate_typed(view).map(Tensor::from_typed),
Self::C32(view) => duplicate_typed(view).map(Tensor::from_typed),
Self::C64(view) => duplicate_typed(view).map(Tensor::from_typed),
}
}
}
macro_rules! tensor_view_mut_constructor {
($name:ident, $variant:ident, $scalar:ty) => {
pub fn $name(shape: &'a [usize], data: &'a mut [$scalar]) -> crate::Result<Self> {
Ok(Self::$variant(TypedTensorViewMut::from_col_major(
shape, data,
)?))
}
};
}
impl<'a> TensorViewMut<'a> {
tensor_view_mut_constructor!(f32, F32, f32);
tensor_view_mut_constructor!(f64, F64, f64);
tensor_view_mut_constructor!(i32, I32, i32);
tensor_view_mut_constructor!(i64, I64, i64);
tensor_view_mut_constructor!(bool, Bool, bool);
tensor_view_mut_constructor!(c32, C32, Complex32);
tensor_view_mut_constructor!(c64, C64, Complex64);
pub fn dtype(&self) -> DType {
match self {
Self::F32(_) => DType::F32,
Self::F64(_) => DType::F64,
Self::I32(_) => DType::I32,
Self::I64(_) => DType::I64,
Self::Bool(_) => DType::Bool,
Self::C32(_) => DType::C32,
Self::C64(_) => DType::C64,
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::F32(t) => t.shape(),
Self::F64(t) => t.shape(),
Self::I32(t) => t.shape(),
Self::I64(t) => t.shape(),
Self::Bool(t) => t.shape(),
Self::C32(t) => t.shape(),
Self::C64(t) => t.shape(),
}
}
pub fn strides(&self) -> &[isize] {
match self {
Self::F32(t) => t.strides(),
Self::F64(t) => t.strides(),
Self::I32(t) => t.strides(),
Self::I64(t) => t.strides(),
Self::Bool(t) => t.strides(),
Self::C32(t) => t.strides(),
Self::C64(t) => t.strides(),
}
}
pub fn offset(&self) -> isize {
match self {
Self::F32(t) => t.offset(),
Self::F64(t) => t.offset(),
Self::I32(t) => t.offset(),
Self::I64(t) => t.offset(),
Self::Bool(t) => t.offset(),
Self::C32(t) => t.offset(),
Self::C64(t) => t.offset(),
}
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
match self {
Self::F32(t) => t.layout_linear_offset(indices),
Self::F64(t) => t.layout_linear_offset(indices),
Self::I32(t) => t.layout_linear_offset(indices),
Self::I64(t) => t.layout_linear_offset(indices),
Self::Bool(t) => t.layout_linear_offset(indices),
Self::C32(t) => t.layout_linear_offset(indices),
Self::C64(t) => t.layout_linear_offset(indices),
}
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
match self {
Self::F32(t) => t.is_col_major_contiguous(),
Self::F64(t) => t.is_col_major_contiguous(),
Self::I32(t) => t.is_col_major_contiguous(),
Self::I64(t) => t.is_col_major_contiguous(),
Self::Bool(t) => t.is_col_major_contiguous(),
Self::C32(t) => t.is_col_major_contiguous(),
Self::C64(t) => t.is_col_major_contiguous(),
}
}
pub fn layout_summary(&self) -> String {
layout_summary(self.shape(), self.strides(), self.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
self.strides(),
self.offset(),
"TensorViewMut::assert_col_major_contiguous",
)
}
pub fn duplicate(&self) -> crate::Result<Tensor> {
self.as_read_only().duplicate()
}
pub fn as_read_only(&self) -> TensorView<'_> {
match self {
Self::F32(t) => TensorView::F32(t.as_read_only()),
Self::F64(t) => TensorView::F64(t.as_read_only()),
Self::I32(t) => TensorView::I32(t.as_read_only()),
Self::I64(t) => TensorView::I64(t.as_read_only()),
Self::Bool(t) => TensorView::Bool(t.as_read_only()),
Self::C32(t) => TensorView::C32(t.as_read_only()),
Self::C64(t) => TensorView::C64(t.as_read_only()),
}
}
}
impl<'a> TensorRead<'a> {
pub fn from_tensor(tensor: &'a Tensor) -> Self {
Self::Tensor(tensor)
}
#[inline]
pub fn from_view(view: TensorView<'a>) -> Self {
Self::View(view)
}
pub fn tensor_view(self) -> TensorView<'a> {
match self {
Self::Tensor(tensor) => tensor_view_with_layout(tensor, tensor_layout(tensor)),
Self::View(view) => view,
}
}
#[doc(hidden)]
pub fn broadcast_in_dim_view(self, shape: &[usize], dims: &[usize]) -> crate::Result<Self> {
let dtype = self.dtype();
if let DType::External(_) = dtype {
return Err(crate::Error::unsupported_dtype(
"TensorRead::broadcast_in_dim_view",
dtype,
"a caller-owned payload has no borrowed broadcast view",
));
}
let view = self.tensor_view();
Ok(Self::View(match view {
TensorView::F32(view) => TensorView::F32(broadcast_typed_view(view, shape, dims)?),
TensorView::F64(view) => TensorView::F64(broadcast_typed_view(view, shape, dims)?),
TensorView::I32(view) => TensorView::I32(broadcast_typed_view(view, shape, dims)?),
TensorView::I64(view) => TensorView::I64(broadcast_typed_view(view, shape, dims)?),
TensorView::Bool(view) => TensorView::Bool(broadcast_typed_view(view, shape, dims)?),
TensorView::C32(view) => TensorView::C32(broadcast_typed_view(view, shape, dims)?),
TensorView::C64(view) => TensorView::C64(broadcast_typed_view(view, shape, dims)?),
}))
}
pub fn as_slice<T: TensorScalar>(&self) -> crate::Result<&'a [T]> {
self.clone().tensor_view().as_slice()
}
pub fn dtype(&self) -> DType {
match self {
Self::Tensor(tensor) => tensor.dtype(),
Self::View(view) => view.dtype(),
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::Tensor(tensor) => tensor.shape(),
Self::View(view) => view.shape(),
}
}
pub fn placement(&self) -> &Placement {
match self {
Self::Tensor(tensor) => tensor.placement(),
Self::View(view) => view.placement(),
}
}
pub fn backend_family(&self) -> Option<&'static str> {
match self {
Self::Tensor(tensor) => match tensor.dtype() {
DType::External(_) => None,
DType::F32 => tensor.as_typed::<f32>().and_then(|t| t.backend_family()),
DType::F64 => tensor.as_typed::<f64>().and_then(|t| t.backend_family()),
DType::I32 => tensor.as_typed::<i32>().and_then(|t| t.backend_family()),
DType::I64 => tensor.as_typed::<i64>().and_then(|t| t.backend_family()),
DType::Bool => tensor.as_typed::<bool>().and_then(|t| t.backend_family()),
DType::C32 => tensor
.as_typed::<Complex<f32>>()
.and_then(|t| t.backend_family()),
DType::C64 => tensor
.as_typed::<Complex<f64>>()
.and_then(|t| t.backend_family()),
},
Self::View(view) => view.backend_family(),
}
}
pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
match self {
Self::Tensor(tensor) => match tensor.dtype() {
DType::External(_) => None,
DType::F32 => tensor.as_typed::<f32>().and_then(|t| t.allocation_domain()),
DType::F64 => tensor.as_typed::<f64>().and_then(|t| t.allocation_domain()),
DType::I32 => tensor.as_typed::<i32>().and_then(|t| t.allocation_domain()),
DType::I64 => tensor.as_typed::<i64>().and_then(|t| t.allocation_domain()),
DType::Bool => tensor
.as_typed::<bool>()
.and_then(|t| t.allocation_domain()),
DType::C32 => tensor
.as_typed::<Complex<f32>>()
.and_then(|t| t.allocation_domain()),
DType::C64 => tensor
.as_typed::<Complex<f64>>()
.and_then(|t| t.allocation_domain()),
},
Self::View(view) => view.allocation_domain(),
}
}
pub fn strides(&self) -> crate::Result<Vec<isize>> {
match self {
Self::Tensor(tensor) => Ok(col_major_strides(tensor.shape())?.into_vec()),
Self::View(view) => Ok(view.strides().to_vec()),
}
}
pub fn offset(&self) -> isize {
match self {
Self::Tensor(_) => 0,
Self::View(view) => view.offset(),
}
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
match self {
Self::Tensor(tensor) => tensor.layout_linear_offset(indices),
Self::View(view) => view.layout_linear_offset(indices),
}
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
match self {
Self::Tensor(tensor) => tensor.is_col_major_contiguous(),
Self::View(view) => view.is_col_major_contiguous(),
}
}
pub fn layout_summary(&self) -> String {
let strides = match self.strides() {
Ok(strides) => strides,
Err(err) => return format!("layout unavailable: {err}"),
};
layout_summary(self.shape(), &strides, self.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
let strides = self.strides()?;
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
&strides,
self.offset(),
"TensorRead::assert_col_major_contiguous",
)
}
pub fn as_tensor(&self) -> Option<&'a Tensor> {
match self {
Self::Tensor(tensor) => Some(*tensor),
Self::View(_) => None,
}
}
}
impl<'a> TensorWrite<'a> {
pub fn from_tensor(tensor: &'a mut Tensor) -> Self {
Self::Tensor(tensor)
}
pub fn from_view(view: TensorViewMut<'a>) -> Self {
Self::View(view)
}
pub fn as_read(&self) -> TensorRead<'_> {
match self {
Self::Tensor(tensor) => TensorRead::from_tensor(tensor),
Self::View(view) => TensorRead::from_view(view.as_read_only()),
}
}
pub fn dtype(&self) -> DType {
match self {
Self::Tensor(tensor) => tensor.dtype(),
Self::View(view) => view.dtype(),
}
}
pub fn shape(&self) -> &[usize] {
match self {
Self::Tensor(tensor) => tensor.shape(),
Self::View(view) => view.shape(),
}
}
pub fn strides(&self) -> crate::Result<Vec<isize>> {
match self {
Self::Tensor(tensor) => Ok(col_major_strides(tensor.shape())?.into_vec()),
Self::View(view) => Ok(view.strides().to_vec()),
}
}
pub fn offset(&self) -> isize {
match self {
Self::Tensor(_) => 0,
Self::View(view) => view.offset(),
}
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
match self {
Self::Tensor(tensor) => tensor.layout_linear_offset(indices),
Self::View(view) => view.layout_linear_offset(indices),
}
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
match self {
Self::Tensor(tensor) => tensor.is_col_major_contiguous(),
Self::View(view) => view.is_col_major_contiguous(),
}
}
pub fn layout_summary(&self) -> String {
let strides = match self.strides() {
Ok(strides) => strides,
Err(err) => return format!("layout unavailable: {err}"),
};
layout_summary(self.shape(), &strides, self.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
let strides = self.strides()?;
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
&strides,
self.offset(),
"TensorWrite::assert_col_major_contiguous",
)
}
}
pub fn col_major_strides(shape: &[usize]) -> crate::Result<StrideVec> {
let mut strides = StrideVec::with_capacity(shape.len());
let mut stride = 1isize;
for &extent in shape {
strides.push(stride);
let extent = isize::try_from(extent).map_err(|_| {
crate::Error::validation("col_major_strides", ValidationError::IntegerOverflow)
})?;
stride = stride.checked_mul(extent).ok_or_else(|| {
crate::Error::validation("col_major_strides", ValidationError::IntegerOverflow)
})?;
}
Ok(strides)
}
fn try_linear_offset_for_shape(
shape: &[usize],
indices: &[usize],
op: &'static str,
) -> crate::Result<usize> {
if indices.len() != shape.len() {
return Err(crate::Error::validation(
op,
ValidationError::RankMismatch {
expected: shape.len(),
actual: indices.len(),
},
));
}
let mut offset = 0usize;
let mut stride = 1usize;
for (axis, (&idx, &extent)) in indices.iter().zip(shape).enumerate() {
if idx >= extent {
return Err(crate::Error::invalid_argument(
op,
"index",
format!("index {idx} out of bounds for axis {axis} extent {extent}"),
));
}
offset =
offset
.checked_add(idx.checked_mul(stride).ok_or_else(|| {
crate::Error::validation(op, ValidationError::IntegerOverflow)
})?)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
stride = stride
.checked_mul(extent)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
}
Ok(offset)
}
fn checked_view_offset_result(
shape: &[usize],
strides: &[isize],
base_offset: isize,
indices: &[usize],
op: &'static str,
) -> crate::Result<usize> {
if indices.len() != shape.len() {
return Err(crate::Error::validation(
op,
ValidationError::RankMismatch {
expected: shape.len(),
actual: indices.len(),
},
));
}
for (axis, (&index, &extent)) in indices.iter().zip(shape).enumerate() {
if index >= extent {
return Err(crate::Error::invalid_argument(
op,
"index",
format!("index {index} out of bounds for axis {axis} extent {extent}"),
));
}
}
checked_view_offset(shape, strides, base_offset, indices)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))
}
fn layout_summary(shape: &[usize], strides: &[isize], offset: isize) -> String {
format!("shape={shape:?} strides={strides:?} offset={offset}")
}
fn assert_layout_col_major_contiguous(
is_contiguous: bool,
shape: &[usize],
strides: &[isize],
offset: isize,
op: &'static str,
) -> crate::Result<()> {
if is_contiguous {
Ok(())
} else {
Err(crate::Error::invalid_argument(
op,
"layout",
format!(
"expected compact column-major layout, got {}",
layout_summary(shape, strides, offset)
),
))
}
}
fn try_shape_product(shape: &[usize], op: &'static str) -> crate::Result<usize> {
shape.iter().try_fold(1usize, |acc, &dim| {
acc.checked_mul(dim)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))
})
}
fn try_checked_shape_len(shape: &[usize], data_len: usize, op: &'static str) -> crate::Result<()> {
let n = try_shape_product(shape, op)?;
if data_len != n {
return Err(crate::Error::validation(
op,
ValidationError::ShapeDataLengthMismatch {
expected: n,
actual: data_len,
},
));
}
Ok(())
}
fn try_compact_layout<R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
op: &'static str,
) -> crate::Result<TensorLayout<R>> {
let shape = shape
.into_rank_shape()
.map_err(|err| tensor_layout_error(op, err))?;
TensorLayout::compact(shape).map_err(|err| tensor_layout_error(op, err))
}
fn tensor_layout_error(
op: &'static str,
err: tenferro_tensor_core::ValidationError,
) -> crate::Error {
crate::Error::validation(op, err)
}
fn checked_view_element_count(shape: &[usize], op: &'static str) -> crate::Result<usize> {
if shape.contains(&0) {
return Ok(0);
}
shape.iter().try_fold(1usize, |product, &dim| {
product
.checked_mul(dim)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))
})
}
fn checked_view_offset(
shape: &[usize],
strides: &[isize],
base_offset: isize,
indices: &[usize],
) -> Option<usize> {
if indices.len() != shape.len() {
return None;
}
let mut offset = base_offset;
for ((&index, &extent), &stride) in indices.iter().zip(shape).zip(strides) {
if index >= extent {
return None;
}
let index = isize::try_from(index).ok()?;
let delta = index.checked_mul(stride)?;
offset = offset.checked_add(delta)?;
}
usize::try_from(offset).ok()
}
fn reachable_layout_span(
shape: &[usize],
strides: &[isize],
offset: isize,
) -> crate::Result<Option<(usize, usize)>> {
if shape.contains(&0) {
return Ok(None);
}
let mut min_offset = offset;
let mut max_offset = offset;
for (&extent, &stride) in shape.iter().zip(strides) {
let steps = isize::try_from(extent.saturating_sub(1)).map_err(|_| {
crate::Error::validation(
"TypedTensorViewMut::try_multi_slice_mut",
ValidationError::IntegerOverflow,
)
})?;
let end = stride.checked_mul(steps).ok_or_else(|| {
crate::Error::validation(
"TypedTensorViewMut::try_multi_slice_mut",
ValidationError::IntegerOverflow,
)
})?;
let (axis_min, axis_max) = if end < 0 { (end, 0) } else { (0, end) };
min_offset = min_offset.checked_add(axis_min).ok_or_else(|| {
crate::Error::validation(
"TypedTensorViewMut::try_multi_slice_mut",
ValidationError::IntegerOverflow,
)
})?;
max_offset = max_offset.checked_add(axis_max).ok_or_else(|| {
crate::Error::validation(
"TypedTensorViewMut::try_multi_slice_mut",
ValidationError::IntegerOverflow,
)
})?;
}
let min_offset = usize::try_from(min_offset).map_err(|_| {
crate::Error::invalid_argument(
"TypedTensorViewMut::try_multi_slice_mut",
"layout",
"minimum reachable offset is negative",
)
})?;
let max_offset = usize::try_from(max_offset).map_err(|_| {
crate::Error::invalid_argument(
"TypedTensorViewMut::try_multi_slice_mut",
"layout",
"maximum reachable offset is negative",
)
})?;
Ok(Some((min_offset, max_offset)))
}
fn split_two_mut_ranges<T>(
data: &mut [T],
first: (usize, usize),
second: (usize, usize),
) -> Option<(&mut [T], &mut [T])> {
if first.1 < second.0 {
let (_, after_first_start) = data.split_at_mut(first.0);
let (first_slice, after_first) = after_first_start.split_at_mut(first.1 - first.0 + 1);
let (_, after_gap) = after_first.split_at_mut(second.0 - first.1 - 1);
let (second_slice, _) = after_gap.split_at_mut(second.1 - second.0 + 1);
Some((first_slice, second_slice))
} else if second.1 < first.0 {
let (_, after_second_start) = data.split_at_mut(second.0);
let (second_slice, after_second) = after_second_start.split_at_mut(second.1 - second.0 + 1);
let (_, after_gap) = after_second.split_at_mut(first.0 - second.1 - 1);
let (first_slice, _) = after_gap.split_at_mut(first.1 - first.0 + 1);
Some((first_slice, second_slice))
} else {
None
}
}
fn adjusted_view_offset(offset: isize, span_start: usize) -> crate::Result<isize> {
let span_start = isize::try_from(span_start).map_err(|_| {
crate::Error::validation(
"TypedTensorViewMut::try_multi_slice_mut",
ValidationError::IntegerOverflow,
)
})?;
offset.checked_sub(span_start).ok_or_else(|| {
crate::Error::validation(
"TypedTensorViewMut::try_multi_slice_mut",
ValidationError::IntegerOverflow,
)
})
}
fn view_mut_from_layout_and_slice<'a, T: 'static, R: TensorRank>(
layout: &TensorLayout<R>,
offset: isize,
data: &'a mut [T],
placement: Placement,
) -> crate::Result<TypedTensorViewMut<'a, T, R>> {
let shape = R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("TypedTensorViewMut::try_multi_slice_mut", err))?;
let strides = R::strides_from_vec(stride_vec(layout.strides()))
.map_err(|err| tensor_layout_error("TypedTensorViewMut::try_multi_slice_mut", err))?;
TypedTensorViewMut::from_buffer_ref_mut(
shape,
strides,
offset,
TensorStorageRefMut::Host(data),
placement,
"TypedTensorViewMut::try_multi_slice_mut",
)
}
fn contiguous_layout_slice<'a, T, R: TensorRank>(
layout: &TensorLayout<R>,
data: &'a [T],
op: &'static str,
) -> crate::Result<&'a [T]> {
if !layout
.is_compact_col_major()
.map_err(|err| tensor_layout_error(op, err))?
{
return Err(crate::Error::validation(
op,
ValidationError::NonContiguousViewAsSlice,
));
}
let len = checked_view_element_count(layout.shape(), op)?;
let start = usize::try_from(layout.offset())
.map_err(|_| crate::Error::invalid_argument(op, "layout", "view offset is negative"))?;
let end = start
.checked_add(len)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
data.get(start..end)
.ok_or_else(|| crate::Error::validation(op, ValidationError::ViewOutOfBounds))
}
fn relaxed_col_major_contiguous(
shape: &[usize],
strides: &[isize],
op: &'static str,
) -> crate::Result<bool> {
if shape.contains(&0) {
return Ok(true);
}
let mut expected = 1isize;
for (&extent, &stride) in shape.iter().zip(strides) {
if extent <= 1 {
continue;
}
if stride != expected {
return Ok(false);
}
let extent = isize::try_from(extent)
.map_err(|_| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
expected = expected
.checked_mul(extent)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
}
Ok(true)
}
fn reshape_layout_dyn<R: TensorRank>(
layout: &TensorLayout<R>,
shape: &[usize],
buffer_len: usize,
op: &'static str,
) -> crate::Result<TensorLayout<DynRank>> {
match layout.reshape_view_as::<DynRank>(shape_vec(shape), buffer_len) {
Ok(layout) => Ok(layout),
Err(err) => {
if !relaxed_col_major_contiguous(layout.shape(), layout.strides(), op)? {
return Err(tensor_layout_error(op, err));
}
let from = checked_view_element_count(layout.shape(), op)?;
let to = checked_view_element_count(shape, op)?;
if from != to {
return Err(tensor_layout_error(
op,
tenferro_tensor_core::ShapeMismatch::ReshapeElementCount { from, to }.into(),
));
}
TensorLayout::<DynRank>::compact(shape_vec(shape))
.and_then(|compact| {
TensorLayout::from_parts(
shape_vec(compact.shape()),
stride_vec(compact.strides()),
layout.offset(),
buffer_len,
)
})
.map_err(|err| tensor_layout_error(op, err))
}
}
}
fn core_slice_specs(
slices: &[StridedSliceSpec],
shape: &[usize],
op: &'static str,
) -> crate::Result<TinyVec<[CoreSliceSpec; 8]>> {
if slices.len() != shape.len() {
return Err(crate::Error::validation(
op,
ValidationError::RankMismatch {
expected: shape.len(),
actual: slices.len(),
},
));
}
let mut specs = TinyVec::<[CoreSliceSpec; 8]>::with_capacity(slices.len());
for (slice, &axis_len) in slices.iter().zip(shape) {
specs.push(core_slice_spec(*slice, axis_len, op)?);
}
Ok(specs)
}
fn core_slice_spec(
slice: StridedSliceSpec,
axis_len: usize,
op: &'static str,
) -> crate::Result<CoreSliceSpec> {
if slice.step() == 0 {
return Err(crate::Error::validation(
op,
ValidationError::InvalidSliceStep { step: slice.step() },
));
}
let start = normalize_strided_bound(slice.start(), axis_len, op, "slice start")?;
let end = match slice.end() {
Some(end) => normalize_strided_bound(end, axis_len, op, "slice end")?,
None => isize::try_from(axis_len)
.map_err(|_| crate::Error::validation(op, ValidationError::IntegerOverflow))?,
};
if slice.step() > 0 {
return Ok(CoreSliceSpec {
start,
end,
step: slice.step(),
});
}
if start >= end {
let empty_start = if axis_len == 0 {
start
} else {
start.min(axis_len as isize - 1)
};
return Ok(CoreSliceSpec {
start: empty_start,
end: empty_start,
step: slice.step(),
});
}
Ok(CoreSliceSpec {
start: end
.checked_sub(1)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?,
end: start
.checked_sub(1)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?,
step: slice.step(),
})
}
fn normalize_strided_bound(
bound: isize,
axis_len: usize,
op: &'static str,
role: &'static str,
) -> crate::Result<isize> {
let original_axis_len = axis_len;
let axis_len = isize::try_from(axis_len)
.map_err(|_| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
let bound = if bound < 0 {
axis_len
.checked_add(bound)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?
} else {
bound
};
if !(0..=axis_len).contains(&bound) {
let (start, end) = if role == "slice start" {
(bound, bound)
} else {
(0, bound)
};
return Err(crate::Error::validation(
op,
ValidationError::InvalidSliceBounds {
start,
end,
axis_len: original_axis_len,
},
));
}
Ok(bound)
}
fn slice_axis_specs(
rank: usize,
axis: usize,
slice: StridedSliceSpec,
op: &'static str,
) -> crate::Result<Vec<StridedSliceSpec>> {
if axis >= rank {
return Err(crate::Error::validation(
op,
ValidationError::AxisOutOfBounds { axis, rank },
));
}
let mut slices = vec![StridedSliceSpec::all(); rank];
slices[axis] = slice;
Ok(slices)
}
pub(crate) fn default_placement() -> Placement {
Placement {
memory_kind: MemoryKind::UnpinnedHost,
device: None,
cpu_affinity: None,
}
}
fn typed_tensor_from_vec_col_major<T, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
op: &'static str,
) -> crate::Result<TypedTensor<T, R>> {
try_typed_tensor_from_vec_col_major(shape, data, op)
}
fn try_typed_tensor_from_vec_col_major<T, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
op: &'static str,
) -> crate::Result<TypedTensor<T, R>> {
let shape = shape
.into_rank_shape()
.map_err(|err| tensor_layout_error(op, err))?;
let layout = tenferro_tensor_core::col_major_strides(shape.as_ref())
.map_err(|err| tensor_layout_error(op, err))?;
try_checked_shape_len(shape.as_ref(), data.len(), op)?;
Ok(TypedTensor {
shape,
placement: default_placement(),
layout,
storage: DynamicStorage::Host(HostStorage {
data: HostData::new(data),
}),
})
}
fn typed_tensor_zeros<T: TensorScalar + Zero, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
) -> crate::Result<TypedTensor<T, R>> {
try_typed_tensor_zeros(shape)
}
fn try_typed_tensor_zeros<T: TensorScalar + Clone + Zero, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
) -> crate::Result<TypedTensor<T, R>> {
let layout = try_compact_layout(shape, "zeros")?;
let n = try_shape_product(layout.shape(), "zeros")?;
typed_tensor_from_vec_col_major(
R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("zeros", err))?,
vec![T::zero(); n],
"zeros",
)
}
fn typed_tensor_ones<T: TensorScalar + One + Zero, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
) -> crate::Result<TypedTensor<T, R>> {
try_typed_tensor_ones(shape)
}
fn try_typed_tensor_ones<T: TensorScalar + Clone + One + Zero, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
) -> crate::Result<TypedTensor<T, R>> {
let layout = try_compact_layout(shape, "ones")?;
let n = try_shape_product(layout.shape(), "ones")?;
typed_tensor_from_vec_col_major(
R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("ones", err))?,
vec![T::one(); n],
"ones",
)
}
fn typed_tensor_from_buffer_col_major<T: TensorScalar + Send + Sync + 'static, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
buffer: StorageBuffer<T>,
placement: Placement,
) -> crate::Result<TypedTensor<T, R>> {
try_typed_tensor_from_buffer_col_major(shape, buffer, placement)
}
#[doc(hidden)]
fn typed_tensor_from_backend_allocation<T: TensorScalar + Send + Sync + 'static, R: TensorRank>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
allocation: Box<dyn crate::BackendAllocation>,
placement: Placement,
) -> crate::Result<TypedTensor<T, R>> {
let layout = try_compact_layout(shape, "from_backend_allocation")?;
let group_shape = R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("from_backend_allocation", err))?;
let shape = group_shape.clone();
let (mut group, slot) =
AllocationGroup::from_backend_allocation::<T, R>(group_shape, allocation)
.map_err(|error| group_error("TypedTensor::from_backend_allocation", error))?;
group
.set_descriptor_placement(slot, placement.clone())
.map_err(|error| group_error("TypedTensor::from_backend_allocation", error))?;
let allocation_index = group
.allocation_index(slot)
.map_err(|error| group_error("TypedTensor::from_backend_allocation", error))?;
let (host_ptr, host_byte_len) = host_metadata::<T>(&group, slot);
Ok(TypedTensor {
shape,
placement,
layout: stride_vec(layout.strides()),
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
}),
}),
})
}
fn try_typed_tensor_from_buffer_col_major<
T: TensorScalar + Send + Sync + 'static,
R: TensorRank,
>(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
buffer: StorageBuffer<T>,
placement: Placement,
) -> crate::Result<TypedTensor<T, R>> {
let layout = try_compact_layout(shape, "from_buffer_col_major")?;
let len = buffer.len();
try_checked_shape_len(layout.shape(), len, "from_buffer_col_major")?;
let group_shape = R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("from_buffer_col_major", err))?;
match buffer {
StorageBuffer::Host(data) => Ok(TypedTensor {
shape: group_shape,
placement,
layout: stride_vec(layout.strides()),
storage: DynamicStorage::Host(HostStorage {
data: HostData::new(data),
}),
}),
StorageBuffer::Backend(buffer) => {
let group = OwnedTensorGroup::from_backend_buffer(
group_shape.clone(),
StorageBuffer::Backend(buffer),
placement.clone(),
)?;
Ok(TypedTensor {
shape: group_shape,
placement,
layout: stride_vec(layout.strides()),
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(group),
}),
})
}
}
}
impl<T: TensorScalar + Zero, R: TensorRank> TypedTensor<T, R> {
pub fn zeros(shape: impl tenferro_tensor_core::IntoRankShape<R>) -> crate::Result<Self> {
typed_tensor_zeros(shape)
}
}
impl<T: TensorScalar + One + Zero, R: TensorRank> TypedTensor<T, R> {
pub fn ones(shape: impl tenferro_tensor_core::IntoRankShape<R>) -> crate::Result<Self> {
typed_tensor_ones(shape)
}
}
impl<T, R: TensorRank> TypedTensor<T, R> {
pub fn from_vec_col_major(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
) -> crate::Result<Self> {
typed_tensor_from_vec_col_major(shape, data, "from_vec_col_major")
}
pub fn from_vec_row_major(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
) -> crate::Result<Self>
where
T: Clone,
{
let op = "from_vec_row_major";
let (shape, reordered) = row_major_reorder(shape, data, op)?;
typed_tensor_from_vec_col_major(shape, reordered, op)
}
pub fn into_host_vec(self) -> std::result::Result<Vec<T>, ReinterpretError<Self>>
where
T: 'static,
{
let TypedTensor {
shape,
placement,
layout,
storage,
} = self;
match storage {
DynamicStorage::Host(HostStorage { data }) => Ok(data.into_vec()),
DynamicStorage::Group(GroupStorage { group }) => {
let range = match group_logical_range(&group, "TypedTensor::into_host_vec") {
Ok(range) => range,
Err(error) => {
return Err(ReinterpretError::new(
TypedTensor {
shape,
placement,
layout,
storage: DynamicStorage::Group(GroupStorage { group }),
},
error,
))
}
};
match group.into_host_vec::<T>() {
Ok(mut data) => {
debug_assert!(range.end <= data.len(), "descriptor range exceeds its root");
data.truncate(range.end);
data.drain(..range.start);
data.shrink_to_fit();
Ok(data)
}
Err((group, error)) => Err(ReinterpretError::new(
TypedTensor {
shape,
placement,
layout,
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(group),
}),
},
error,
)),
}
}
}
}
pub fn into_vec_col_major(
self,
) -> std::result::Result<(Vec<usize>, Vec<T>), ReinterpretError<Self>>
where
T: 'static,
{
let shape = self.shape().to_vec();
match self.into_host_vec() {
Ok(data) => Ok((shape, data)),
Err(failure) => Err(failure),
}
}
pub fn host_data(&self) -> crate::Result<&[T]> {
match &self.storage {
DynamicStorage::Host(HostStorage { data }) => Ok(data.as_slice()),
DynamicStorage::Group(core) => {
let data = core.group.host_slice::<T>()?;
let range = group_logical_range(&core.group, "TypedTensor::host_data")?;
data.get(range).ok_or_else(|| {
crate::Error::validation(
"TypedTensor::host_data",
ValidationError::ViewOutOfBounds,
)
})
}
}
}
pub fn as_slice(&self) -> crate::Result<&[T]> {
self.host_data()
}
pub fn host_data_mut(&mut self) -> crate::Result<&mut [T]> {
match &mut self.storage {
DynamicStorage::Host(HostStorage { data }) => Ok(data.as_mut_slice()),
DynamicStorage::Group(core) => {
let range = group_logical_range(&core.group, "TypedTensor::host_data_mut")?;
let data = core.group.host_slice_mut::<T>()?;
data.get_mut(range).ok_or_else(|| {
crate::Error::validation(
"TypedTensor::host_data_mut",
ValidationError::ViewOutOfBounds,
)
})
}
}
}
pub fn get(&self, indices: &[usize]) -> crate::Result<&T> {
let offset = self.linear_offset(indices)?;
self.host_data()?.get(offset).ok_or_else(|| {
crate::Error::validation("TypedTensor::get", ValidationError::ViewOutOfBounds)
})
}
pub fn get_mut(&mut self, indices: &[usize]) -> crate::Result<&mut T> {
let offset = self.linear_offset(indices)?;
self.host_data_mut()?.get_mut(offset).ok_or_else(|| {
crate::Error::validation("TypedTensor::get_mut", ValidationError::ViewOutOfBounds)
})
}
pub fn duplicate(&self) -> crate::Result<Self>
where
T: Clone,
{
let mut copy = Self::from_vec_col_major(
R::shape_from_vec(shape_vec(self.shape()))
.map_err(|err| tensor_layout_error("TypedTensor::duplicate", err))?,
self.host_data()?.to_vec(),
)?;
copy.set_placement(self.placement().clone());
Ok(copy)
}
pub fn from_buffer_col_major(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
buffer: StorageBuffer<T>,
placement: Placement,
) -> crate::Result<Self>
where
T: TensorScalar + Send + Sync + 'static,
{
typed_tensor_from_buffer_col_major(shape, buffer, placement)
}
#[doc(hidden)]
pub fn from_backend_allocation(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
allocation: Box<dyn crate::BackendAllocation>,
placement: Placement,
) -> crate::Result<Self>
where
T: TensorScalar + Send + Sync + 'static,
{
typed_tensor_from_backend_allocation(shape, allocation, placement)
}
pub fn try_into_rank<const N: usize>(
self,
) -> std::result::Result<TypedTensor<T, Rank<N>>, ReinterpretError<Self>> {
let op = "TypedTensor::try_into_rank";
let actual = self.shape().len();
let shape: [usize; N] = match self.shape().try_into() {
Ok(shape) => shape,
Err(_) => {
return Err(ReinterpretError::new(
self,
tensor_layout_error(
op,
ValidationError::RankMismatch {
expected: N,
actual,
},
),
))
}
};
let layout = match TensorLayout::<Rank<N>>::compact(shape) {
Ok(layout) => layout,
Err(error) => return Err(ReinterpretError::new(self, tensor_layout_error(op, error))),
};
let TypedTensor {
placement, storage, ..
} = self;
let storage = match storage {
DynamicStorage::Host(host) => DynamicStorage::Host(host),
DynamicStorage::Group(GroupStorage { group }) => DynamicStorage::Group(GroupStorage {
group: Box::new(OwnedTensorGroup {
group: group.group,
slot: group.slot,
allocation_index: group.allocation_index,
host_ptr: group.host_ptr,
host_byte_len: group.host_byte_len,
_rank: PhantomData,
}),
}),
};
Ok(TypedTensor {
shape,
placement,
layout: stride_vec(layout.strides()),
storage,
})
}
pub fn buffer(&self) -> &StorageBuffer<T>
where
T: 'static,
{
match &self.storage {
DynamicStorage::Host(HostStorage { data }) => &data.buffer,
DynamicStorage::Group(core) => core
.group
.host_buffer::<T>()
.or_else(|| core.group.backend_buffer::<T>())
.unwrap_or_else(|| unreachable!("typed tensor group storage mismatch")),
}
}
#[doc(hidden)]
pub fn backend_family(&self) -> Option<&'static str>
where
T: TensorScalar + 'static,
{
self.as_view().backend_family()
}
#[doc(hidden)]
pub fn backend_buffer(&self) -> Option<&dyn BackendStorage<T>>
where
T: 'static,
{
match &self.storage {
DynamicStorage::Host(..) => None,
DynamicStorage::Group(core) => match core.group.backend_buffer::<T>() {
Some(StorageBuffer::Backend(buffer)) => Some(buffer.as_ref()),
Some(StorageBuffer::Host(_)) | None => None,
},
}
}
#[doc(hidden)]
pub fn backend_buffer_mut(&mut self) -> Option<&mut dyn BackendStorage<T>>
where
T: 'static,
{
let DynamicStorage::Group(core) = &mut self.storage else {
return None;
};
match core.group.backend_buffer_mut::<T>()? {
StorageBuffer::Host(_) => None,
StorageBuffer::Backend(buffer) => Some(buffer.as_mut()),
}
}
#[doc(hidden)]
pub fn prepare_device_read(
&self,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>>
where
T: TensorScalar + 'static,
{
let DynamicStorage::Group(core) = &self.storage else {
return Err(crate::Error::runtime_state_source(
op,
crate::AccessError::Unsupported { backend: "host" },
));
};
let layout = self.physical_layout();
core.group
.prepare_device_read_for_layout::<T>(&layout)
.map_err(|error| crate::Error::runtime_state_source(op, error))
}
#[doc(hidden)]
pub fn prepare_device_write(
&mut self,
op: &'static str,
) -> crate::Result<Box<dyn PreparedDeviceAccess + '_>>
where
T: TensorScalar + 'static,
{
let layout = self.physical_layout();
let DynamicStorage::Group(core) = &mut self.storage else {
return Err(crate::Error::runtime_state_source(
op,
crate::AccessError::Unsupported { backend: "host" },
));
};
core.group
.prepare_device_write_for_layout::<T>(&layout)
.map_err(|error| crate::Error::runtime_state_source(op, error))
}
pub(crate) fn buffer_len(&self) -> usize
where
T: 'static,
{
match &self.storage {
DynamicStorage::Host(HostStorage { data }) => data.as_slice().len(),
DynamicStorage::Group(core) => core
.group
.group
.descriptor_backing_len(core.group.slot)
.unwrap_or_else(|| unreachable!("typed tensor group descriptor mismatch")),
}
}
pub fn allocation_domain(&self) -> Option<AllocationDomainId>
where
T: 'static,
{
match &self.storage {
DynamicStorage::Host(..) => None,
DynamicStorage::Group(core) => core
.group
.group
.backend_identity(core.group.slot)
.map(|(domain, _)| domain),
}
}
pub fn allocation_id(&self) -> Option<AllocationId>
where
T: 'static,
{
match &self.storage {
DynamicStorage::Host(..) => None,
DynamicStorage::Group(core) => core
.group
.group
.backend_identity(core.group.slot)
.map(|(_, allocation)| allocation),
}
}
pub fn as_view(&self) -> TypedTensorView<'_, T, R>
where
T: 'static,
{
let layout = self.physical_layout();
let placement = self.placement.clone();
match &self.storage {
DynamicStorage::Host(HostStorage { data }) => TypedTensorView {
buffer: TensorStorageRef::Host(data.as_slice()),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
},
DynamicStorage::Group(core) => {
let root = core
.group
.group
.view_raw::<T, R>(core.group.slot)
.unwrap_or_else(|error| {
unreachable!("typed tensor group descriptor mismatch: {error}")
});
let buffer = if let Some(allocation) = root.backend_allocation() {
TensorStorageRef::Root(allocation)
} else {
TensorStorageRef::Host(core.group.host_slice::<T>().unwrap_or_default())
};
TypedTensorView {
buffer,
root: Some(root),
layout,
placement,
_representation: std::marker::PhantomData,
}
}
}
}
pub fn as_view_mut(&mut self) -> TypedTensorViewMut<'_, T, R>
where
T: TensorScalar + 'static,
{
let layout = self.physical_layout();
let placement = self.placement.clone();
match &mut self.storage {
DynamicStorage::Host(HostStorage { data }) => TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(data.as_mut_slice()),
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
},
DynamicStorage::Group(core) => {
let mut root = core.group.view_mut::<T>().unwrap_or_else(|error| {
unreachable!("typed tensor group descriptor mismatch: {error}")
});
let buffer = if let Some(StorageBuffer::Backend(buffer)) = root.backend_buffer_mut()
{
TensorStorageRefMut::Backend(buffer.as_mut())
} else {
TensorStorageRefMut::Host(root.host_slice_mut().unwrap_or_else(|error| {
unreachable!("typed tensor group descriptor is not host-backed: {error}")
}))
};
TypedTensorViewMut {
buffer,
root: Some(root),
layout,
placement,
_representation: std::marker::PhantomData,
}
}
}
}
pub fn backend_region_view(
&self,
shape: Vec<usize>,
strides: Vec<isize>,
offset: isize,
) -> crate::Result<TypedTensorView<'_, T, DynRank>>
where
T: TensorScalar + 'static,
{
let op = "TypedTensor::backend_region_view";
let DynamicStorage::Group(core) = &self.storage else {
return Err(crate::Error::runtime_state(op, "expected a backend (device) allocation; host tensors use TypedTensorView::from_slice over host storage"));
};
let root = core.group.view_dyn::<T>()?;
let Some(allocation) = root.backend_allocation() else {
return Err(crate::Error::runtime_state(
op,
"expected a backend (device) allocation; host tensors use \
TypedTensorView::from_slice over host storage",
));
};
let element_len = allocation
.root_extent()
.byte_len()
.checked_div(size_of::<T>())
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
let layout =
TensorLayout::from_parts(shape_vec(&shape), stride_vec(&strides), offset, element_len)
.map_err(|err| tensor_layout_error(op, err))?;
Ok(TypedTensorView {
buffer: TensorStorageRef::Root(allocation),
root: Some(root),
layout,
placement: self.placement.clone(),
_representation: std::marker::PhantomData,
})
}
pub fn backend_region_view_mut(
&mut self,
shape: Vec<usize>,
strides: Vec<isize>,
offset: isize,
) -> crate::Result<TypedTensorViewMut<'_, T, DynRank>>
where
T: TensorScalar + 'static,
{
let op = "TypedTensor::backend_region_view_mut";
let placement = self.placement.clone();
let DynamicStorage::Group(core) = &mut self.storage else {
return Err(crate::Error::runtime_state(op, "expected a backend (device) buffer; mutable host regions use TypedTensorViewMut host constructors or try_multi_slice_mut"));
};
let mut root = core.group.view_mut_dyn::<T>()?;
let Some(StorageBuffer::Backend(buffer)) = root.backend_buffer_mut() else {
return Err(crate::Error::runtime_state(
op,
"expected a backend (device) buffer; mutable host regions use \
TypedTensorViewMut host constructors or try_multi_slice_mut",
));
};
let layout = TensorLayout::from_parts(
shape_vec(&shape),
stride_vec(&strides),
offset,
buffer.len(),
)
.map_err(|err| tensor_layout_error(op, err))?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(buffer.as_mut()),
root: Some(root),
layout,
placement,
_representation: std::marker::PhantomData,
})
}
pub fn into_parts(
self,
) -> std::result::Result<(StorageBuffer<T>, TensorLayout<R>, Placement), ReinterpretError<Self>>
where
T: TensorScalar,
{
let placement = self.placement().clone();
let layout = TensorLayout::<R>::compact(
R::shape_from_vec(shape_vec(self.shape()))
.unwrap_or_else(|err| unreachable!("validated owned shape: {err:?}")),
)
.unwrap_or_else(|err| unreachable!("validated owned shape: {err}"));
match self.into_host_vec() {
Ok(data) => Ok((StorageBuffer::Host(data), layout, placement)),
Err(error) => Err(error),
}
}
}
impl<T: TensorScalar, R: TensorRank> TypedTensor<T, R> {
fn into_managed(self) -> Self {
let TypedTensor {
shape,
placement,
layout,
storage,
} = self;
match storage {
DynamicStorage::Host(HostStorage { data }) => {
let group = promote_host_group::<T, R>(shape.clone(), data)
.unwrap_or_else(|err| unreachable!("a validated host owner promotes: {err}"));
Self {
shape,
placement,
layout,
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(group),
}),
}
}
DynamicStorage::Group(core) => Self {
shape,
placement,
layout,
storage: DynamicStorage::Group(core),
},
}
}
fn into_group_parts(self) -> (AllocationGroup, DescriptorSlot) {
let TypedTensor {
placement, storage, ..
} = self.into_managed();
let DynamicStorage::Group(GroupStorage { group }) = storage else {
unreachable!("explicit group promotion produces managed storage")
};
let (mut group, slot) = group.into_parts();
group.publish_live_descriptor_placement(slot, placement);
(group, slot)
}
#[doc(hidden)]
pub fn from_vec_col_major_with_recycler(
shape: impl tenferro_tensor_core::IntoRankShape<R>,
data: Vec<T>,
recycler: std::sync::Weak<dyn crate::HostBufferRecycler<T>>,
) -> crate::Result<Self> {
let mut tensor = Self::from_vec_col_major(shape, data)?;
if let DynamicStorage::Host(HostStorage { data }) = &mut tensor.storage {
data.recycler = Some(recycler);
}
Ok(tensor)
}
#[doc(hidden)]
pub fn with_host_read<U>(&self, f: impl FnOnce(&[T]) -> U) -> crate::Result<U>
where
T: TensorScalar + 'static,
{
self.as_view().with_host_read(f)
}
#[doc(hidden)]
pub fn with_host_write<U>(&mut self, f: impl FnOnce(&mut [T]) -> U) -> crate::Result<U>
where
T: TensorScalar + 'static,
{
match &mut self.storage {
DynamicStorage::Host(HostStorage { data }) => Ok(f(data.as_mut_slice())),
DynamicStorage::Group(core) => {
let slot = core.group.slot;
let mut view = core
.group
.group
.view_mut::<T, R>(slot)
.map_err(|error| group_error("TypedTensor::with_host_write", error))?;
let mut prepared = view.prepare_host_write().map_err(|error| {
crate::Error::runtime_state("TypedTensor::with_host_write", error.to_string())
})?;
let slice = prepared.as_slice_mut().ok_or_else(|| {
crate::Error::unsupported(
"TypedTensor::with_host_write",
"host guard access requires a compact descriptor",
)
})?;
Ok(f(slice))
}
}
}
fn backing_host_slice_mut(&mut self) -> crate::Result<&mut [T]> {
match &mut self.storage {
DynamicStorage::Host(HostStorage { data }) => Ok(data.as_mut_slice()),
DynamicStorage::Group(core) => core.group.host_slice_mut::<T>(),
}
}
fn group_host_slice(&self) -> &[T] {
let DynamicStorage::Group(core) = &self.storage else {
return &[];
};
core.group
.view::<T>()
.ok()
.and_then(|view| view.host_slice().ok())
.unwrap_or_default()
}
}
type ReinterpretLayoutFn =
fn(&[usize], &[isize], isize, usize, &'static str) -> crate::Result<TensorLayout<DynRank>>;
fn reinterpret_owned<S: TensorScalar, U: TensorScalar, R: TensorRank>(
owner: TypedTensor<S, R>,
from: DType,
to: DType,
make_layout: ReinterpretLayoutFn,
op: &'static str,
) -> Result<TypedTensor<U>, ReinterpretError<TypedTensor<S, R>>> {
if let Err(error) = validate_representation_pair(op, from, to) {
return Err(ReinterpretError::new(owner, error));
}
if matches!(to, DType::C32 | DType::C64) && !owner.buffer_len().is_multiple_of(2) {
return Err(ReinterpretError::new(
owner,
crate::Error::invalid_argument(
op,
"buffer",
"the owned real buffer must contain an even number of elements",
),
));
}
let source_layout = owner.physical_layout();
let target_layout = match make_layout(
source_layout.shape(),
source_layout.strides(),
source_layout.offset(),
owner.buffer_len(),
op,
) {
Ok(layout) => layout,
Err(error) => return Err(ReinterpretError::new(owner, error)),
};
let TypedTensor {
shape,
placement,
layout,
storage,
} = owner.into_managed();
let DynamicStorage::Group(GroupStorage { group }) = storage else {
unreachable!("explicit group promotion produces managed storage")
};
match group.reinterpret::<S, U>(
target_layout.shape().to_vec(),
target_layout.strides().to_vec(),
target_layout.offset(),
) {
Ok(group) => Ok(TypedTensor {
shape: shape_vec(target_layout.shape()),
placement,
layout: stride_vec(target_layout.strides()),
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(group),
}),
}),
Err((group, error)) => Err(ReinterpretError::new(
TypedTensor {
shape,
placement,
layout,
storage: DynamicStorage::Group(GroupStorage {
group: Box::new(group),
}),
},
error,
)),
}
}
impl<R: TensorRank> TypedTensor<Complex32, R> {
pub fn as_real_view(&self) -> crate::Result<TypedTensorView<'_, f32, DynRank>> {
self.as_view().as_real_view()
}
pub fn as_real_view_mut(&mut self) -> crate::Result<TypedTensorViewMut<'_, f32, DynRank>> {
let op = "TypedTensor::as_real_view_mut";
validate_representation_pair(op, DType::C32, DType::F32)?;
let source_layout = self.physical_layout();
let layout = reinterpret_complex_to_real_layout(
self.shape(),
source_layout.strides(),
source_layout.offset(),
self.buffer_len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
if self.backend_buffer().is_some() {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
));
}
let placement = self.placement().clone();
let buffer = TensorStorageRefMut::Host(reinterpret_host_slice_mut::<Complex32, f32>(
self.backing_host_slice_mut()?,
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
})
}
pub fn into_real(self) -> Result<TypedTensor<f32, DynRank>, ReinterpretError<Self>> {
reinterpret_owned(
self,
DType::C32,
DType::F32,
reinterpret_complex_to_real_layout,
"TypedTensor::into_real",
)
}
}
impl<R: TensorRank> TypedTensor<Complex64, R> {
pub fn as_real_view(&self) -> crate::Result<TypedTensorView<'_, f64, DynRank>> {
self.as_view().as_real_view()
}
pub fn as_real_view_mut(&mut self) -> crate::Result<TypedTensorViewMut<'_, f64, DynRank>> {
let op = "TypedTensor::as_real_view_mut";
validate_representation_pair(op, DType::C64, DType::F64)?;
let source_layout = self.physical_layout();
let layout = reinterpret_complex_to_real_layout(
self.shape(),
source_layout.strides(),
source_layout.offset(),
self.buffer_len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
if self.backend_buffer().is_some() {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
));
}
let placement = self.placement().clone();
let buffer = TensorStorageRefMut::Host(reinterpret_host_slice_mut::<Complex64, f64>(
self.backing_host_slice_mut()?,
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
})
}
pub fn into_real(self) -> Result<TypedTensor<f64, DynRank>, ReinterpretError<Self>> {
reinterpret_owned(
self,
DType::C64,
DType::F64,
reinterpret_complex_to_real_layout,
"TypedTensor::into_real",
)
}
}
impl<R: TensorRank> TypedTensor<f32, R> {
pub fn as_complex_view(&self) -> crate::Result<TypedTensorView<'_, Complex32, DynRank>> {
self.as_view().as_complex_view()
}
pub fn as_complex_view_mut(
&mut self,
) -> crate::Result<TypedTensorViewMut<'_, Complex32, DynRank>> {
let op = "TypedTensor::as_complex_view_mut";
validate_representation_pair(op, DType::F32, DType::C32)?;
let source_layout = self.physical_layout();
let layout = reinterpret_real_to_complex_layout(
self.shape(),
source_layout.strides(),
source_layout.offset(),
self.buffer_len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
if self.backend_buffer().is_some() {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
));
}
let placement = self.placement().clone();
let buffer = TensorStorageRefMut::Host(reinterpret_host_slice_mut::<f32, Complex32>(
self.backing_host_slice_mut()?,
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
})
}
pub fn into_complex(self) -> Result<TypedTensor<Complex32, DynRank>, ReinterpretError<Self>> {
reinterpret_owned(
self,
DType::F32,
DType::C32,
reinterpret_real_to_complex_layout,
"TypedTensor::into_complex",
)
}
}
impl<R: TensorRank> TypedTensor<f64, R> {
pub fn as_complex_view(&self) -> crate::Result<TypedTensorView<'_, Complex64, DynRank>> {
self.as_view().as_complex_view()
}
pub fn as_complex_view_mut(
&mut self,
) -> crate::Result<TypedTensorViewMut<'_, Complex64, DynRank>> {
let op = "TypedTensor::as_complex_view_mut";
validate_representation_pair(op, DType::F64, DType::C64)?;
let source_layout = self.physical_layout();
let layout = reinterpret_real_to_complex_layout(
self.shape(),
source_layout.strides(),
source_layout.offset(),
self.buffer_len(),
op,
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error(op, err))?;
if self.backend_buffer().is_some() {
return Err(crate::Error::unsupported(
op,
"backend representation reinterpretation is enabled by the provider phases",
));
}
let placement = self.placement().clone();
let buffer = TensorStorageRefMut::Host(reinterpret_host_slice_mut::<f64, Complex64>(
self.backing_host_slice_mut()?,
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
_representation: std::marker::PhantomData,
})
}
pub fn into_complex(self) -> Result<TypedTensor<Complex64, DynRank>, ReinterpretError<Self>> {
reinterpret_owned(
self,
DType::F64,
DType::C64,
reinterpret_real_to_complex_layout,
"TypedTensor::into_complex",
)
}
}
impl Tensor {
pub fn as_real_view(&self) -> crate::Result<TensorView<'_>> {
match self.dtype() {
DType::C32 => self
.as_typed::<Complex<f32>>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_real_view()
.map(TensorView::F32),
DType::C64 => self
.as_typed::<Complex<f64>>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_real_view()
.map(TensorView::F64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_real_view",
other,
DType::F32,
"only complex tensors have a sealed real representation view",
)),
}
}
pub fn as_real_view_mut(&mut self) -> crate::Result<TensorViewMut<'_>> {
match self.dtype() {
DType::C32 => self
.as_typed_mut::<Complex<f32>>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_real_view_mut()
.map(TensorViewMut::F32),
DType::C64 => self
.as_typed_mut::<Complex<f64>>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_real_view_mut()
.map(TensorViewMut::F64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_real_view_mut",
other,
DType::F32,
"only complex tensors have a sealed real representation view",
)),
}
}
pub fn into_real(self) -> Result<Self, ReinterpretError<Self>> {
let dtype = self.dtype();
match dtype {
DType::C32 => {
let tensor = <Complex<f32> as TensorScalar>::into_typed(self)
.unwrap_or_else(|_| unreachable!("the dtype guard selects this arm"));
tensor
.into_real()
.map(Tensor::from_typed::<f32>)
.map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::from_typed::<Complex<f32>>(owner), error)
})
}
DType::C64 => {
let tensor = <Complex<f64> as TensorScalar>::into_typed(self)
.unwrap_or_else(|_| unreachable!("the dtype guard selects this arm"));
tensor
.into_real()
.map(Tensor::from_typed::<f64>)
.map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::from_typed::<Complex<f64>>(owner), error)
})
}
_ => Err(ReinterpretError::new(
self,
crate::Error::unsupported(
"Tensor::into_real",
"only complex tensors have a sealed real representation",
),
)),
}
}
pub fn as_complex_view(&self) -> crate::Result<TensorView<'_>> {
match self.dtype() {
DType::F32 => self
.as_typed::<f32>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_complex_view()
.map(TensorView::C32),
DType::F64 => self
.as_typed::<f64>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_complex_view()
.map(TensorView::C64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_complex_view",
other,
DType::C32,
"only real tensors can have a sealed complex representation view",
)),
}
}
pub fn as_complex_view_mut(&mut self) -> crate::Result<TensorViewMut<'_>> {
match self.dtype() {
DType::F32 => self
.as_typed_mut::<f32>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_complex_view_mut()
.map(TensorViewMut::C32),
DType::F64 => self
.as_typed_mut::<f64>()
.unwrap_or_else(|| unreachable!("the dtype guard selects this arm"))
.as_complex_view_mut()
.map(TensorViewMut::C64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_complex_view_mut",
other,
DType::C32,
"only real tensors can have a sealed complex representation view",
)),
}
}
pub fn into_complex(self) -> Result<Self, ReinterpretError<Self>> {
let dtype = self.dtype();
match dtype {
DType::F32 => {
let tensor = <f32 as TensorScalar>::into_typed(self)
.unwrap_or_else(|_| unreachable!("the dtype guard selects this arm"));
tensor
.into_complex()
.map(Tensor::from_typed::<Complex<f32>>)
.map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::from_typed::<f32>(owner), error)
})
}
DType::F64 => {
let tensor = <f64 as TensorScalar>::into_typed(self)
.unwrap_or_else(|_| unreachable!("the dtype guard selects this arm"));
tensor
.into_complex()
.map(Tensor::from_typed::<Complex<f64>>)
.map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::from_typed::<f64>(owner), error)
})
}
_ => Err(ReinterpretError::new(
self,
crate::Error::unsupported(
"Tensor::into_complex",
"only real tensors have a sealed complex representation",
),
)),
}
}
pub fn duplicate(&self) -> crate::Result<Self> {
match &self.payload {
TensorPayload::External(payload, placement) => Ok(Self {
payload: TensorPayload::External(payload.duplicate(), placement.clone()),
}),
TensorPayload::Native(_) => match self.dtype() {
DType::F32 => self.duplicate_typed::<f32>(),
DType::F64 => self.duplicate_typed::<f64>(),
DType::I32 => self.duplicate_typed::<i32>(),
DType::I64 => self.duplicate_typed::<i64>(),
DType::Bool => self.duplicate_typed::<bool>(),
DType::C32 => self.duplicate_typed::<Complex<f32>>(),
DType::C64 => self.duplicate_typed::<Complex<f64>>(),
DType::External(_) => unreachable!("a native payload carries a preset dtype"),
},
}
}
fn duplicate_typed<T: TensorScalar>(&self) -> crate::Result<Self> {
self.as_typed::<T>()
.unwrap_or_else(|| unreachable!("the caller selected this arm from the dtype"))
.duplicate()
.map(Self::from_typed)
}
pub fn from_vec_col_major<T: TensorScalar>(
shape: impl tenferro_tensor_core::IntoShapeVec,
data: Vec<T>,
) -> crate::Result<Self> {
Ok(Self::from_typed(TypedTensor::<T>::from_vec_col_major(
shape, data,
)?))
}
pub fn from_vec_row_major<T: TensorScalar>(
shape: impl tenferro_tensor_core::IntoShapeVec,
data: Vec<T>,
) -> crate::Result<Self> {
let (shape, reordered) = row_major_reorder::<T, tenferro_tensor_core::DynRank>(
shape,
data,
"Tensor::from_vec_row_major",
)?;
Ok(Self::from_typed(TypedTensor::<T>::from_vec_col_major(
shape, reordered,
)?))
}
pub fn shape(&self) -> &[usize] {
match &self.payload {
TensorPayload::Native(preset) => with_preset!(preset, |typed| typed.shape()),
TensorPayload::External(payload, _) => payload.shape(),
}
}
pub fn dtype(&self) -> DType {
match &self.payload {
TensorPayload::Native(preset) => match preset {
PresetTensor::F32(_) => DType::F32,
PresetTensor::F64(_) => DType::F64,
PresetTensor::I32(_) => DType::I32,
PresetTensor::I64(_) => DType::I64,
PresetTensor::Bool(_) => DType::Bool,
PresetTensor::C32(_) => DType::C32,
PresetTensor::C64(_) => DType::C64,
},
TensorPayload::External(payload, _) => DType::External(payload.element_type_id()),
}
}
pub fn placement(&self) -> &Placement {
match &self.payload {
TensorPayload::Native(preset) => with_preset!(preset, |typed| typed.placement()),
TensorPayload::External(_, placement) => placement,
}
}
pub fn is_backend_buffer(&self) -> bool {
match &self.payload {
TensorPayload::Native(preset) => with_preset!(preset, |typed| {
matches!(&typed.storage, DynamicStorage::Group(core) if core.group.is_backend_buffer())
}),
TensorPayload::External(..) => false,
}
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
match &self.payload {
TensorPayload::Native(_) => {
let layout = tensor_layout(self);
checked_view_offset_result(
layout.shape(),
layout.strides(),
layout.offset(),
indices,
"Tensor::layout_linear_offset",
)
}
TensorPayload::External(..) => Err(crate::Error::unsupported_dtype(
"layout_linear_offset",
self.dtype(),
"an externally defined payload resolves its own offsets",
)),
}
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
match &self.payload {
TensorPayload::Native(_) => tensor_layout(self)
.is_compact_col_major()
.map_err(|err| tensor_layout_error("Tensor::is_col_major_contiguous", err)),
TensorPayload::External(..) => Ok(true),
}
}
pub fn layout_summary(&self) -> String {
let layout = tensor_layout(self);
layout_summary(layout.shape(), layout.strides(), layout.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
let layout = tensor_layout(self);
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
layout.shape(),
layout.strides(),
layout.offset(),
"Tensor::assert_col_major_contiguous",
)
}
pub fn as_slice<T: TensorScalar>(&self) -> crate::Result<&[T]> {
T::as_slice(self)
}
#[must_use]
pub fn as_typed<T: TensorScalar>(&self) -> Option<&TypedTensor<T>> {
if self.dtype() != T::dtype() {
return None;
}
match &self.payload {
TensorPayload::Native(preset) => with_preset!(preset, |typed| {
(typed as &dyn Any).downcast_ref::<TypedTensor<T>>()
}),
TensorPayload::External(..) => None,
}
}
#[must_use]
pub fn as_typed_mut<T: TensorScalar>(&mut self) -> Option<&mut TypedTensor<T>> {
if self.dtype() != T::dtype() {
return None;
}
match &mut self.payload {
TensorPayload::Native(preset) => with_preset!(preset, |typed| {
(typed as &mut dyn Any).downcast_mut::<TypedTensor<T>>()
}),
TensorPayload::External(..) => None,
}
}
pub fn into_typed<T: TensorScalar>(
self,
) -> std::result::Result<TypedTensor<T>, ReinterpretError<Self>> {
T::into_typed(self)
}
pub fn into_vec_col_major<T: TensorScalar>(
self,
) -> std::result::Result<(Vec<usize>, Vec<T>), ReinterpretError<Self>> {
let typed = T::into_typed(self)?;
match typed.into_vec_col_major() {
Ok(parts) => Ok(parts),
Err(failure) => {
let (owner, error) = failure.into_parts();
Err(ReinterpretError::new(Tensor::from_typed(owner), error))
}
}
}
}
#[allow(dead_code)]
pub(crate) fn flat_to_multi(mut flat: usize, shape: &[usize], out: &mut [usize]) {
for i in 0..shape.len() {
if shape[i] == 0 {
out[i] = 0;
} else {
out[i] = flat % shape[i];
flat /= shape[i];
}
}
}