use crate::config::{
CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
};
use crate::types::{
TensorRank, TensorScalar, TensorView, TypedTensor, TypedTensorView, TypedTensorViewMut,
};
use crate::validate::validate_convert_dtype;
use crate::{
AllocationDomainId, AllocationId, DType, Error, ShapeMismatch, ShapeVec, Tensor, TensorRead,
TensorValue, TensorWrite, ValidationError,
};
use num_complex::{Complex32, Complex64};
#[cfg(test)]
mod tests;
fn read_boundary_error(op: &'static str) -> crate::Error {
crate::Error::unsupported(
op,
"backend does not accept borrowed tensor views at this execution boundary",
)
}
fn validation(op: &'static str, source: ValidationError) -> crate::Error {
Error::validation(op, source)
}
fn default_materialize_host_view(view: TensorView<'_>) -> crate::Result<Tensor> {
fn typed<T: TensorScalar>(view: &TypedTensorView<'_, T>) -> crate::Result<Tensor> {
view.to_col_major().map(Tensor::from_typed)
}
match view {
TensorView::F32(view) => typed(&view),
TensorView::F64(view) => typed(&view),
TensorView::I32(view) => typed(&view),
TensorView::I64(view) => typed(&view),
TensorView::Bool(view) => typed(&view),
TensorView::C32(view) => typed(&view),
TensorView::C64(view) => typed(&view),
}
}
fn invalid_argument(op: &'static str, argument: &'static str, message: impl Into<String>) -> Error {
Error::invalid_argument(op, argument, message)
}
fn read_tensor<'a>(op: &'static str, input: TensorRead<'a>) -> crate::Result<&'a Tensor> {
input.as_tensor().ok_or_else(|| read_boundary_error(op))
}
#[doc(hidden)]
pub fn read_owned_tensor<'a>(op: &'static str, input: TensorRead<'a>) -> crate::Result<&'a Tensor> {
read_tensor(op, input)
}
fn validate_axis_list(
op: &'static str,
role: &'static str,
axes: &[usize],
rank: usize,
) -> crate::Result<()> {
for (index, &axis) in axes.iter().enumerate() {
if axis >= rank {
return Err(validation(
op,
ValidationError::AxisOutOfBounds { axis, rank },
));
}
if axes[..index].contains(&axis) {
return Err(validation(
op,
ValidationError::DuplicateAxis { axis, role },
));
}
}
Ok(())
}
fn validate_role_disjoint(
op: &'static str,
first_role: &'static str,
first_axes: &[usize],
second_role: &'static str,
second_axes: &[usize],
) -> crate::Result<()> {
for &axis in first_axes {
if second_axes.contains(&axis) {
return Err(validation(
op,
ValidationError::AxisRoleConflict {
axis,
first_role,
second_role,
},
));
}
}
Ok(())
}
#[doc(hidden)]
pub fn dot_general_output_shape(
lhs_shape: &[usize],
rhs_shape: &[usize],
config: &DotGeneralConfig,
op: &'static str,
) -> crate::Result<ShapeVec> {
if config.lhs_contracting_dims.len() != config.rhs_contracting_dims.len() {
return Err(invalid_argument(
op,
"contracting_dims",
"lhs/rhs contracting dim counts differ",
));
}
if config.lhs_batch_dims.len() != config.rhs_batch_dims.len() {
return Err(invalid_argument(
op,
"batch_dims",
"lhs/rhs batch dim counts differ",
));
}
let lhs_rank = lhs_shape.len();
let rhs_rank = rhs_shape.len();
validate_axis_list(
op,
"lhs_contracting",
&config.lhs_contracting_dims,
lhs_rank,
)?;
validate_axis_list(
op,
"rhs_contracting",
&config.rhs_contracting_dims,
rhs_rank,
)?;
validate_axis_list(op, "lhs_batch", &config.lhs_batch_dims, lhs_rank)?;
validate_axis_list(op, "rhs_batch", &config.rhs_batch_dims, rhs_rank)?;
validate_role_disjoint(
op,
"lhs_contracting",
&config.lhs_contracting_dims,
"lhs_batch",
&config.lhs_batch_dims,
)?;
validate_role_disjoint(
op,
"rhs_contracting",
&config.rhs_contracting_dims,
"rhs_batch",
&config.rhs_batch_dims,
)?;
for (&lhs_axis, &rhs_axis) in config
.lhs_contracting_dims
.iter()
.zip(&config.rhs_contracting_dims)
{
if lhs_shape[lhs_axis] != rhs_shape[rhs_axis] {
return Err(validation(
op,
ShapeMismatch::ContractedDimensions {
lhs_axis,
lhs_size: lhs_shape[lhs_axis],
rhs_axis,
rhs_size: rhs_shape[rhs_axis],
}
.into(),
));
}
}
for (&lhs_axis, &rhs_axis) in config.lhs_batch_dims.iter().zip(&config.rhs_batch_dims) {
if lhs_shape[lhs_axis] != rhs_shape[rhs_axis] {
return Err(validation(
op,
ShapeMismatch::ContractedDimensions {
lhs_axis,
lhs_size: lhs_shape[lhs_axis],
rhs_axis,
rhs_size: rhs_shape[rhs_axis],
}
.into(),
));
}
}
let lhs_free = (0..lhs_rank)
.filter(|axis| {
!config.lhs_contracting_dims.contains(axis) && !config.lhs_batch_dims.contains(axis)
})
.map(|axis| lhs_shape[axis]);
let rhs_free = (0..rhs_rank)
.filter(|axis| {
!config.rhs_contracting_dims.contains(axis) && !config.rhs_batch_dims.contains(axis)
})
.map(|axis| rhs_shape[axis]);
let batch = config.lhs_batch_dims.iter().map(|&axis| lhs_shape[axis]);
Ok(lhs_free.chain(rhs_free).chain(batch).collect())
}
#[doc(hidden)]
pub fn validate_dot_general_read_into(
lhs: &TensorRead<'_>,
rhs: &TensorRead<'_>,
config: &DotGeneralConfig,
out: &TensorWrite<'_>,
op: &'static str,
) -> crate::Result<ShapeVec> {
if lhs.dtype() != rhs.dtype() {
return Err(validation(
op,
ValidationError::DTypeMismatch {
expected: lhs.dtype(),
actual: rhs.dtype(),
},
));
}
if lhs.dtype() != out.dtype() {
return Err(validation(
op,
ValidationError::DTypeMismatch {
expected: lhs.dtype(),
actual: out.dtype(),
},
));
}
let expected = dot_general_output_shape(lhs.shape(), rhs.shape(), config, op)?;
if out.shape() != expected.as_slice() {
return Err(validation(
op,
ShapeMismatch::ExpectedActual {
expected: expected.clone(),
actual: out.shape().iter().copied().collect(),
}
.into(),
));
}
Ok(expected)
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ContractionScalar {
F32(f32),
F64(f64),
C32(Complex32),
C64(Complex64),
}
impl ContractionScalar {
pub fn dtype(self) -> DType {
match self {
Self::F32(_) => DType::F32,
Self::F64(_) => DType::F64,
Self::C32(_) => DType::C32,
Self::C64(_) => DType::C64,
}
}
pub fn one(dtype: DType) -> crate::Result<Self> {
match dtype {
DType::F32 => Ok(Self::F32(1.0)),
DType::F64 => Ok(Self::F64(1.0)),
DType::C32 => Ok(Self::C32(Complex32::new(1.0, 0.0))),
DType::C64 => Ok(Self::C64(Complex64::new(1.0, 0.0))),
DType::I32 | DType::I64 | DType::Bool | DType::External(_) => Err(validation(
"ContractionScalar::one",
ValidationError::DTypeMismatch {
expected: dtype,
actual: DType::F32,
},
)),
}
}
pub fn zero(dtype: DType) -> crate::Result<Self> {
match dtype {
DType::F32 => Ok(Self::F32(0.0)),
DType::F64 => Ok(Self::F64(0.0)),
DType::C32 => Ok(Self::C32(Complex32::new(0.0, 0.0))),
DType::C64 => Ok(Self::C64(Complex64::new(0.0, 0.0))),
DType::I32 | DType::I64 | DType::Bool | DType::External(_) => Err(validation(
"ContractionScalar::zero",
ValidationError::DTypeMismatch {
expected: dtype,
actual: DType::F32,
},
)),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DotGeneralAccumulation {
pub lhs_conj: bool,
pub rhs_conj: bool,
pub alpha: ContractionScalar,
pub beta: ContractionScalar,
}
impl DotGeneralAccumulation {
fn identity(
op: &'static str,
dtype: DType,
multiplicative: bool,
) -> crate::Result<ContractionScalar> {
let result = if multiplicative {
ContractionScalar::one(dtype)
} else {
ContractionScalar::zero(dtype)
};
result.map_err(|error| match error {
Error::Validation { source, .. } => validation(op, source),
error => error,
})
}
pub fn overwrite(dtype: DType) -> crate::Result<Self> {
Ok(Self {
lhs_conj: false,
rhs_conj: false,
alpha: Self::identity("DotGeneralAccumulation::overwrite", dtype, true)?,
beta: Self::identity("DotGeneralAccumulation::overwrite", dtype, false)?,
})
}
pub fn add_to(dtype: DType) -> crate::Result<Self> {
Ok(Self {
lhs_conj: false,
rhs_conj: false,
alpha: Self::identity("DotGeneralAccumulation::add_to", dtype, true)?,
beta: Self::identity("DotGeneralAccumulation::add_to", dtype, true)?,
})
}
pub fn scaled(alpha: ContractionScalar, beta: ContractionScalar) -> crate::Result<Self> {
if alpha.dtype() != beta.dtype() {
return Err(validation(
"DotGeneralAccumulation::scaled",
ValidationError::DTypeMismatch {
expected: alpha.dtype(),
actual: beta.dtype(),
},
));
}
Ok(Self {
lhs_conj: false,
rhs_conj: false,
alpha,
beta,
})
}
fn validate_for_dtype(self, dtype: DType) -> crate::Result<()> {
for scalar in [self.alpha, self.beta] {
if scalar.dtype() != dtype {
return Err(validation(
"dot_general",
ValidationError::DTypeMismatch {
expected: scalar.dtype(),
actual: dtype,
},
));
}
}
Ok(())
}
}
#[doc(hidden)]
pub fn validate_dot_general_accumulation(
lhs: &TensorRead<'_>,
rhs: &TensorRead<'_>,
config: &DotGeneralConfig,
accumulation: DotGeneralAccumulation,
out: &TensorWrite<'_>,
op: &'static str,
) -> crate::Result<ShapeVec> {
let shape = validate_dot_general_read_into(lhs, rhs, config, out, op)?;
accumulation.validate_for_dtype(lhs.dtype())?;
Ok(shape)
}
#[doc(hidden)]
pub fn dot_general_accum_via_temp<B: TensorDot + ?Sized>(
backend: &mut B,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
accumulation: DotGeneralAccumulation,
mut out: TensorWrite<'_>,
) -> crate::Result<()> {
validate_dot_general_accumulation(&lhs, &rhs, config, accumulation, &out, "dot_general")?;
let dot = backend.dot_general_with_conj_read(
lhs,
rhs,
config,
accumulation.lhs_conj,
accumulation.rhs_conj,
)?;
accumulate_dot_result_into(&dot, accumulation, &mut out)
}
#[doc(hidden)]
pub fn accumulate_dot_result_into(
dot: &Tensor,
accumulation: DotGeneralAccumulation,
out: &mut TensorWrite<'_>,
) -> crate::Result<()> {
macro_rules! dispatch {
($variant:ident, $ty:ty) => {
if dot.dtype() == <$ty as crate::TensorScalar>::dtype() {
let (ContractionScalar::$variant(alpha), ContractionScalar::$variant(beta)) =
(accumulation.alpha, accumulation.beta)
else {
return Err(validation(
"dot_general",
ValidationError::DTypeMismatch {
expected: dot.dtype(),
actual: accumulation.alpha.dtype(),
},
));
};
let dot = dot
.as_typed::<$ty>()
.expect("the dtype guard selects this arm");
match out {
TensorWrite::Tensor(out) => {
let out = out
.as_typed_mut::<$ty>()
.expect("the dtype guard selects this arm");
let mut out = out.as_view_mut();
accumulate_typed(dot.as_slice()?, alpha, beta, &mut out)?;
return Ok(());
}
TensorWrite::View(crate::TensorViewMut::$variant(out)) => {
accumulate_typed(dot.as_slice()?, alpha, beta, out)?;
return Ok(());
}
_ => {}
}
}
};
}
dispatch!(F32, f32);
dispatch!(F64, f64);
dispatch!(C32, Complex32);
dispatch!(C64, Complex64);
Err(validation(
"dot_general",
ValidationError::DTypeMismatch {
expected: accumulation.alpha.dtype(),
actual: dot.dtype(),
},
))
}
fn accumulate_typed<T>(
dot: &[T],
alpha: T,
beta: T,
out: &mut TypedTensorViewMut<'_, T>,
) -> crate::Result<()>
where
T: Copy
+ PartialEq
+ std::ops::Add<Output = T>
+ std::ops::Mul<Output = T>
+ num_traits::Zero
+ 'static,
{
let beta_is_zero = beta == T::zero();
if let Some(output) = compact_host_accumulation_slice(out, dot.len())? {
for (output, dot_value) in output.iter_mut().zip(dot.iter().copied()) {
*output = if beta_is_zero {
alpha * dot_value
} else {
alpha * dot_value + beta * *output
};
}
return Ok(());
}
for (linear, dot_value) in dot.iter().copied().enumerate() {
let indices = flat_to_multi_for_shape(out.shape(), linear);
let output = out.get_mut(&indices).ok_or_else(|| {
invalid_argument(
"dot_general",
"output",
format!("index {indices:?} is outside accumulation target"),
)
})?;
*output = if beta_is_zero {
alpha * dot_value
} else {
alpha * dot_value + beta * *output
};
}
Ok(())
}
fn compact_host_accumulation_slice<'a, T: 'static>(
out: &'a mut TypedTensorViewMut<'_, T>,
expected_len: usize,
) -> crate::Result<Option<&'a mut [T]>> {
if out.backend_buffer().is_some()
|| out.n_elements() != expected_len
|| !out.is_col_major_contiguous()?
{
return Ok(None);
}
let start = usize::try_from(out.offset()).map_err(|_| {
invalid_argument("dot_general", "output", "compact output offset is negative")
})?;
let end = start
.checked_add(expected_len)
.ok_or_else(|| validation("dot_general", ValidationError::IntegerOverflow))?;
out.host_storage_mut()?
.get_mut(start..end)
.map(Some)
.ok_or_else(|| {
invalid_argument(
"dot_general",
"output",
"compact output is outside its backing storage",
)
})
}
fn flat_to_multi_for_shape(shape: &[usize], mut linear: usize) -> Vec<usize> {
let mut indices = Vec::with_capacity(shape.len());
for &dim in shape {
if dim == 0 {
indices.push(0);
} else {
indices.push(linear % dim);
linear /= dim;
}
}
indices
}
#[doc(hidden)]
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct ElementwiseFusionPlan {
dtype: crate::DType,
input_count: usize,
input_views: Vec<ElementwiseFusionInputView>,
outputs: Vec<usize>,
ops: Vec<ElementwiseFusionInst>,
}
#[doc(hidden)]
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub enum ElementwiseFusionInputView {
Identity,
BroadcastInDim {
shape: Vec<usize>,
dims: Vec<usize>,
},
}
#[doc(hidden)]
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct ElementwiseFusionInst {
op: ElementwiseFusionOp,
inputs: Vec<usize>,
}
tenferro_core_ops::define_elementwise_fusion_op!();
impl ElementwiseFusionPlan {
pub fn new(
dtype: crate::DType,
input_count: usize,
outputs: Vec<usize>,
ops: Vec<ElementwiseFusionInst>,
) -> Self {
Self::with_input_views(
dtype,
vec![ElementwiseFusionInputView::Identity; input_count],
outputs,
ops,
)
}
pub fn with_input_views(
dtype: crate::DType,
input_views: impl IntoIterator<Item = ElementwiseFusionInputView>,
outputs: Vec<usize>,
ops: Vec<ElementwiseFusionInst>,
) -> Self {
let input_views = input_views.into_iter().collect::<Vec<_>>();
let input_count = input_views.len();
Self {
dtype,
input_count,
input_views,
outputs,
ops,
}
}
pub fn dtype(&self) -> crate::DType {
self.dtype
}
pub fn input_count(&self) -> usize {
self.input_count
}
pub fn input_views(&self) -> &[ElementwiseFusionInputView] {
&self.input_views
}
pub fn outputs(&self) -> &[usize] {
&self.outputs
}
pub fn ops(&self) -> &[ElementwiseFusionInst] {
&self.ops
}
}
impl ElementwiseFusionInputView {
pub fn broadcast_in_dim(
shape: impl IntoIterator<Item = usize>,
dims: impl IntoIterator<Item = usize>,
) -> Self {
Self::BroadcastInDim {
shape: shape.into_iter().collect(),
dims: dims.into_iter().collect(),
}
}
pub fn is_identity(&self) -> bool {
matches!(self, Self::Identity)
}
}
impl ElementwiseFusionInst {
pub fn new(op: ElementwiseFusionOp, inputs: Vec<usize>) -> Self {
Self { op, inputs }
}
pub fn op(&self) -> ElementwiseFusionOp {
self.op
}
pub fn inputs(&self) -> &[usize] {
&self.inputs
}
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ElementwiseReadOp {
Add,
Subtract,
Multiply,
Negate,
Conj,
Divide,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ActivationOp {
Sigmoid,
Silu,
Softplus,
Gelu,
GeluTanh,
}
impl ActivationOp {
#[doc(hidden)]
pub fn label(self) -> &'static str {
match self {
Self::Sigmoid => "sigmoid",
Self::Silu => "silu",
Self::Softplus => "softplus",
Self::Gelu => "gelu",
Self::GeluTanh => "gelu_tanh",
}
}
}
impl ElementwiseReadOp {
#[doc(hidden)]
pub fn label(self) -> &'static str {
match self {
Self::Add => "add",
Self::Subtract => "sub",
Self::Multiply => "mul",
Self::Negate => "neg",
Self::Conj => "conj",
Self::Divide => "div",
}
}
#[doc(hidden)]
pub fn arity(self) -> usize {
match self {
Self::Negate | Self::Conj => 1,
Self::Add | Self::Subtract | Self::Multiply | Self::Divide => 2,
}
}
}
#[derive(Clone, Copy, Debug)]
enum StorageIdentity {
Host {
start: usize,
end: usize,
},
Backend {
domain: Option<AllocationDomainId>,
allocation: Option<AllocationId>,
family: &'static str,
object: usize,
},
}
fn host_storage_identity<T>(data: &[T]) -> StorageIdentity {
let start = data.as_ptr() as usize;
let bytes = std::mem::size_of_val(data);
StorageIdentity::Host {
start,
end: start.saturating_add(bytes),
}
}
fn backend_storage_identity<T: 'static>(buffer: &dyn crate::BackendStorage<T>) -> StorageIdentity {
StorageIdentity::Backend {
domain: buffer.allocation_domain(),
allocation: buffer.allocation_id(),
family: buffer.backend_family(),
object: buffer as *const dyn crate::BackendStorage<T> as *const () as usize,
}
}
fn typed_tensor_storage_identity<T: crate::TensorScalar>(
tensor: &TypedTensor<T>,
) -> crate::Result<StorageIdentity> {
if tensor.backend_buffer().is_some() {
let buffer = tensor.backend_buffer().ok_or_else(|| {
crate::Error::runtime_state("typed_tensor_storage_identity", "backend buffer missing")
})?;
Ok(backend_storage_identity(buffer))
} else {
Ok(host_storage_identity(tensor.host_data()?))
}
}
fn typed_view_storage_identity<T: crate::TensorScalar + 'static>(
view: &TypedTensorView<'_, T>,
) -> crate::Result<StorageIdentity> {
match view.backend_buffer() {
Some(buffer) => Ok(backend_storage_identity(buffer)),
None => view.host_storage().map(host_storage_identity),
}
}
fn identity_operand<T: TensorScalar>(value: &Tensor) -> crate::Result<&TypedTensor<T>> {
value.as_typed::<T>().ok_or_else(|| {
crate::Error::unsupported_dtype(
"storage_identity",
value.dtype(),
"an externally defined payload has no storage identity",
)
})
}
fn tensor_read_storage_identity(input: &TensorRead<'_>) -> crate::Result<StorageIdentity> {
macro_rules! typed_identity {
($value:expr) => {
match $value.dtype() {
DType::F32 => typed_tensor_storage_identity(identity_operand::<f32>($value)?),
DType::F64 => typed_tensor_storage_identity(identity_operand::<f64>($value)?),
DType::I32 => typed_tensor_storage_identity(identity_operand::<i32>($value)?),
DType::I64 => typed_tensor_storage_identity(identity_operand::<i64>($value)?),
DType::Bool => typed_tensor_storage_identity(identity_operand::<bool>($value)?),
DType::C32 => typed_tensor_storage_identity(identity_operand::<Complex32>($value)?),
DType::C64 => typed_tensor_storage_identity(identity_operand::<Complex64>($value)?),
DType::External(_) => {
return Err(crate::Error::unsupported_dtype(
"storage_identity",
$value.dtype(),
"an externally defined payload has no storage identity",
));
}
}
};
}
macro_rules! view_identity {
($value:expr) => {
match $value {
TensorView::F32(value) => typed_view_storage_identity(value),
TensorView::F64(value) => typed_view_storage_identity(value),
TensorView::I32(value) => typed_view_storage_identity(value),
TensorView::I64(value) => typed_view_storage_identity(value),
TensorView::Bool(value) => typed_view_storage_identity(value),
TensorView::C32(value) => typed_view_storage_identity(value),
TensorView::C64(value) => typed_view_storage_identity(value),
}
};
}
match input {
TensorRead::Tensor(tensor) => typed_identity!(tensor),
TensorRead::View(view) => view_identity!(view),
}
}
fn storage_overlaps(lhs: StorageIdentity, rhs: StorageIdentity) -> bool {
match (lhs, rhs) {
(
StorageIdentity::Host {
start: lhs_start,
end: lhs_end,
},
StorageIdentity::Host {
start: rhs_start,
end: rhs_end,
},
) => lhs_start < rhs_end && rhs_start < lhs_end,
(
StorageIdentity::Backend {
domain: lhs_domain,
allocation: lhs_allocation,
family: lhs_family,
object: lhs_object,
},
StorageIdentity::Backend {
domain: rhs_domain,
allocation: rhs_allocation,
family: rhs_family,
object: rhs_object,
},
) => {
lhs_object == rhs_object
|| matches!(
(lhs_domain, rhs_domain, lhs_allocation, rhs_allocation),
(Some(lhs_domain), Some(rhs_domain), Some(lhs), Some(rhs))
if lhs_domain == rhs_domain && lhs == rhs
)
|| matches!(
(lhs_domain, rhs_domain, lhs_allocation, rhs_allocation),
(None, None, Some(lhs), Some(rhs)) if lhs_family == rhs_family && lhs == rhs
)
}
_ => false,
}
}
pub fn validate_read_into_destination(
op: &'static str,
inputs: &[TensorRead<'_>],
out: &TensorWrite<'_>,
) -> crate::Result<()> {
let output_identity = tensor_read_storage_identity(&out.as_read())?;
for (index, input) in inputs.iter().enumerate() {
if storage_overlaps(tensor_read_storage_identity(input)?, output_identity) {
return Err(Error::invalid_argument(
op,
"out",
format!("destination storage overlaps input {index}"),
));
}
}
Ok(())
}
#[doc(hidden)]
pub fn elementwise_read_into_via_allocating_ops<B: TensorElementwise + ?Sized>(
backend: &mut B,
op: ElementwiseReadOp,
inputs: &[TensorRead<'_>],
out: TensorWrite<'_>,
) -> crate::Result<()> {
if inputs.len() != op.arity() {
return Err(Error::invalid_argument(
op.label(),
"inputs",
format!("expected {} inputs, got {}", op.arity(), inputs.len()),
));
}
validate_read_into_destination(op.label(), inputs, &out)?;
let result = match op {
ElementwiseReadOp::Add => backend.add_read(inputs[0].clone(), inputs[1].clone())?,
ElementwiseReadOp::Subtract => backend.sub_read(inputs[0].clone(), inputs[1].clone())?,
ElementwiseReadOp::Multiply => backend.mul_read(inputs[0].clone(), inputs[1].clone())?,
ElementwiseReadOp::Negate => backend.neg_read(inputs[0].clone())?,
ElementwiseReadOp::Conj => backend.conj_read(inputs[0].clone())?,
ElementwiseReadOp::Divide => backend.div_read(inputs[0].clone(), inputs[1].clone())?,
};
backend.copy_read_into(TensorRead::from_tensor(&result), out)
}
pub trait TensorElementwise: TensorStructural {
fn elementwise_read_into(
&mut self,
op: ElementwiseReadOp,
inputs: &[TensorRead<'_>],
out: TensorWrite<'_>,
) -> crate::Result<()>;
fn add_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn add_into(&mut self, lhs: &Tensor, rhs: &Tensor, out: TensorWrite<'_>) -> crate::Result<()> {
self.add_read_into(
TensorRead::from_tensor(lhs),
TensorRead::from_tensor(rhs),
out,
)
}
fn add_read_into(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.elementwise_read_into(ElementwiseReadOp::Add, &[lhs, rhs], out)
}
fn sub_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn sub_into(&mut self, lhs: &Tensor, rhs: &Tensor, out: TensorWrite<'_>) -> crate::Result<()> {
self.sub_read_into(
TensorRead::from_tensor(lhs),
TensorRead::from_tensor(rhs),
out,
)
}
fn sub_read_into(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.elementwise_read_into(ElementwiseReadOp::Subtract, &[lhs, rhs], out)
}
fn mul_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn mul_into(&mut self, lhs: &Tensor, rhs: &Tensor, out: TensorWrite<'_>) -> crate::Result<()> {
self.mul_read_into(
TensorRead::from_tensor(lhs),
TensorRead::from_tensor(rhs),
out,
)
}
fn mul_read_into(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.elementwise_read_into(ElementwiseReadOp::Multiply, &[lhs, rhs], out)
}
fn neg_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn neg_into(&mut self, input: &Tensor, out: TensorWrite<'_>) -> crate::Result<()> {
self.neg_read_into(TensorRead::from_tensor(input), out)
}
fn neg_read_into(&mut self, input: TensorRead<'_>, out: TensorWrite<'_>) -> crate::Result<()> {
self.elementwise_read_into(ElementwiseReadOp::Negate, &[input], out)
}
fn conj_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn conj_into(&mut self, input: &Tensor, out: TensorWrite<'_>) -> crate::Result<()> {
self.conj_read_into(TensorRead::from_tensor(input), out)
}
fn conj_read_into(&mut self, input: TensorRead<'_>, out: TensorWrite<'_>) -> crate::Result<()> {
self.elementwise_read_into(ElementwiseReadOp::Conj, &[input], out)
}
fn div_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn div_into(&mut self, lhs: &Tensor, rhs: &Tensor, out: TensorWrite<'_>) -> crate::Result<()> {
self.div_read_into(
TensorRead::from_tensor(lhs),
TensorRead::from_tensor(rhs),
out,
)
}
fn div_read_into(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.elementwise_read_into(ElementwiseReadOp::Divide, &[lhs, rhs], out)
}
fn rem(&mut self, lhs: &Tensor, _rhs: &Tensor) -> crate::Result<Tensor> {
Err(crate::Error::unsupported(
"rem",
format!("backend does not implement rem for dtype {:?}", lhs.dtype()),
))
}
fn rem_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
self.rem(read_tensor("rem", lhs)?, read_tensor("rem", rhs)?)
}
fn abs_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn sign_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn maximum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn minimum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn compare_read(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
dir: &CompareDir,
) -> crate::Result<Tensor>;
fn select_read(
&mut self,
pred: TensorRead<'_>,
on_true: TensorRead<'_>,
on_false: TensorRead<'_>,
) -> crate::Result<Tensor>;
fn clamp_read(
&mut self,
input: TensorRead<'_>,
lower: TensorRead<'_>,
upper: TensorRead<'_>,
) -> crate::Result<Tensor>;
}
pub trait TensorAnalytic {
fn exp_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn log_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn sin_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn cos_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn tanh_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn sqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn rsqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn pow_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
fn expm1_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn log1p_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
fn erf_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor>;
}
pub trait TensorStructural {
fn to_contiguous_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
match input {
TensorRead::Tensor(input) => {
if input.is_backend_buffer()
|| !matches!(
input.placement().memory_kind,
crate::MemoryKind::PinnedHost | crate::MemoryKind::UnpinnedHost
)
{
return Err(crate::Error::runtime_state(
"to_contiguous_read",
"default materialization accepts only host-owned tensors; use the storage's owning backend",
));
}
input.duplicate()
}
TensorRead::View(view) => {
if view.backend_family().is_some()
|| !matches!(
view.placement().memory_kind,
crate::MemoryKind::PinnedHost | crate::MemoryKind::UnpinnedHost
)
{
return Err(crate::Error::runtime_state(
"to_contiguous_read",
"default materialization accepts only host-owned tensors; use the storage's owning backend",
));
}
default_materialize_host_view(view)
}
}
}
fn copy_read_into(&mut self, _src: TensorRead<'_>, _dst: TensorWrite<'_>) -> crate::Result<()> {
Err(crate::Error::unsupported(
"copy_read_into",
"backend-owned runtime copy is unsupported by this backend",
))
}
fn transpose_read(&mut self, input: TensorRead<'_>, perm: &[usize]) -> crate::Result<Tensor>;
fn reshape_read(&mut self, input: TensorRead<'_>, shape: &[usize]) -> crate::Result<Tensor>;
fn broadcast_in_dim_read(
&mut self,
input: TensorRead<'_>,
shape: &[usize],
dims: &[usize],
) -> crate::Result<Tensor>;
fn cast(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor>;
fn convert(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor> {
validate_convert_dtype("convert", input.dtype(), to)?;
self.cast(input, to)
}
fn extract_diagonal(
&mut self,
input: &Tensor,
axis_a: usize,
axis_b: usize,
) -> crate::Result<Tensor>;
fn extract_diagonal_read(
&mut self,
input: TensorRead<'_>,
axis_a: usize,
axis_b: usize,
) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.extract_diagonal(&input, axis_a, axis_b)
}
fn embed_diagonal(
&mut self,
input: &Tensor,
axis_a: usize,
axis_b: usize,
) -> crate::Result<Tensor>;
fn embed_diagonal_read(
&mut self,
input: TensorRead<'_>,
axis_a: usize,
axis_b: usize,
) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.embed_diagonal(&input, axis_a, axis_b)
}
fn tril(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor>;
fn tril_read(&mut self, input: TensorRead<'_>, k: i64) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.tril(&input, k)
}
fn triu(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor>;
fn triu_read(&mut self, input: TensorRead<'_>, k: i64) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.triu(&input, k)
}
}
pub trait TensorReduction {
fn reduce_sum_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
#[doc(hidden)]
fn reduce_sum_squares_read(
&mut self,
_input: TensorRead<'_>,
_axes: &[usize],
) -> crate::Result<Tensor> {
Err(crate::Error::unsupported(
"reduce_sum_squares",
"backend does not implement fused sum-of-squares reduction",
))
}
fn reduce_prod_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
fn reduce_max_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
fn reduce_min_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
}
pub trait TensorDot: TensorElementwise {
#[doc(hidden)]
fn dot_general_read(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
) -> crate::Result<Tensor>;
fn dot_general_read_into(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
out: TensorWrite<'_>,
) -> crate::Result<()> {
let accumulation = DotGeneralAccumulation::overwrite(lhs.dtype())?;
self.dot_general_read_into_accum(lhs, rhs, config, accumulation, out)
}
#[doc(hidden)]
fn dot_general_with_conj(
&mut self,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> crate::Result<Tensor> {
if !lhs_conj && !rhs_conj {
return self.dot_general_read(
TensorRead::from_tensor(lhs),
TensorRead::from_tensor(rhs),
config,
);
}
let lhs_tmp;
let lhs_ref = if lhs_conj {
lhs_tmp = self.conj_read(TensorRead::from_tensor(lhs))?;
&lhs_tmp
} else {
lhs
};
let rhs_tmp;
let rhs_ref = if rhs_conj {
rhs_tmp = self.conj_read(TensorRead::from_tensor(rhs))?;
&rhs_tmp
} else {
rhs
};
self.dot_general_read(
TensorRead::from_tensor(lhs_ref),
TensorRead::from_tensor(rhs_ref),
config,
)
}
#[allow(clippy::too_many_arguments)]
#[doc(hidden)]
fn dot_general_with_conj_read(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> crate::Result<Tensor> {
if !lhs_conj && !rhs_conj {
return self.dot_general_read(lhs, rhs, config);
}
let lhs_tmp;
let lhs_ref = if let Some(tensor) = lhs.as_tensor() {
tensor
} else {
lhs_tmp = self.to_contiguous_read(lhs)?;
&lhs_tmp
};
let rhs_tmp;
let rhs_ref = if let Some(tensor) = rhs.as_tensor() {
tensor
} else {
rhs_tmp = self.to_contiguous_read(rhs)?;
&rhs_tmp
};
self.dot_general_with_conj(lhs_ref, rhs_ref, config, lhs_conj, rhs_conj)
}
fn dot_general_read_into_accum(
&mut self,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
accumulation: DotGeneralAccumulation,
out: TensorWrite<'_>,
) -> crate::Result<()> {
dot_general_accum_via_temp(self, lhs, rhs, config, accumulation, out)
}
}
pub trait SessionCachedDot: TensorDot {
#[doc(hidden)]
fn dot_general_cached(
&mut self,
_cache_slot: Option<usize>,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
) -> crate::Result<Tensor> {
self.dot_general_read(
TensorRead::from_tensor(lhs),
TensorRead::from_tensor(rhs),
config,
)
}
#[doc(hidden)]
fn dot_general_read_cached(
&mut self,
cache_slot: Option<usize>,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
) -> crate::Result<Tensor> {
match (lhs.as_tensor(), rhs.as_tensor()) {
(Some(lhs), Some(rhs)) => self.dot_general_cached(cache_slot, lhs, rhs, config),
_ => {
let lhs = self.to_contiguous_read(lhs)?;
let rhs = self.to_contiguous_read(rhs)?;
self.dot_general_cached(cache_slot, &lhs, &rhs, config)
}
}
}
#[allow(clippy::too_many_arguments)]
#[doc(hidden)]
fn dot_general_with_conj_cached(
&mut self,
_cache_slot: Option<usize>,
lhs: &Tensor,
rhs: &Tensor,
config: &DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> crate::Result<Tensor> {
self.dot_general_with_conj(lhs, rhs, config, lhs_conj, rhs_conj)
}
#[allow(clippy::too_many_arguments)]
#[doc(hidden)]
fn dot_general_with_conj_read_cached(
&mut self,
cache_slot: Option<usize>,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> crate::Result<Tensor> {
if !lhs_conj && !rhs_conj {
return self.dot_general_read_cached(cache_slot, lhs, rhs, config);
}
let lhs_tmp;
let lhs_ref = if let Some(tensor) = lhs.as_tensor() {
tensor
} else {
lhs_tmp = self.to_contiguous_read(lhs)?;
&lhs_tmp
};
let rhs_tmp;
let rhs_ref = if let Some(tensor) = rhs.as_tensor() {
tensor
} else {
rhs_tmp = self.to_contiguous_read(rhs)?;
&rhs_tmp
};
self.dot_general_with_conj_cached(cache_slot, lhs_ref, rhs_ref, config, lhs_conj, rhs_conj)
}
fn dot_general_read_into_accum_cached(
&mut self,
_cache_slot: Option<usize>,
lhs: TensorRead<'_>,
rhs: TensorRead<'_>,
config: &DotGeneralConfig,
accumulation: DotGeneralAccumulation,
out: TensorWrite<'_>,
) -> crate::Result<()> {
self.dot_general_read_into_accum(lhs, rhs, config, accumulation, out)
}
}
pub trait TensorIndexing: TensorStructural {
fn gather(
&mut self,
operand: &Tensor,
start_indices: &Tensor,
config: &GatherConfig,
) -> crate::Result<Tensor>;
fn scatter(
&mut self,
operand: &Tensor,
scatter_indices: &Tensor,
updates: &Tensor,
config: &ScatterConfig,
) -> crate::Result<Tensor>;
fn scatter_read(
&mut self,
operand: TensorRead<'_>,
scatter_indices: TensorRead<'_>,
updates: TensorRead<'_>,
config: &ScatterConfig,
) -> crate::Result<Tensor> {
let operand = stage_view(self, operand)?;
let scatter_indices = stage_view(self, scatter_indices)?;
let updates = stage_view(self, updates)?;
self.scatter(&operand, &scatter_indices, &updates, config)
}
fn slice(&mut self, input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor>;
fn slice_read(&mut self, input: TensorRead<'_>, config: &SliceConfig) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.slice(&input, config)
}
fn dynamic_slice(
&mut self,
input: &Tensor,
starts: &Tensor,
slice_sizes: &[usize],
) -> crate::Result<Tensor>;
fn dynamic_slice_read(
&mut self,
input: TensorRead<'_>,
starts: TensorRead<'_>,
slice_sizes: &[usize],
) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
let starts = stage_view(self, starts)?;
self.dynamic_slice(&input, &starts, slice_sizes)
}
fn dynamic_update_slice(
&mut self,
operand: &Tensor,
update: &Tensor,
starts: &Tensor,
) -> crate::Result<Tensor>;
fn dynamic_update_slice_read(
&mut self,
operand: TensorRead<'_>,
update: TensorRead<'_>,
starts: TensorRead<'_>,
) -> crate::Result<Tensor> {
let operand = stage_view(self, operand)?;
let update = stage_view(self, update)?;
let starts = stage_view(self, starts)?;
self.dynamic_update_slice(&operand, &update, &starts)
}
fn pad(&mut self, input: &Tensor, config: &PadConfig) -> crate::Result<Tensor>;
fn pad_read(&mut self, input: TensorRead<'_>, config: &PadConfig) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.pad(&input, config)
}
fn concatenate(&mut self, inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor>;
fn concatenate_read(
&mut self,
inputs: &[TensorRead<'_>],
axis: usize,
) -> crate::Result<Tensor> {
let staged = inputs
.iter()
.map(|input| stage_view(self, input.clone()))
.collect::<crate::Result<Vec<_>>>()?;
let refs: Vec<&Tensor> = staged.iter().map(|input| &**input).collect();
self.concatenate(&refs, axis)
}
fn reverse(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
fn reverse_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
let input = stage_view(self, input)?;
self.reverse(&input, axes)
}
}
pub trait TensorViewCanonicalization<T: TensorScalar, R: TensorRank> {
fn to_contiguous(
&mut self,
view: &TypedTensorView<'_, T, R>,
) -> crate::Result<TypedTensor<T, R>>;
fn copy_into(
&mut self,
src: &TypedTensorView<'_, T, R>,
dst: &mut TypedTensorViewMut<'_, T, R>,
) -> crate::Result<()>;
}
pub trait TensorFusion {
#[doc(hidden)]
fn execute_elementwise_fusion(
&mut self,
_inputs: &[&Tensor],
_plan: &ElementwiseFusionPlan,
) -> crate::Result<Option<Vec<Tensor>>> {
Ok(None)
}
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
fn execute_broadcast_multiply(
&mut self,
_lhs: TensorRead<'_>,
_lhs_shape: &[usize],
_lhs_dims: &[usize],
_rhs: TensorRead<'_>,
_rhs_shape: &[usize],
_rhs_dims: &[usize],
) -> crate::Result<Option<Tensor>> {
Ok(None)
}
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
fn execute_broadcast_multiply_value(
&mut self,
lhs: TensorRead<'_>,
lhs_shape: &[usize],
lhs_dims: &[usize],
rhs: TensorRead<'_>,
rhs_shape: &[usize],
rhs_dims: &[usize],
) -> crate::Result<Option<TensorValue>> {
self.execute_broadcast_multiply(lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims)
.map(|tensor| tensor.map(TensorValue::from_tensor))
}
}
pub trait TensorBuffer {
fn reclaim_buffer(&mut self, _tensor: Tensor) {}
}
pub trait TensorDeviceTransfer {
fn download_to_host(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor>;
fn upload_host_tensor(&mut self, tensor: TensorRead<'_>) -> crate::Result<Tensor>;
}
#[doc(hidden)]
pub trait TensorBackendOps:
TensorElementwise
+ TensorAnalytic
+ TensorStructural
+ TensorReduction
+ TensorIndexing
+ TensorDot
+ TensorFusion
+ TensorBuffer
{
}
impl<T> TensorBackendOps for T where
T: TensorElementwise
+ TensorAnalytic
+ TensorStructural
+ TensorReduction
+ TensorIndexing
+ TensorDot
+ TensorFusion
+ TensorBuffer
+ ?Sized
{
}
#[doc(hidden)]
pub fn validate_vdot_read(lhs: &TensorRead<'_>, rhs: &TensorRead<'_>) -> crate::Result<()> {
let op = "Engine::vdot_read";
validate_supported_blas1_dtype(op, lhs.dtype())?;
if rhs.dtype() != lhs.dtype() {
return Err(Error::dtype_mismatch(op, lhs.dtype(), rhs.dtype()));
}
if lhs.shape() != rhs.shape() {
return Err(Error::shape_mismatch(
op,
lhs.shape().to_vec(),
rhs.shape().to_vec(),
));
}
validate_shape_product(op, lhs.shape())?;
validate_compatible_placement(op, lhs, rhs)
}
#[doc(hidden)]
pub fn validate_norm_squared_read(input: &TensorRead<'_>) -> crate::Result<()> {
validate_supported_blas1_dtype("Engine::norm_squared_read", input.dtype())?;
validate_shape_product("Engine::norm_squared_read", input.shape()).map(|_| ())
}
#[doc(hidden)]
pub fn validate_axpby_read_into_accum(
alpha: ContractionScalar,
x: &TensorRead<'_>,
beta: ContractionScalar,
y: &TensorWrite<'_>,
) -> crate::Result<()> {
let op = "Engine::axpby_read_into_accum";
validate_supported_blas1_dtype(op, x.dtype())?;
if y.dtype() != x.dtype() {
return Err(Error::dtype_mismatch(op, x.dtype(), y.dtype()));
}
if alpha.dtype() != x.dtype() {
return Err(Error::dtype_mismatch(op, x.dtype(), alpha.dtype()));
}
if beta.dtype() != x.dtype() {
return Err(Error::dtype_mismatch(op, x.dtype(), beta.dtype()));
}
if x.shape() != y.shape() {
return Err(Error::shape_mismatch(
op,
x.shape().to_vec(),
y.shape().to_vec(),
));
}
validate_shape_product(op, x.shape())?;
validate_compatible_placement(op, x, &y.as_read())?;
validate_read_into_destination(op, std::slice::from_ref(x), y)
}
fn validate_supported_blas1_dtype(op: &'static str, dtype: DType) -> crate::Result<()> {
if matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
Ok(())
} else {
Err(Error::unsupported(
op,
format!("unsupported dtype {dtype:?}; supported dtypes: F32/F64/C32/C64"),
))
}
}
fn validate_shape_product(op: &'static str, shape: &[usize]) -> crate::Result<usize> {
crate::validate::checked_shape_product(op, "shape", shape)
}
fn validate_compatible_placement(
op: &'static str,
lhs: &TensorRead<'_>,
rhs: &TensorRead<'_>,
) -> crate::Result<()> {
if lhs.placement() != rhs.placement() {
return Err(Error::invalid_argument(
op,
"placement",
"tensor inputs must have compatible placement",
));
}
Ok(())
}
pub trait Engine: TensorBackendOps + SessionCachedDot + TensorDeviceTransfer {
fn gather_read(
&mut self,
operand: TensorRead<'_>,
start_indices: TensorRead<'_>,
config: &GatherConfig,
) -> crate::Result<Tensor> {
let operand = stage_view(self, operand)?;
let start_indices = stage_view(self, start_indices)?;
self.gather(&operand, &start_indices, config)
}
fn vdot_read(&mut self, _lhs: TensorRead<'_>, _rhs: TensorRead<'_>) -> crate::Result<Tensor> {
Err(Error::unsupported(
"Engine::vdot_read",
"backend session does not implement vdot_read",
))
}
fn norm_squared_read(&mut self, _input: TensorRead<'_>) -> crate::Result<Tensor> {
Err(Error::unsupported(
"Engine::norm_squared_read",
"backend session does not implement norm_squared_read",
))
}
fn fused_activation_read(
&mut self,
op: ActivationOp,
input: TensorRead<'_>,
) -> crate::Result<Option<Tensor>> {
let _ = (op, input);
Ok(None)
}
fn softmax_read(
&mut self,
input: TensorRead<'_>,
axis: usize,
) -> crate::Result<Option<Tensor>> {
let _ = (input, axis);
Ok(None)
}
fn axpby_read_into_accum(
&mut self,
_alpha: ContractionScalar,
_x: TensorRead<'_>,
_beta: ContractionScalar,
_y: TensorWrite<'_>,
) -> crate::Result<()> {
Err(Error::unsupported(
"Engine::axpby_read_into_accum",
"backend session does not implement axpby_read_into_accum",
))
}
fn native_engine(&mut self) -> Option<crate::NativeSessionRef<'_>> {
None
}
}
enum StagedRead<'a> {
Borrowed(&'a Tensor),
Staged(Box<Tensor>),
}
impl std::ops::Deref for StagedRead<'_> {
type Target = Tensor;
fn deref(&self) -> &Tensor {
match self {
Self::Borrowed(tensor) => tensor,
Self::Staged(tensor) => tensor,
}
}
}
fn stage_view<'a, S: TensorStructural + ?Sized>(
session: &mut S,
read: TensorRead<'a>,
) -> crate::Result<StagedRead<'a>> {
match read {
TensorRead::Tensor(tensor) => Ok(StagedRead::Borrowed(tensor)),
view => Ok(StagedRead::Staged(Box::new(
session.to_contiguous_read(view)?,
))),
}
}