#![cfg_attr(
all(test, feature = "provider-inject"),
allow(dead_code, unused_imports)
)]
macro_rules! preset_scalar {
(F32) => {
f32
};
(F64) => {
f64
};
(I32) => {
i32
};
(I64) => {
i64
};
(Bool) => {
bool
};
(C32) => {
num_complex::Complex32
};
(C64) => {
num_complex::Complex64
};
}
#[cfg(not(any(feature = "cpu-faer", feature = "cpu-blas")))]
compile_error!("enable at least one CPU backend: cpu-faer or cpu-blas");
#[cfg(all(feature = "provider-inject", not(feature = "cpu-blas")))]
compile_error!("provider-inject requires cpu-blas");
#[cfg(any(
all(feature = "blas-openblas", feature = "blas-accelerate"),
all(feature = "blas-openblas", feature = "blas-mkl"),
all(feature = "blas-accelerate", feature = "blas-mkl"),
))]
compile_error!(
"enable at most one explicit BLAS provider feature: blas-openblas, blas-accelerate, or blas-mkl"
);
#[cfg(all(
feature = "provider-inject",
any(
feature = "blas-openblas",
feature = "blas-accelerate",
feature = "blas-mkl"
)
))]
compile_error!("provider-inject cannot be combined with explicit BLAS provider features");
pub mod affinity;
mod affinity_policy;
mod analytic;
mod arbiter;
pub mod backend;
mod blas1;
pub(crate) mod buffer_pool {
pub use tenferro_cpu_basic::buffer_pool::*;
}
mod capability;
pub mod context;
#[allow(dead_code)]
mod domain_executor;
#[allow(dead_code)]
mod dot_runtime;
pub(crate) use tenferro_cpu_basic::PooledUninitOutput;
pub(crate) use tenferro_cpu_basic::{erased_raw_strided_ref, erased_raw_strided_uninit_mut};
pub(crate) use tenferro_internal_cpu_kernels::elementwise;
mod engine;
mod exec_session;
mod gemm;
mod indexed_plan_cache;
mod indexing;
#[cfg(feature = "provider-inject")]
pub mod inject;
mod placement;
pub mod provider;
mod provider_capability;
mod reduction;
mod resource_domain;
mod runtime_adapter;
mod structural;
mod topology;
use num_complex::{Complex32, Complex64};
use std::ptr::NonNull;
#[cfg(test)]
use strided_kernel::col_major_strides as kernel_col_major_strides;
#[cfg(test)]
use strided_kernel::StridedArray;
use crate::buffer_pool::BufferPool;
pub(crate) use tenferro_tensor::*;
pub(crate) fn cpu_contraction_unsupported_dtype_message(dtype: DType) -> String {
let remedy = matches!(dtype, DType::I32 | DType::I64)
.then_some(format!("; convert {dtype:?} to F64 before contraction"));
format!(
"CPU contraction providers support F32/F64/C32/C64{}",
remedy.unwrap_or_default()
)
}
pub(crate) fn erased_raw_strided_mut<'a>(
dtype: strided_kernel::KernelDType,
data: &'a mut [u8],
dims: &'a [usize],
strides: &'a [isize],
offset: isize,
) -> strided_kernel::Result<strided_kernel::ErasedRawStridedMut<'a>> {
let data_ptr = NonNull::new(data.as_mut_ptr()).unwrap_or_else(NonNull::dangling);
unsafe {
strided_kernel::ErasedRawStridedMut::from_raw_parts(
dtype,
data_ptr,
data.len(),
dims,
strides,
offset,
)
}
}
#[cfg(feature = "provider-src")]
extern crate blas_src as _;
#[cfg(feature = "provider-inject")]
extern crate cblas_inject as _;
#[cfg(feature = "provider-inject")]
extern crate lapack_inject as _;
#[cfg(feature = "provider-src")]
extern crate lapack_src as _;
pub use affinity::{
available_parallelism, process_cpu_affinity, process_cpu_affinity_count, CpuAffinityError,
};
pub use affinity_policy::{
resolve_cpu_affinity, resolve_cpu_affinity_with_override, CpuAffinityInput,
CpuAffinityInputError, CpuAffinityPolicy, CpuAffinityResolutionError, CpuAffinitySelection,
CpuAffinitySelectionReason,
};
pub use backend::{
CpuBackend, CpuBackendError, CpuBackendKind, CpuExecutionInfo, CpuExecutionMode,
CpuRuntimeIdentity, ExternalCpuDomainRegistryError,
};
pub use buffer_pool::BufferPoolStats;
pub use capability::cpu_capabilities;
pub use context::{CpuContext, CpuContextError, DEFAULT_WORKER_STACK_BYTES};
pub use domain_executor::{
CpuDomainExecutor, CpuDomainExecutorCapabilities, CpuDomainExecutorError, CpuExecutorAffinity,
CpuExecutorReentrancy, CpuExecutorShutdown, CpuInnerParallelism, RayonCpuDomainExecutor,
ScopedCpuJob, ScopedCpuJobs,
};
pub use dot_runtime::{
CpuProviderBundle, CpuProviderBundleBuildError, CpuProviderBundleBuilder,
CpuProviderBundleInstallError, CpuProviderSlot, GeneralContractionPolicy,
};
#[doc(hidden)]
pub use exec_session::CpuExecSession;
pub use indexed_plan_cache::IndexedPlanCacheLimits;
pub use placement::{
CpuEngineConstructionError, CpuPlacement, CpuPlacementError, CpuPlacementGuarantee,
ResolvedCpuPlacement,
};
pub use provider::{CpuExecutionContext, ParallelMode};
pub use provider_capability::{
CpuPlacementControl, CpuProviderDomainError, CpuProviderExecutionCapabilities,
CpuThreadCountControl,
};
pub use resource_domain::{
CpuAdmissionMode, CpuDomainOwnership, ExternalCpuDomain, ExternalCpuDomainError,
};
pub use runtime_adapter::{
runtime_engine_id, runtime_engine_registration, runtime_engine_registration_with_id,
runtime_hardware_class,
};
pub use tenferro_internal_cpu_kernels::scalar_ops::{
scalar_binary_into, scalar_fold, AddOp, BinaryScalarOp, MulOp, SubOp,
};
pub use tenferro_internal_cpu_kernels::{same_variant_pair, same_variant_unary};
pub use topology::{
discover_cpu_topology, CpuId, CpuNode, CpuSet, CpuSetError, CpuTopology, CpuTopologyError,
NumaNodeId,
};
#[doc(hidden)]
pub fn with_cpu_exec_session<B, R>(
session: &mut B,
f: impl for<'a> FnOnce(&'a mut CpuExecSession<'a>) -> R,
) -> Option<R>
where
B: tenferro_tensor::BackendSession + ?Sized,
{
if session.session_type_id() != std::any::TypeId::of::<exec_session::CpuExecSessionMarker>() {
return None;
}
let data = unsafe { session.session_data_mut() };
Some(unsafe { f(&mut *(data.cast::<CpuExecSession<'static>>())) })
}
#[cfg(feature = "cpu-faer")]
#[cfg_attr(docsrs, doc(cfg(feature = "cpu-faer")))]
pub trait FaerParallelismExt {
fn with_faer_parallelism(
&mut self,
callback: impl FnOnce(faer::Par) -> tenferro_tensor::Result<()> + Send,
) -> tenferro_tensor::Result<()>;
}
#[cfg(feature = "cpu-faer")]
impl<S> FaerParallelismExt for S
where
S: tenferro_tensor::BackendSession + ?Sized,
{
fn with_faer_parallelism(
&mut self,
callback: impl FnOnce(faer::Par) -> tenferro_tensor::Result<()> + Send,
) -> tenferro_tensor::Result<()> {
with_cpu_exec_session(self, |session| session.with_faer_parallelism(callback))
.unwrap_or_else(|| {
Err(tenferro_tensor::Error::unsupported(
"with_faer_parallelism",
"selected session is not a CPU/faer execution session",
))
})
}
}
#[cfg(test)]
pub(crate) use analytic::pow;
#[cfg(test)]
macro_rules! test_elementwise_wrapper {
($name:ident($($arg:ident: $ty:ty),*) => $with_pool:ident) => {
pub(crate) fn $name($($arg: $ty),*) -> crate::Result<Tensor> {
let mut buffers = BufferPool::new();
elementwise::$with_pool(&mut buffers, &strided_kernel::ExecContext::serial(), $($arg),*)
}
};
}
#[cfg(test)]
test_elementwise_wrapper!(abs(input: &Tensor) => abs_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(add(lhs: &Tensor, rhs: &Tensor) => add_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) => clamp_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) => compare_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(conj(input: &Tensor) => conj_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(div(lhs: &Tensor, rhs: &Tensor) => div_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(maximum(lhs: &Tensor, rhs: &Tensor) => maximum_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(minimum(lhs: &Tensor, rhs: &Tensor) => minimum_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(mul(lhs: &Tensor, rhs: &Tensor) => mul_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(neg(input: &Tensor) => neg_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(rem(lhs: &Tensor, rhs: &Tensor) => rem_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) => select_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(sign(input: &Tensor) => sign_with_pool);
#[cfg(test)]
test_elementwise_wrapper!(sub(lhs: &Tensor, rhs: &Tensor) => sub_with_pool);
#[cfg(test)]
pub(crate) use indexing::{dynamic_slice, dynamic_update_slice, gather, pad, scatter};
#[cfg(test)]
pub(crate) use reduction::{reduce_max, reduce_min, reduce_prod, reduce_sum, reduce_sum_squares};
#[cfg(test)]
pub(crate) use structural::{
broadcast_in_dim, embed_diagonal, extract_diagonal, reshape, transpose, tril, triu,
};
#[doc(hidden)]
pub mod linalg_interop {
pub use crate::buffer_pool::{BufferPool, PoolScalar};
pub use tenferro_cpu_basic::PooledUninitOutput;
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum CpuNumericalError {
#[error("{op} received a negative integer exponent for dtype {dtype:?}")]
NegativeIntegerExponent { op: &'static str, dtype: DType },
}
pub(crate) fn cpu_negative_integer_exponent(op: &'static str, dtype: DType) -> crate::Error {
crate::Error::extension(
op,
"cpu",
ErrorKind::NumericalFailure,
CpuNumericalError::NegativeIntegerExponent { op, dtype },
)
}
pub(crate) use tenferro_cpu_basic::{
cpu_backend_buffer_error, typed_host_data, typed_view, typed_view_from_view, ConjElem,
};
pub(crate) fn materialize_tensor_read_in_domain(
buffers: &mut BufferPool,
op: &'static str,
input: TensorRead<'_>,
domain: Option<&dyn SharedTensorAllocationDomain>,
) -> crate::Result<Tensor> {
if let Some(domain) = domain {
if input.backend_family().is_some() {
return materialize_managed_read(op, input, domain);
}
}
materialize_tensor_read(buffers, op, input)
}
pub(crate) fn materialize_tensor_read(
buffers: &mut BufferPool,
op: &'static str,
input: TensorRead<'_>,
) -> crate::Result<Tensor> {
match input {
TensorRead::Tensor(tensor) => clone_host_tensor_read(op, tensor),
TensorRead::View(view) => materialize_tensor_view(buffers, op, view),
}
}
fn materialize_managed_read(
op: &'static str,
input: TensorRead<'_>,
domain: &dyn SharedTensorAllocationDomain,
) -> crate::Result<Tensor> {
if input.placement().memory_kind != MemoryKind::Managed {
return Err(Error::host_access(
op,
HostAccessError::Unsupported { backend: "backend" },
));
}
match input.allocation_domain() {
Some(actual) if actual == domain.id() => {}
Some(actual) => {
return Err(Error::host_access(
op,
HostAccessError::ForeignDomain {
expected: domain.id(),
actual,
},
))
}
None => {
return Err(Error::host_access(
op,
HostAccessError::Unsupported { backend: "backend" },
))
}
}
fn copy<T: TensorScalar>(
op: &'static str,
input: TypedTensorView<'_, T>,
domain: &dyn SharedTensorAllocationDomain,
) -> crate::Result<Tensor> {
input.with_host_read(|source| {
let mut output = domain.allocate(T::dtype(), input.shape())?;
if output.shape() != input.shape()
|| output.placement().memory_kind != MemoryKind::Managed
|| TensorRead::from_tensor(&output).allocation_domain() != Some(domain.id())
{
return Err(Error::runtime_state(
op,
"shared allocator returned incompatible output",
));
}
let typed = output.as_typed_mut::<T>().ok_or_else(|| {
Error::runtime_state(op, "shared allocator returned the wrong dtype")
})?;
if let Some(buffer) = typed.backend_buffer_mut() {
buffer
.map_write()
.map_err(|error| Error::host_access(op, error))?
.copy_from_slice(source)
.map_err(|error| Error::host_access(op, error))?;
} else {
typed.with_host_write(|target| target.copy_from_slice(source))?;
}
Ok(output)
})?
}
match input.tensor_view() {
TensorView::F32(view) => copy(op, view, domain),
TensorView::F64(view) => copy(op, view, domain),
TensorView::I32(view) => copy(op, view, domain),
TensorView::I64(view) => copy(op, view, domain),
TensorView::Bool(view) => copy(op, view, domain),
TensorView::C32(view) => copy(op, view, domain),
TensorView::C64(view) => copy(op, view, domain),
}
}
pub(crate) fn copy_tensor_read_into(
op: &'static str,
src: TensorRead<'_>,
dst: TensorWrite<'_>,
) -> crate::Result<()> {
let src_dtype = src.dtype();
let dst_dtype = dst.dtype();
macro_rules! copy_source {
($variant:ident, $src:expr) => {{
let src = $src;
match dst {
TensorWrite::Tensor(dst)
if dst.dtype()
== <preset_scalar!($variant) as tenferro_tensor::TensorScalar>::dtype() =>
{
let dst = dst
.as_typed_mut::<preset_scalar!($variant)>()
.expect("the dtype guard selects this arm");
let mut dst = dst.as_view_mut();
structural::typed_copy_view_into(&src, &mut dst, op)
}
TensorWrite::View(TensorViewMut::$variant(mut dst)) => {
structural::typed_copy_view_into(&src, &mut dst, op)
}
_ => Err(crate::Error::dtype_mismatch(op, src_dtype, dst_dtype)),
}
}};
}
fn read_refusal(tensor: &Tensor) -> crate::Error {
crate::Error::unsupported_dtype(
"copy_tensor_read_into",
tensor.dtype(),
"an externally defined payload is not a runtime read",
)
}
match src {
TensorRead::Tensor(tensor) => match tensor.dtype() {
DType::F32 => copy_source!(
F32,
tensor
.as_typed::<f32>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::F64 => copy_source!(
F64,
tensor
.as_typed::<f64>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::I32 => copy_source!(
I32,
tensor
.as_typed::<i32>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::I64 => copy_source!(
I64,
tensor
.as_typed::<i64>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::Bool => copy_source!(
Bool,
tensor
.as_typed::<bool>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::C32 => copy_source!(
C32,
tensor
.as_typed::<Complex32>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::C64 => copy_source!(
C64,
tensor
.as_typed::<Complex64>()
.ok_or_else(|| read_refusal(tensor))?
.as_view()
),
DType::External(_) => Err(read_refusal(tensor)),
},
TensorRead::View(TensorView::F32(src)) => copy_source!(F32, src),
TensorRead::View(TensorView::F64(src)) => copy_source!(F64, src),
TensorRead::View(TensorView::I32(src)) => copy_source!(I32, src),
TensorRead::View(TensorView::I64(src)) => copy_source!(I64, src),
TensorRead::View(TensorView::Bool(src)) => copy_source!(Bool, src),
TensorRead::View(TensorView::C32(src)) => copy_source!(C32, src),
TensorRead::View(TensorView::C64(src)) => copy_source!(C64, src),
}
}
fn clone_host_tensor_read(op: &'static str, tensor: &Tensor) -> crate::Result<Tensor> {
macro_rules! clone_host {
($variant:ident, $tensor:expr) => {{
structural::validate_cpu_host_placement(op, "source", $tensor.placement())?;
typed_host_data(op, $tensor)?;
$tensor
.duplicate()
.map(Tensor::from_typed::<preset_scalar!($variant)>)
}};
}
match tensor.dtype() {
DType::F32 => {
let tensor = host_typed::<f32>(op, tensor)?;
clone_host!(F32, tensor)
}
DType::F64 => {
let tensor = host_typed::<f64>(op, tensor)?;
clone_host!(F64, tensor)
}
DType::I32 => {
let tensor = host_typed::<i32>(op, tensor)?;
clone_host!(I32, tensor)
}
DType::I64 => {
let tensor = host_typed::<i64>(op, tensor)?;
clone_host!(I64, tensor)
}
DType::Bool => {
let tensor = host_typed::<bool>(op, tensor)?;
clone_host!(Bool, tensor)
}
DType::C32 => {
let tensor = host_typed::<Complex32>(op, tensor)?;
clone_host!(C32, tensor)
}
DType::C64 => {
let tensor = host_typed::<Complex64>(op, tensor)?;
clone_host!(C64, tensor)
}
DType::External(_) => Ok(Tensor::external_with_placement(
tensor
.external_payload()
.ok_or_else(|| host_typed_error(op, tensor))?
.duplicate(),
tensor.placement().clone(),
)),
}
}
fn host_typed<'a, T: TensorScalar>(
op: &'static str,
tensor: &'a Tensor,
) -> crate::Result<&'a TypedTensor<T>> {
tensor
.as_typed::<T>()
.ok_or_else(|| host_typed_error(op, tensor))
}
fn host_typed_error(op: &'static str, tensor: &Tensor) -> crate::Error {
crate::Error::unsupported_dtype(
op,
tensor.dtype(),
"the CPU host clone requires a preset scalar",
)
}
fn materialize_tensor_view(
buffers: &mut BufferPool,
op: &'static str,
view: TensorView<'_>,
) -> crate::Result<Tensor> {
macro_rules! materialize {
($variant:ident, $view:expr) => {{
Ok(Tensor::from_typed::<preset_scalar!($variant)>(
structural::typed_materialize_view_with_pool(buffers, &$view, op)?,
))
}};
}
match view {
TensorView::F32(view) => materialize!(F32, view),
TensorView::F64(view) => materialize!(F64, view),
TensorView::I32(view) => materialize!(I32, view),
TensorView::I64(view) => materialize!(I64, view),
TensorView::Bool(view) => materialize!(Bool, view),
TensorView::C32(view) => materialize!(C32, view),
TensorView::C64(view) => materialize!(C64, view),
}
}
#[allow(clippy::uninit_vec)]
#[cfg(test)]
pub(crate) unsafe fn typed_array_uninit<T>(shape: &[usize]) -> StridedArray<T> {
let total: usize = shape.iter().product();
let strides = kernel_col_major_strides(shape);
let mut data = Vec::with_capacity(total);
unsafe { data.set_len(total) };
StridedArray::from_parts(data, shape, &strides, 0).expect("column-major output array")
}
#[cfg(test)]
pub(crate) fn tensor_from_array<T: Clone + tenferro_tensor::TensorScalar>(
array: StridedArray<T>,
) -> TypedTensor<T> {
TypedTensor::from_vec_col_major(array.dims().to_vec(), array.into_data())
.expect("strided array dimensions match owned data length")
}
#[cfg(all(test, not(feature = "provider-inject")))]
mod tests;