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::ops::Deref;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::config::SliceConfig;
use crate::error::ReinterpretError;
pub use tenferro_tensor_core::{DynRank, Rank, TensorLayout, TensorRank};
use tenferro_tensor_core::{ShapeVec, StrideVec};
use tenferro_tensor_core::{SliceSpec as CoreSliceSpec, ValidationError};
use crate::storage::{
AllocationGroup, BackendAllocation, DescriptorSlot, GroupError, GroupReadView, GroupWriteView,
};
mod accessors;
mod shape_packing;
mod strided_view;
#[cfg(test)]
mod tests;
pub use strided_view::StridedSliceSpec;
fn shape_vec(shape: &[usize]) -> ShapeVec {
shape.iter().copied().collect()
}
fn stride_vec(strides: &[isize]) -> StrideVec {
strides.iter().copied().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 mut real_shape = Vec::with_capacity(
shape
.len()
.checked_add(1)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?,
);
real_shape.push(2);
real_shape.extend_from_slice(shape);
let real_strides = strides
.iter()
.map(|&stride| {
stride
.checked_mul(2)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))
})
.collect::<crate::Result<Vec<_>>>()?;
let mut all_strides = Vec::with_capacity(
real_strides
.len()
.checked_add(1)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?,
);
all_strides.push(1);
all_strides.extend(real_strides);
let real_offset = offset
.checked_mul(2)
.ok_or_else(|| crate::Error::validation(op, ValidationError::IntegerOverflow))?;
TensorLayout::from_parts(
real_shape.into(),
all_strides.into(),
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 complex_strides = strides[1..]
.iter()
.map(|&stride| {
if stride % 2 != 0 {
return Err(crate::Error::invalid_argument(
op,
"strides",
"all non-leading strides must be divisible by 2",
));
}
Ok(stride / 2)
})
.collect::<crate::Result<Vec<_>>>()?;
let complex_buffer_len = real_buffer_len / 2;
TensorLayout::from_parts(
shape[1..].to_vec().into(),
complex_strides.into(),
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),
}
}
}
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 TypedTensor<T, R: TensorRank = DynRank> {
group: OwnedTensorGroup<R>,
layout: TensorLayout<R>,
placement: Placement,
_scalar: PhantomData<T>,
}
struct OwnedTensorGroup<R: TensorRank> {
group: AllocationGroup,
slot: DescriptorSlot,
allocation_index: usize,
host_ptr: Option<usize>,
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("TypedTensor::prepare_device_read", error.to_string())
})
}
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("TypedTensor::prepare_device_write", error.to_string())
})
}
fn host_buffer<T: 'static>(&self) -> Option<&StorageBuffer<T>> {
self.group.host_buffer_at::<T>(self.allocation_index)
}
fn host_slice<T: 'static>(&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 as *const T, element_count) })
}
fn host_slice_mut<T: 'static>(&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 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)
}
fn into_host_vec<T: TensorScalar>(self) -> crate::Result<Vec<T>> {
self.group
.into_host_vec::<T>(self.slot)
.map_err(|error| crate::Error::runtime_state("TypedTensor::into_vec_col_major", 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: 'static>(
group: &AllocationGroup,
slot: DescriptorSlot,
) -> (Option<usize>, usize) {
group
.host_root_metadata::<T>(slot)
.map_or((None, 0), |(pointer, byte_len)| (Some(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
}
}
#[derive(Clone, Debug)]
pub struct TypedTensorView<'a, T, R: TensorRank = DynRank> {
buffer: TensorStorageRef<'a, T>,
root: Option<GroupReadView<'a, T, R>>,
layout: TensorLayout<R>,
placement: Placement,
}
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> TypedTensorView<'a, T, R> {
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",
)
}
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,
})
}
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(op, error.to_string()));
}
let buffer = self
.backend_buffer()
.ok_or_else(|| crate::Error::runtime_state(op, "expected a backend tensor view"))?;
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")
}
pub fn duplicate(&self) -> crate::Result<TypedTensor<T, R>>
where
T: TensorScalar,
{
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.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(),
})
}
pub fn try_slice(&self, slices: &[StridedSliceSpec]) -> crate::Result<Self> {
let specs = core_slice_specs(slices, self.shape(), "TypedTensorView::try_slice")?;
let layout = self
.layout
.slice_view(specs, self.buffer.len())
.map_err(|err| tensor_layout_error("TypedTensorView::try_slice", err))?;
Ok(Self {
buffer: self.buffer.clone(),
root: self.root.clone(),
layout,
placement: self.placement.clone(),
})
}
pub fn try_slice_axis(&self, axis: usize, slice: StridedSliceSpec) -> crate::Result<Self> {
let slices = slice_axis_specs(
self.shape().len(),
axis,
slice,
"TypedTensorView::try_slice_axis",
)?;
self.try_slice(&slices)
}
pub fn try_reshape(&self, shape: &[usize]) -> crate::Result<TypedTensorView<'a, T, DynRank>> {
let layout = reshape_layout_dyn(
&self.layout,
shape,
self.buffer.len(),
"TypedTensorView::try_reshape",
)?;
Ok(TypedTensorView {
buffer: self.buffer.clone(),
root: self.root.as_ref().map(GroupReadView::clone_dyn),
layout,
placement: self.placement.clone(),
})
}
}
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(),
})
}
}
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(),
})
}
}
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(),
})
}
}
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(),
})
}
}
#[derive(Debug)]
pub struct TypedTensorViewMut<'a, T, R: TensorRank = DynRank> {
buffer: TensorStorageRefMut<'a, T>,
root: Option<GroupWriteView<'a, T, R>>,
layout: TensorLayout<R>,
placement: Placement,
}
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> TypedTensorViewMut<'a, T, R> {
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",
)
}
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,
})
}
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(op, "expected a backend tensor view"))?;
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(op, error.to_string()))
}
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: TensorScalar,
{
self.as_read_only().duplicate()
}
pub fn as_read_only(&self) -> TypedTensorView<'_, T, R> {
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(),
}
}
pub fn into_read_only(self) -> TypedTensorView<'a, T, R> {
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,
}
}
pub fn transpose_view(
self,
axes: impl AsRef<[usize]>,
) -> crate::Result<TypedTensorViewMut<'a, T, R>> {
let Self {
buffer,
root,
layout,
placement,
} = 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,
}),
TensorStorageRefMut::Backend(buffer) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(buffer),
root,
layout,
placement,
}),
}
}
pub fn try_slice(
&mut self,
slices: &[StridedSliceSpec],
) -> crate::Result<TypedTensorViewMut<'_, T, R>> {
let specs = core_slice_specs(slices, self.shape(), "TypedTensorViewMut::try_slice")?;
let layout = self
.layout
.slice_view(specs, self.buffer.len())
.map_err(|err| tensor_layout_error("TypedTensorViewMut::try_slice", err))?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::try_slice", err))?;
let placement = self.placement.clone();
match &mut self.buffer {
TensorStorageRefMut::Host(data) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(data),
root: None,
layout,
placement,
}),
TensorStorageRefMut::Backend(buffer) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(*buffer),
root: None,
layout,
placement,
}),
}
}
pub fn try_slice_axis(
&mut self,
axis: usize,
slice: StridedSliceSpec,
) -> crate::Result<TypedTensorViewMut<'_, T, R>> {
let slices = slice_axis_specs(
self.shape().len(),
axis,
slice,
"TypedTensorViewMut::try_slice_axis",
)?;
self.try_slice(&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 try_reshape(
&mut self,
shape: &[usize],
) -> crate::Result<TypedTensorViewMut<'_, T, DynRank>> {
let layout = reshape_layout_dyn(
&self.layout,
shape,
self.buffer.len(),
"TypedTensorViewMut::try_reshape",
)?;
layout
.validate_mutable_no_overlap()
.map_err(|err| tensor_layout_error("TypedTensorViewMut::try_reshape", err))?;
let placement = self.placement.clone();
match &mut self.buffer {
TensorStorageRefMut::Host(data) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Host(data),
root: None,
layout,
placement,
}),
TensorStorageRefMut::Backend(buffer) => Ok(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(*buffer),
root: None,
layout,
placement,
}),
}
}
}
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(),
})
}
}
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(),
})
}
}
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(),
})
}
}
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(),
})
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum DType {
F32,
F64,
I32,
I64,
Bool,
C32,
C64,
}
pub trait TensorScalar: Copy + Clone + Send + Sync + 'static + private::Sealed {
type Real: TensorScalar;
fn dtype() -> DType;
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) -> crate::Result<TypedTensor<Self>>;
}
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;
fn dtype() -> DType {
DType::$dtype
}
fn into_tensor(shape: Vec<usize>, data: Vec<Self>) -> crate::Result<Tensor> {
TypedTensor::from_vec_col_major(shape, data).map(Tensor::$variant)
}
fn typed_tensor_into_tensor(tensor: TypedTensor<Self>) -> Tensor {
Tensor::$variant(tensor)
}
fn tensor_read(tensor: &TypedTensor<Self>) -> TensorRead<'_> {
TensorRead::from_view(TensorView::$variant(tensor.as_view()))
}
fn tensor_view<'a>(view: TypedTensorView<'a, Self>) -> TensorView<'a> {
TensorView::$variant(view)
}
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]> {
let actual = tensor.dtype();
match tensor {
Tensor::$variant(t) => t.host_data(),
_ => Err(crate::Error::validation(
"Tensor::as_slice",
ValidationError::DTypeMismatch {
expected: crate::core_dtype(Self::dtype()),
actual: crate::core_dtype(actual),
},
)),
}
}
fn as_slice_mut(tensor: &mut Tensor) -> crate::Result<&mut [Self]> {
let actual = tensor.dtype();
match tensor {
Tensor::$variant(t) => t.host_data_mut(),
_ => Err(crate::Error::validation(
"Tensor::as_slice_mut",
ValidationError::DTypeMismatch {
expected: crate::core_dtype(Self::dtype()),
actual: crate::core_dtype(actual),
},
)),
}
}
fn into_typed(tensor: Tensor) -> crate::Result<TypedTensor<Self>> {
let actual = tensor.dtype();
match tensor {
Tensor::$variant(inner) => Ok(inner),
_ => Err(crate::Error::validation(
"TensorScalar::into_typed",
ValidationError::DTypeMismatch {
expected: crate::core_dtype(Self::dtype()),
actual: crate::core_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)]
pub enum Tensor {
F32(TypedTensor<f32>),
F64(TypedTensor<f64>),
I32(TypedTensor<i32>),
I64(TypedTensor<i64>),
Bool(TypedTensor<bool>),
C32(TypedTensor<Complex<f32>>),
C64(TypedTensor<Complex<f64>>),
}
impl Tensor {
pub(crate) fn into_group_parts(self) -> (AllocationGroup, DescriptorSlot) {
match self {
Self::F32(tensor) => tensor.group.into_parts(),
Self::F64(tensor) => tensor.group.into_parts(),
Self::I32(tensor) => tensor.group.into_parts(),
Self::I64(tensor) => tensor.group.into_parts(),
Self::Bool(tensor) => tensor.group.into_parts(),
Self::C32(tensor) => tensor.group.into_parts(),
Self::C64(tensor) => tensor.group.into_parts(),
}
}
}
#[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 = self.owner.duplicate()?;
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.into(),
strides.into(),
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;
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 slices = Vec::with_capacity(self.shape().len());
for ((&start, &limit), &stride) in config
.starts
.iter()
.zip(config.limits.iter())
.zip(config.strides.iter())
{
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",
)
})?;
slices.push(StridedSliceSpec::new(start, Some(limit), stride));
}
let specs = core_slice_specs(&slices, self.shape(), 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 tensor_layout(tensor: &Tensor) -> TensorLayout<DynRank> {
match tensor {
Tensor::F32(tensor) => tensor.layout.clone(),
Tensor::F64(tensor) => tensor.layout.clone(),
Tensor::I32(tensor) => tensor.layout.clone(),
Tensor::I64(tensor) => tensor.layout.clone(),
Tensor::Bool(tensor) => tensor.layout.clone(),
Tensor::C32(tensor) => tensor.layout.clone(),
Tensor::C64(tensor) => tensor.layout.clone(),
}
}
fn tensor_buffer_len(tensor: &Tensor) -> usize {
match tensor {
Tensor::F32(tensor) => tensor.buffer_len(),
Tensor::F64(tensor) => tensor.buffer_len(),
Tensor::I32(tensor) => tensor.buffer_len(),
Tensor::I64(tensor) => tensor.buffer_len(),
Tensor::Bool(tensor) => tensor.buffer_len(),
Tensor::C32(tensor) => tensor.buffer_len(),
Tensor::C64(tensor) => tensor.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(op, error.to_string()))
}
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 tensor_view_with_layout(tensor: &Tensor, layout: TensorLayout<DynRank>) -> TensorView<'_> {
match tensor {
Tensor::F32(tensor) => TensorView::F32(typed_view_with_layout(tensor, layout)),
Tensor::F64(tensor) => TensorView::F64(typed_view_with_layout(tensor, layout)),
Tensor::I32(tensor) => TensorView::I32(typed_view_with_layout(tensor, layout)),
Tensor::I64(tensor) => TensorView::I64(typed_view_with_layout(tensor, layout)),
Tensor::Bool(tensor) => TensorView::Bool(typed_view_with_layout(tensor, layout)),
Tensor::C32(tensor) => TensorView::C32(typed_view_with_layout(tensor, layout)),
Tensor::C64(tensor) => TensorView::C64(typed_view_with_layout(tensor, layout)),
}
}
fn typed_view_with_layout<T: TensorScalar + 'static>(
tensor: &TypedTensor<T>,
layout: TensorLayout<DynRank>,
) -> TypedTensorView<'_, T> {
let root = match tensor.group.view::<T>() {
Ok(root) => root,
Err(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(),
}
}
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().ok_or_else(|| {
crate::Error::runtime_state(
"AllocationGroup::tensor_read",
"group descriptor has no backing storage",
)
})?;
match storage {
StorageBuffer::Host(data) => TensorStorageRef::Host(data),
StorageBuffer::Backend(buffer) => TensorStorageRef::Backend(buffer.as_ref()),
}
};
let layout = view.descriptor().layout().clone();
let placement = view.descriptor().placement().clone();
let typed = TypedTensorView {
buffer,
root: Some(view.clone()),
layout,
placement,
};
Ok(T::tensor_view(typed))
}
pub(crate) fn tensor_from_group(
group: AllocationGroup,
slot: DescriptorSlot,
allocation_index: usize,
dtype: DType,
layout: TensorLayout<DynRank>,
placement: Placement,
) -> Tensor {
fn typed<T: TensorScalar>(
group: AllocationGroup,
slot: DescriptorSlot,
allocation_index: usize,
layout: TensorLayout<DynRank>,
placement: Placement,
) -> TypedTensor<T> {
let (host_ptr, host_byte_len) = host_metadata::<T>(&group, slot);
TypedTensor {
group: OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
},
layout,
placement,
_scalar: PhantomData,
}
}
match dtype {
DType::F32 => Tensor::F32(typed(group, slot, allocation_index, layout, placement)),
DType::F64 => Tensor::F64(typed(group, slot, allocation_index, layout, placement)),
DType::I32 => Tensor::I32(typed(group, slot, allocation_index, layout, placement)),
DType::I64 => Tensor::I64(typed(group, slot, allocation_index, layout, placement)),
DType::Bool => Tensor::Bool(typed(group, slot, allocation_index, layout, placement)),
DType::C32 => Tensor::C32(typed(group, slot, allocation_index, layout, placement)),
DType::C64 => Tensor::C64(typed(group, slot, allocation_index, layout, placement)),
}
}
impl From<TypedTensor<f64>> for Tensor {
fn from(t: TypedTensor<f64>) -> Self {
Tensor::F64(t)
}
}
impl From<TypedTensor<f32>> for Tensor {
fn from(t: TypedTensor<f32>) -> Self {
Tensor::F32(t)
}
}
impl From<TypedTensor<i64>> for Tensor {
fn from(t: TypedTensor<i64>) -> Self {
Tensor::I64(t)
}
}
impl From<TypedTensor<i32>> for Tensor {
fn from(t: TypedTensor<i32>) -> Self {
Tensor::I32(t)
}
}
impl From<TypedTensor<bool>> for Tensor {
fn from(t: TypedTensor<bool>) -> Self {
Tensor::Bool(t)
}
}
impl From<TypedTensor<Complex<f64>>> for Tensor {
fn from(t: TypedTensor<Complex<f64>>) -> Self {
Tensor::C64(t)
}
}
impl From<TypedTensor<Complex<f32>>> for Tensor {
fn from(t: TypedTensor<Complex<f32>>) -> Self {
Tensor::C32(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: crate::core_dtype(T::dtype()),
actual: crate::core_dtype(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::F32),
Self::F64(view) => duplicate_typed(view).map(Tensor::F64),
Self::I32(view) => duplicate_typed(view).map(Tensor::I32),
Self::I64(view) => duplicate_typed(view).map(Tensor::I64),
Self::Bool(view) => duplicate_typed(view).map(Tensor::Bool),
Self::C32(view) => duplicate_typed(view).map(Tensor::C32),
Self::C64(view) => duplicate_typed(view).map(Tensor::C64),
}
}
}
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)
}
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,
}
}
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 {
Tensor::F32(t) => t.backend_family(),
Tensor::F64(t) => t.backend_family(),
Tensor::I32(t) => t.backend_family(),
Tensor::I64(t) => t.backend_family(),
Tensor::Bool(t) => t.backend_family(),
Tensor::C32(t) => t.backend_family(),
Tensor::C64(t) => t.backend_family(),
},
Self::View(view) => view.backend_family(),
}
}
pub fn allocation_domain(&self) -> Option<AllocationDomainId> {
match self {
Self::Tensor(tensor) => match tensor {
Tensor::F32(t) => t.allocation_domain(),
Tensor::F64(t) => t.allocation_domain(),
Tensor::I32(t) => t.allocation_domain(),
Tensor::I64(t) => t.allocation_domain(),
Tensor::Bool(t) => t.allocation_domain(),
Tensor::C32(t) => t.allocation_domain(),
Tensor::C64(t) => t.allocation_domain(),
},
Self::View(view) => view.allocation_domain(),
}
}
pub fn strides(&self) -> crate::Result<Vec<isize>> {
match self {
Self::Tensor(tensor) => col_major_strides(tensor.shape()),
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) => col_major_strides(tensor.shape()),
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<Vec<isize>> {
let mut strides = Vec::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 Into<R::Shape>,
op: &'static str,
) -> crate::Result<TensorLayout<R>> {
TensorLayout::compact(shape.into()).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::invalid_argument(
op,
"layout",
"view is not contiguous column-major",
));
}
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<Vec<CoreSliceSpec>> {
if slices.len() != shape.len() {
return Err(crate::Error::validation(
op,
ValidationError::RankMismatch {
expected: shape.len(),
actual: slices.len(),
},
));
}
let mut specs = Vec::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 {
return Ok(CoreSliceSpec {
start,
end: 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: TensorScalar, R: TensorRank>(
shape: impl Into<R::Shape>,
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 Into<R::Shape>,
data: Vec<T>,
op: &'static str,
) -> crate::Result<TypedTensor<T, R>>
where
T: TensorScalar,
{
let layout = try_compact_layout(shape, op)?;
try_checked_shape_len(layout.shape(), data.len(), op)?;
let group_shape =
R::shape_from_vec(shape_vec(layout.shape())).map_err(|err| tensor_layout_error(op, err))?;
let group = OwnedTensorGroup::from_host_vec(group_shape, data)?;
Ok(TypedTensor {
group,
layout,
placement: default_placement(),
_scalar: PhantomData,
})
}
fn typed_tensor_zeros<T: TensorScalar + Zero, R: TensorRank>(
shape: impl Into<R::Shape>,
) -> crate::Result<TypedTensor<T, R>> {
try_typed_tensor_zeros(shape)
}
fn try_typed_tensor_zeros<T: TensorScalar + Clone + Zero, R: TensorRank>(
shape: impl Into<R::Shape>,
) -> crate::Result<TypedTensor<T, R>> {
let layout = try_compact_layout(shape, "zeros")?;
let n = try_shape_product(layout.shape(), "zeros")?;
let group_shape = R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("zeros", err))?;
let group = OwnedTensorGroup::from_host_vec(group_shape, vec![T::zero(); n])?;
Ok(TypedTensor {
group,
layout,
placement: default_placement(),
_scalar: PhantomData,
})
}
fn typed_tensor_ones<T: TensorScalar + One + Zero, R: TensorRank>(
shape: impl Into<R::Shape>,
) -> crate::Result<TypedTensor<T, R>> {
try_typed_tensor_ones(shape)
}
fn try_typed_tensor_ones<T: TensorScalar + Clone + One + Zero, R: TensorRank>(
shape: impl Into<R::Shape>,
) -> crate::Result<TypedTensor<T, R>> {
let layout = try_compact_layout(shape, "ones")?;
let n = try_shape_product(layout.shape(), "ones")?;
let group_shape = R::shape_from_vec(shape_vec(layout.shape()))
.map_err(|err| tensor_layout_error("ones", err))?;
let group = OwnedTensorGroup::from_host_vec(group_shape, vec![T::one(); n])?;
Ok(TypedTensor {
group,
layout,
placement: default_placement(),
_scalar: PhantomData,
})
}
fn typed_tensor_from_buffer_col_major<T: TensorScalar + Send + Sync + 'static, R: TensorRank>(
shape: impl Into<R::Shape>,
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 Into<R::Shape>,
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 (group, slot) =
AllocationGroup::from_backend_allocation::<T, R>(group_shape, allocation)
.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 {
group: OwnedTensorGroup {
group,
slot,
allocation_index,
host_ptr,
host_byte_len,
_rank: PhantomData,
},
layout,
placement,
_scalar: PhantomData,
})
}
fn try_typed_tensor_from_buffer_col_major<
T: TensorScalar + Send + Sync + 'static,
R: TensorRank,
>(
shape: impl Into<R::Shape>,
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))?;
let group = match buffer {
StorageBuffer::Host(data) => OwnedTensorGroup::from_host_vec(group_shape, data)?,
StorageBuffer::Backend(buffer) => OwnedTensorGroup::from_backend_buffer(
group_shape,
StorageBuffer::Backend(buffer),
placement.clone(),
)?,
};
Ok(TypedTensor {
group,
layout,
placement,
_scalar: PhantomData,
})
}
impl<T: TensorScalar + Zero, R: TensorRank> TypedTensor<T, R> {
pub fn zeros(shape: impl Into<R::Shape>) -> crate::Result<Self> {
typed_tensor_zeros(shape)
}
}
impl<T: TensorScalar + One + Zero, R: TensorRank> TypedTensor<T, R> {
pub fn ones(shape: impl Into<R::Shape>) -> crate::Result<Self> {
typed_tensor_ones(shape)
}
}
impl<T, R: TensorRank> TypedTensor<T, R> {
pub fn from_buffer_col_major(
shape: impl Into<R::Shape>,
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 Into<R::Shape>,
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) -> crate::Result<TypedTensor<T, Rank<N>>> {
let op = "TypedTensor::try_into_rank";
let actual = self.shape().len();
let shape: [usize; N] = self.shape().try_into().map_err(|_| {
tensor_layout_error(
op,
ValidationError::RankMismatch {
expected: N,
actual,
},
)
})?;
let layout =
TensorLayout::<Rank<N>>::compact(shape).map_err(|err| tensor_layout_error(op, err))?;
let owned = self.group;
Ok(TypedTensor {
group: OwnedTensorGroup {
group: owned.group,
slot: owned.slot,
allocation_index: owned.allocation_index,
host_ptr: owned.host_ptr,
host_byte_len: owned.host_byte_len,
_rank: PhantomData,
},
layout,
placement: self.placement,
_scalar: PhantomData,
})
}
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.layout.shape()
}
pub fn rank(&self) -> usize {
self.shape().len()
}
pub fn layout(&self) -> &TensorLayout<R> {
&self.layout
}
pub fn buffer(&self) -> &StorageBuffer<T>
where
T: 'static,
{
match self
.group
.host_buffer::<T>()
.or_else(|| self.group.backend_buffer::<T>())
{
Some(buffer) => buffer,
None => 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.buffer() {
StorageBuffer::Host(_) => None,
StorageBuffer::Backend(buffer) => Some(buffer.as_ref()),
}
}
#[doc(hidden)]
pub fn backend_buffer_mut(&mut self) -> Option<&mut dyn BackendStorage<T>>
where
T: 'static,
{
let buffer = self.group.backend_buffer_mut::<T>()?;
match buffer {
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,
{
self.group
.prepare_device_read_for_layout::<T>(&self.layout)
.map_err(|error| crate::Error::runtime_state(op, error.to_string()))
}
#[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();
self.group
.prepare_device_write_for_layout::<T>(&layout)
.map_err(|error| crate::Error::runtime_state(op, error.to_string()))
}
pub(crate) fn buffer_len(&self) -> usize
where
T: 'static,
{
self.group
.group
.descriptor_len(self.group.slot)
.unwrap_or_else(|| unreachable!("typed tensor group descriptor mismatch"))
}
pub fn allocation_domain(&self) -> Option<AllocationDomainId>
where
T: 'static,
{
self.group
.group
.backend_identity(self.group.slot)
.map(|(domain, _)| domain)
}
pub fn allocation_id(&self) -> Option<AllocationId>
where
T: 'static,
{
self.group
.group
.backend_identity(self.group.slot)
.map(|(_, allocation)| allocation)
}
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 as_view(&self) -> TypedTensorView<'_, T, R>
where
T: TensorScalar + 'static,
{
let root = match self.group.view::<T>() {
Ok(root) => root,
Err(error) => unreachable!("typed tensor group descriptor mismatch: {error}"),
};
let buffer = if let Some(allocation) = root.backend_allocation() {
TensorStorageRef::Root(allocation)
} else {
TensorStorageRef::Host(self.group_host_slice())
};
let root = Some(root);
TypedTensorView {
buffer,
root,
layout: self.layout.clone(),
placement: self.placement.clone(),
}
}
pub fn as_view_mut(&mut self) -> TypedTensorViewMut<'_, T, R>
where
T: TensorScalar + 'static,
{
let layout = self.layout.clone();
let placement = self.placement.clone();
let mut root = match self.group.view_mut::<T>() {
Ok(root) => root,
Err(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(match root.host_slice_mut() {
Ok(slice) => slice,
Err(error) => {
unreachable!("typed tensor group descriptor is not host-backed: {error}")
}
})
};
TypedTensorViewMut {
buffer,
root: Some(root),
layout,
placement,
}
}
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 root = self.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.into(), strides.into(), 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(),
})
}
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 mut root = self.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.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(TypedTensorViewMut {
buffer: TensorStorageRefMut::Backend(buffer.as_mut()),
root: Some(root),
layout,
placement: self.placement.clone(),
})
}
pub fn into_layout(self) -> TensorLayout<R> {
self.layout
}
pub fn into_parts(self) -> (StorageBuffer<T>, TensorLayout<R>, Placement)
where
T: TensorScalar,
{
let TypedTensor {
group,
layout,
placement,
..
} = self;
let buffer = match group.into_host_vec::<T>() {
Ok(data) => StorageBuffer::Host(data),
Err(_) => StorageBuffer::Host(Vec::new()),
};
(buffer, layout, placement)
}
}
impl<T: TensorScalar, R: TensorRank> TypedTensor<T, R> {
pub fn from_vec_col_major(shape: impl Into<R::Shape>, data: Vec<T>) -> crate::Result<Self> {
typed_tensor_from_vec_col_major(shape, data, "from_vec_col_major")
}
pub fn duplicate(&self) -> crate::Result<Self> {
self.as_view().duplicate()
}
pub fn into_vec_col_major(self) -> crate::Result<(Vec<usize>, Vec<T>)> {
let shape = self.shape().to_vec();
if self.group.backend_buffer::<T>().is_some() {
return Err(crate::Error::runtime_state(
"into_vec_col_major",
"backend buffers cannot be exported as host Vec",
));
}
Ok((shape, self.group.into_host_vec::<T>()?))
}
pub fn into_host_vec(self) -> crate::Result<Vec<T>> {
if self.group.backend_buffer::<T>().is_some() {
return Err(crate::Error::runtime_state(
"into_host_vec",
"backend buffers cannot be exported as host Vec",
));
}
self.group.into_host_vec::<T>()
}
pub fn host_data(&self) -> crate::Result<&[T]> {
self.group.host_slice::<T>()
}
#[doc(hidden)]
pub fn with_host_read<U>(&self, f: impl FnOnce(&[T]) -> U) -> crate::Result<U>
where
T: TensorScalar + 'static,
{
let view = self
.group
.group
.view::<T, R>(self.group.slot)
.map_err(|error| group_error("TypedTensor::with_host_read", error))?;
let prepared = view.prepare_host_read().map_err(|error| {
crate::Error::runtime_state("TypedTensor::with_host_read", error.to_string())
})?;
let slice = prepared.as_slice().ok_or_else(|| {
crate::Error::unsupported(
"TypedTensor::with_host_read",
"host guard access requires a compact descriptor",
)
})?;
Ok(f(slice))
}
#[doc(hidden)]
pub fn with_host_write<U>(&mut self, f: impl FnOnce(&mut [T]) -> U) -> crate::Result<U>
where
T: TensorScalar + 'static,
{
let mut view = self
.group
.group
.view_mut::<T, R>(self.group.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))
}
pub fn as_slice(&self) -> crate::Result<&[T]> {
self.host_data()
}
pub fn host_data_mut(&mut self) -> crate::Result<&mut [T]> {
self.group.host_slice_mut::<T>()
}
fn group_host_slice(&self) -> &[T] {
self.group
.view::<T>()
.ok()
.and_then(|view| view.host_slice().ok())
.unwrap_or_default()
}
fn group_host_slice_mut(&mut self) -> &mut [T] {
self.group
.view_mut::<T>()
.ok()
.and_then(|mut view| view.host_slice_mut().ok())
.unwrap_or_default()
}
pub fn linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
try_linear_offset_for_shape(self.shape(), indices, "TypedTensor::linear_offset")
}
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 {
layout_summary(self.shape(), self.layout.strides(), self.layout.offset())
}
pub fn assert_col_major_contiguous(&self) -> crate::Result<()> {
assert_layout_col_major_contiguous(
self.is_col_major_contiguous()?,
self.shape(),
self.layout.strides(),
self.layout.offset(),
"TypedTensor::assert_col_major_contiguous",
)
}
pub fn get(&self, indices: &[usize]) -> crate::Result<&T> {
let off = self.linear_offset(indices)?;
self.host_data()?.get(off).ok_or_else(|| {
crate::Error::validation("TypedTensor::get", ValidationError::ViewOutOfBounds)
})
}
pub fn get_mut(&mut self, indices: &[usize]) -> crate::Result<&mut T> {
let off = self.linear_offset(indices)?;
self.host_data_mut()?.get_mut(off).ok_or_else(|| {
crate::Error::validation("TypedTensor::get_mut", ValidationError::ViewOutOfBounds)
})
}
}
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 layout = reinterpret_complex_to_real_layout(
self.shape(),
self.layout.strides(),
self.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.group_host_slice_mut(),
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
})
}
pub fn into_real(self) -> Result<TypedTensor<f32, DynRank>, ReinterpretError<Self>> {
let op = "TypedTensor::into_real";
if let Err(error) = validate_representation_pair(op, DType::C32, DType::F32) {
return Err(ReinterpretError::new(self, error));
}
let source_shape = self.shape().to_vec();
let source_strides = self.layout.strides().to_vec();
let source_offset = self.layout.offset();
let target_layout = match reinterpret_complex_to_real_layout(
&source_shape,
&source_strides,
source_offset,
self.buffer_len(),
op,
) {
Ok(layout) => layout,
Err(error) => return Err(ReinterpretError::new(self, error)),
};
let TypedTensor {
group,
layout: source_layout,
placement,
..
} = self;
match group.reinterpret::<Complex32, f32>(
target_layout.shape().to_vec(),
target_layout.strides().to_vec(),
target_layout.offset(),
) {
Ok(group) => Ok(TypedTensor {
group,
layout: target_layout,
placement,
_scalar: PhantomData,
}),
Err((group, error)) => Err(ReinterpretError::new(
TypedTensor {
group,
layout: source_layout,
placement,
_scalar: PhantomData,
},
error,
)),
}
}
}
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 layout = reinterpret_complex_to_real_layout(
self.shape(),
self.layout.strides(),
self.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.group_host_slice_mut(),
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
})
}
pub fn into_real(self) -> Result<TypedTensor<f64, DynRank>, ReinterpretError<Self>> {
let op = "TypedTensor::into_real";
if let Err(error) = validate_representation_pair(op, DType::C64, DType::F64) {
return Err(ReinterpretError::new(self, error));
}
let source_shape = self.shape().to_vec();
let source_strides = self.layout.strides().to_vec();
let source_offset = self.layout.offset();
let target_layout = match reinterpret_complex_to_real_layout(
&source_shape,
&source_strides,
source_offset,
self.buffer_len(),
op,
) {
Ok(layout) => layout,
Err(error) => return Err(ReinterpretError::new(self, error)),
};
let TypedTensor {
group,
layout: source_layout,
placement,
..
} = self;
match group.reinterpret::<Complex64, f64>(
target_layout.shape().to_vec(),
target_layout.strides().to_vec(),
target_layout.offset(),
) {
Ok(group) => Ok(TypedTensor {
group,
layout: target_layout,
placement,
_scalar: PhantomData,
}),
Err((group, error)) => Err(ReinterpretError::new(
TypedTensor {
group,
layout: source_layout,
placement,
_scalar: PhantomData,
},
error,
)),
}
}
}
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 layout = reinterpret_real_to_complex_layout(
self.shape(),
self.layout.strides(),
self.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.group_host_slice_mut(),
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
})
}
pub fn into_complex(self) -> Result<TypedTensor<Complex32, DynRank>, ReinterpretError<Self>> {
let op = "TypedTensor::into_complex";
if let Err(error) = validate_representation_pair(op, DType::F32, DType::C32) {
return Err(ReinterpretError::new(self, error));
}
if !self.buffer_len().is_multiple_of(2) {
return Err(ReinterpretError::new(
self,
crate::Error::invalid_argument(
op,
"buffer",
"the owned real buffer must contain an even number of elements",
),
));
}
let source_shape = self.shape().to_vec();
let source_strides = self.layout.strides().to_vec();
let source_offset = self.layout.offset();
let target_layout = match reinterpret_real_to_complex_layout(
&source_shape,
&source_strides,
source_offset,
self.buffer_len(),
op,
) {
Ok(layout) => layout,
Err(error) => return Err(ReinterpretError::new(self, error)),
};
let TypedTensor {
group,
layout: source_layout,
placement,
..
} = self;
match group.reinterpret::<f32, Complex32>(
target_layout.shape().to_vec(),
target_layout.strides().to_vec(),
target_layout.offset(),
) {
Ok(group) => Ok(TypedTensor {
group,
layout: target_layout,
placement,
_scalar: PhantomData,
}),
Err((group, error)) => Err(ReinterpretError::new(
TypedTensor {
group,
layout: source_layout,
placement,
_scalar: PhantomData,
},
error,
)),
}
}
}
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 layout = reinterpret_real_to_complex_layout(
self.shape(),
self.layout.strides(),
self.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.group_host_slice_mut(),
op,
)?);
Ok(TypedTensorViewMut {
buffer,
root: None,
layout,
placement,
})
}
pub fn into_complex(self) -> Result<TypedTensor<Complex64, DynRank>, ReinterpretError<Self>> {
let op = "TypedTensor::into_complex";
if let Err(error) = validate_representation_pair(op, DType::F64, DType::C64) {
return Err(ReinterpretError::new(self, error));
}
if !self.buffer_len().is_multiple_of(2) {
return Err(ReinterpretError::new(
self,
crate::Error::invalid_argument(
op,
"buffer",
"the owned real buffer must contain an even number of elements",
),
));
}
let source_shape = self.shape().to_vec();
let source_strides = self.layout.strides().to_vec();
let source_offset = self.layout.offset();
let target_layout = match reinterpret_real_to_complex_layout(
&source_shape,
&source_strides,
source_offset,
self.buffer_len(),
op,
) {
Ok(layout) => layout,
Err(error) => return Err(ReinterpretError::new(self, error)),
};
let TypedTensor {
group,
layout: source_layout,
placement,
..
} = self;
match group.reinterpret::<f64, Complex64>(
target_layout.shape().to_vec(),
target_layout.strides().to_vec(),
target_layout.offset(),
) {
Ok(group) => Ok(TypedTensor {
group,
layout: target_layout,
placement,
_scalar: PhantomData,
}),
Err((group, error)) => Err(ReinterpretError::new(
TypedTensor {
group,
layout: source_layout,
placement,
_scalar: PhantomData,
},
error,
)),
}
}
}
impl Tensor {
pub fn as_real_view(&self) -> crate::Result<TensorView<'_>> {
match self {
Tensor::C32(tensor) => tensor.as_real_view().map(TensorView::F32),
Tensor::C64(tensor) => tensor.as_real_view().map(TensorView::F64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_real_view",
other.dtype(),
DType::F32,
"only complex tensors have a sealed real representation view",
)),
}
}
pub fn as_real_view_mut(&mut self) -> crate::Result<TensorViewMut<'_>> {
match self {
Tensor::C32(tensor) => tensor.as_real_view_mut().map(TensorViewMut::F32),
Tensor::C64(tensor) => tensor.as_real_view_mut().map(TensorViewMut::F64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_real_view_mut",
other.dtype(),
DType::F32,
"only complex tensors have a sealed real representation view",
)),
}
}
pub fn into_real(self) -> Result<Self, ReinterpretError<Self>> {
match self {
Tensor::C32(tensor) => tensor.into_real().map(Tensor::F32).map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::C32(owner), error)
}),
Tensor::C64(tensor) => tensor.into_real().map(Tensor::F64).map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::C64(owner), error)
}),
tensor => Err(ReinterpretError::new(
tensor,
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 {
Tensor::F32(tensor) => tensor.as_complex_view().map(TensorView::C32),
Tensor::F64(tensor) => tensor.as_complex_view().map(TensorView::C64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_complex_view",
other.dtype(),
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 {
Tensor::F32(tensor) => tensor.as_complex_view_mut().map(TensorViewMut::C32),
Tensor::F64(tensor) => tensor.as_complex_view_mut().map(TensorViewMut::C64),
other => Err(crate::Error::unsupported_dtype_conversion(
"Tensor::as_complex_view_mut",
other.dtype(),
DType::C32,
"only real tensors can have a sealed complex representation view",
)),
}
}
pub fn into_complex(self) -> Result<Self, ReinterpretError<Self>> {
match self {
Tensor::F32(tensor) => tensor.into_complex().map(Tensor::C32).map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::F32(owner), error)
}),
Tensor::F64(tensor) => tensor.into_complex().map(Tensor::C64).map_err(|error| {
let (owner, error) = error.into_parts();
ReinterpretError::new(Tensor::F64(owner), error)
}),
tensor => Err(ReinterpretError::new(
tensor,
crate::Error::unsupported(
"Tensor::into_complex",
"only real tensors have a sealed complex representation",
),
)),
}
}
pub fn duplicate(&self) -> crate::Result<Self> {
match self {
Tensor::F32(t) => t.duplicate().map(Tensor::F32),
Tensor::F64(t) => t.duplicate().map(Tensor::F64),
Tensor::I32(t) => t.duplicate().map(Tensor::I32),
Tensor::I64(t) => t.duplicate().map(Tensor::I64),
Tensor::Bool(t) => t.duplicate().map(Tensor::Bool),
Tensor::C32(t) => t.duplicate().map(Tensor::C32),
Tensor::C64(t) => t.duplicate().map(Tensor::C64),
}
}
pub fn from_vec_col_major<T: TensorScalar>(
shape: impl tenferro_tensor_core::IntoShapeVec,
data: Vec<T>,
) -> crate::Result<Self> {
T::into_tensor(shape.into_shape_vec().to_vec(), data)
}
pub fn shape(&self) -> &[usize] {
match self {
Tensor::F32(t) => t.shape(),
Tensor::F64(t) => t.shape(),
Tensor::I32(t) => t.shape(),
Tensor::I64(t) => t.shape(),
Tensor::Bool(t) => t.shape(),
Tensor::C32(t) => t.shape(),
Tensor::C64(t) => t.shape(),
}
}
pub fn dtype(&self) -> DType {
match self {
Tensor::F32(_) => DType::F32,
Tensor::F64(_) => DType::F64,
Tensor::I32(_) => DType::I32,
Tensor::I64(_) => DType::I64,
Tensor::Bool(_) => DType::Bool,
Tensor::C32(_) => DType::C32,
Tensor::C64(_) => DType::C64,
}
}
pub fn placement(&self) -> &Placement {
match self {
Tensor::F32(t) => t.placement(),
Tensor::F64(t) => t.placement(),
Tensor::I32(t) => t.placement(),
Tensor::I64(t) => t.placement(),
Tensor::Bool(t) => t.placement(),
Tensor::C32(t) => t.placement(),
Tensor::C64(t) => t.placement(),
}
}
pub fn is_backend_buffer(&self) -> bool {
match self {
Tensor::F32(t) => t.backend_family().is_some(),
Tensor::F64(t) => t.backend_family().is_some(),
Tensor::I32(t) => t.backend_family().is_some(),
Tensor::I64(t) => t.backend_family().is_some(),
Tensor::Bool(t) => t.backend_family().is_some(),
Tensor::C32(t) => t.backend_family().is_some(),
Tensor::C64(t) => t.backend_family().is_some(),
}
}
pub fn layout_linear_offset(&self, indices: &[usize]) -> crate::Result<usize> {
match self {
Tensor::F32(t) => t.layout_linear_offset(indices),
Tensor::F64(t) => t.layout_linear_offset(indices),
Tensor::I32(t) => t.layout_linear_offset(indices),
Tensor::I64(t) => t.layout_linear_offset(indices),
Tensor::Bool(t) => t.layout_linear_offset(indices),
Tensor::C32(t) => t.layout_linear_offset(indices),
Tensor::C64(t) => t.layout_linear_offset(indices),
}
}
pub fn is_col_major_contiguous(&self) -> crate::Result<bool> {
match self {
Tensor::F32(t) => t.is_col_major_contiguous(),
Tensor::F64(t) => t.is_col_major_contiguous(),
Tensor::I32(t) => t.is_col_major_contiguous(),
Tensor::I64(t) => t.is_col_major_contiguous(),
Tensor::Bool(t) => t.is_col_major_contiguous(),
Tensor::C32(t) => t.is_col_major_contiguous(),
Tensor::C64(t) => t.is_col_major_contiguous(),
}
}
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)
}
pub fn into_vec_col_major<T: TensorScalar>(self) -> crate::Result<(Vec<usize>, Vec<T>)> {
let typed = T::into_typed(self)?;
typed.into_vec_col_major()
}
}
#[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];
}
}
}