tenferro_tensor/
backend.rs

1use crate::config::{
2    CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
3};
4use crate::types::{TensorRank, TypedTensor, TypedTensorView, TypedTensorViewMut};
5use crate::validate::validate_convert_dtype;
6use crate::{RuntimeCacheControl, Tensor, TensorRead, TensorValue, TensorWrite};
7
8fn read_boundary_error(op: &'static str) -> crate::Error {
9    crate::Error::backend_failure(
10        op,
11        "backend does not accept borrowed tensor views at this execution boundary",
12    )
13}
14
15fn read_tensor<'a>(op: &'static str, input: TensorRead<'a>) -> crate::Result<&'a Tensor> {
16    input.as_tensor().ok_or_else(|| read_boundary_error(op))
17}
18
19fn validate_axis_list(
20    op: &'static str,
21    role: &'static str,
22    axes: &[usize],
23    rank: usize,
24) -> crate::Result<()> {
25    let mut seen = vec![false; rank];
26    for &axis in axes {
27        if axis >= rank {
28            return Err(crate::Error::AxisOutOfBounds { op, axis, rank });
29        }
30        if seen[axis] {
31            return Err(crate::Error::DuplicateAxis { op, axis, role });
32        }
33        seen[axis] = true;
34    }
35    Ok(())
36}
37
38fn validate_role_disjoint(
39    op: &'static str,
40    first_role: &'static str,
41    first_axes: &[usize],
42    second_role: &'static str,
43    second_axes: &[usize],
44) -> crate::Result<()> {
45    for &axis in first_axes {
46        if second_axes.contains(&axis) {
47            return Err(crate::Error::AxisRoleConflict {
48                op,
49                axis,
50                first_role,
51                second_role,
52            });
53        }
54    }
55    Ok(())
56}
57
58/// Infer the output shape for a validated dot-general operation.
59#[doc(hidden)]
60pub fn dot_general_output_shape(
61    lhs_shape: &[usize],
62    rhs_shape: &[usize],
63    config: &DotGeneralConfig,
64    op: &'static str,
65) -> crate::Result<Vec<usize>> {
66    if config.lhs_contracting_dims.len() != config.rhs_contracting_dims.len() {
67        return Err(crate::Error::InvalidConfig {
68            op,
69            message: "lhs/rhs contracting dim counts differ".into(),
70        });
71    }
72    if config.lhs_batch_dims.len() != config.rhs_batch_dims.len() {
73        return Err(crate::Error::InvalidConfig {
74            op,
75            message: "lhs/rhs batch dim counts differ".into(),
76        });
77    }
78
79    let lhs_rank = lhs_shape.len();
80    let rhs_rank = rhs_shape.len();
81    validate_axis_list(
82        op,
83        "lhs_contracting",
84        &config.lhs_contracting_dims,
85        lhs_rank,
86    )?;
87    validate_axis_list(
88        op,
89        "rhs_contracting",
90        &config.rhs_contracting_dims,
91        rhs_rank,
92    )?;
93    validate_axis_list(op, "lhs_batch", &config.lhs_batch_dims, lhs_rank)?;
94    validate_axis_list(op, "rhs_batch", &config.rhs_batch_dims, rhs_rank)?;
95    validate_role_disjoint(
96        op,
97        "lhs_contracting",
98        &config.lhs_contracting_dims,
99        "lhs_batch",
100        &config.lhs_batch_dims,
101    )?;
102    validate_role_disjoint(
103        op,
104        "rhs_contracting",
105        &config.rhs_contracting_dims,
106        "rhs_batch",
107        &config.rhs_batch_dims,
108    )?;
109
110    for (&lhs_axis, &rhs_axis) in config
111        .lhs_contracting_dims
112        .iter()
113        .zip(&config.rhs_contracting_dims)
114    {
115        if lhs_shape[lhs_axis] != rhs_shape[rhs_axis] {
116            return Err(crate::Error::ShapeMismatch {
117                op,
118                lhs: lhs_shape.to_vec(),
119                rhs: rhs_shape.to_vec(),
120            });
121        }
122    }
123    for (&lhs_axis, &rhs_axis) in config.lhs_batch_dims.iter().zip(&config.rhs_batch_dims) {
124        if lhs_shape[lhs_axis] != rhs_shape[rhs_axis] {
125            return Err(crate::Error::ShapeMismatch {
126                op,
127                lhs: lhs_shape.to_vec(),
128                rhs: rhs_shape.to_vec(),
129            });
130        }
131    }
132
133    let lhs_free = (0..lhs_rank)
134        .filter(|axis| {
135            !config.lhs_contracting_dims.contains(axis) && !config.lhs_batch_dims.contains(axis)
136        })
137        .map(|axis| lhs_shape[axis]);
138    let rhs_free = (0..rhs_rank)
139        .filter(|axis| {
140            !config.rhs_contracting_dims.contains(axis) && !config.rhs_batch_dims.contains(axis)
141        })
142        .map(|axis| rhs_shape[axis]);
143    let batch = config.lhs_batch_dims.iter().map(|&axis| lhs_shape[axis]);
144
145    Ok(lhs_free.chain(rhs_free).chain(batch).collect())
146}
147
148/// Validate output dtype and shape for dot-general read-into dispatch.
149#[doc(hidden)]
150pub fn validate_dot_general_read_into(
151    lhs: &TensorRead<'_>,
152    rhs: &TensorRead<'_>,
153    config: &DotGeneralConfig,
154    out: &TensorWrite<'_>,
155    op: &'static str,
156) -> crate::Result<Vec<usize>> {
157    if lhs.dtype() != rhs.dtype() {
158        return Err(crate::Error::DTypeMismatch {
159            op,
160            lhs: lhs.dtype(),
161            rhs: rhs.dtype(),
162        });
163    }
164    if out.dtype() != lhs.dtype() {
165        return Err(crate::Error::DTypeMismatch {
166            op,
167            lhs: out.dtype(),
168            rhs: lhs.dtype(),
169        });
170    }
171    let expected = dot_general_output_shape(lhs.shape(), rhs.shape(), config, op)?;
172    if out.shape() != expected.as_slice() {
173        return Err(crate::Error::ShapeMismatch {
174            op,
175            lhs: out.shape().to_vec(),
176            rhs: expected.clone(),
177        });
178    }
179    Ok(expected)
180}
181
182/// Canonical elementwise fusion plan shared between segmented execution and backends.
183#[doc(hidden)]
184#[derive(Clone, Debug, Hash, PartialEq, Eq)]
185pub struct ElementwiseFusionPlan {
186    dtype: crate::DType,
187    input_count: usize,
188    outputs: Vec<usize>,
189    ops: Vec<ElementwiseFusionInst>,
190}
191
192/// One node in a canonical elementwise fusion plan.
193#[doc(hidden)]
194#[derive(Clone, Debug, Hash, PartialEq, Eq)]
195pub struct ElementwiseFusionInst {
196    op: ElementwiseFusionOp,
197    inputs: Vec<usize>,
198}
199
200tenferro_core_ops::define_elementwise_fusion_op!();
201
202impl ElementwiseFusionPlan {
203    /// Build a backend elementwise fusion plan.
204    ///
205    /// # Examples
206    ///
207    /// ```rust
208    /// use tenferro_tensor::backend::{
209    ///     ElementwiseFusionInst, ElementwiseFusionOp, ElementwiseFusionPlan,
210    /// };
211    /// use tenferro_tensor::DType;
212    ///
213    /// let plan = ElementwiseFusionPlan::new(
214    ///     DType::F64,
215    ///     2,
216    ///     vec![2],
217    ///     vec![ElementwiseFusionInst::new(ElementwiseFusionOp::Add, vec![0, 1])],
218    /// );
219    /// assert_eq!(plan.input_count(), 2);
220    /// ```
221    pub fn new(
222        dtype: crate::DType,
223        input_count: usize,
224        outputs: Vec<usize>,
225        ops: Vec<ElementwiseFusionInst>,
226    ) -> Self {
227        Self {
228            dtype,
229            input_count,
230            outputs,
231            ops,
232        }
233    }
234
235    /// Return the scalar dtype expected by this fusion plan.
236    ///
237    /// # Examples
238    ///
239    /// ```rust
240    /// use tenferro_tensor::backend::ElementwiseFusionPlan;
241    /// use tenferro_tensor::DType;
242    ///
243    /// let plan = ElementwiseFusionPlan::new(DType::F32, 0, Vec::new(), Vec::new());
244    /// assert_eq!(plan.dtype(), DType::F32);
245    /// ```
246    pub fn dtype(&self) -> crate::DType {
247        self.dtype
248    }
249
250    /// Return the number of input tensors expected by this plan.
251    ///
252    /// # Examples
253    ///
254    /// ```rust
255    /// use tenferro_tensor::backend::ElementwiseFusionPlan;
256    /// use tenferro_tensor::DType;
257    ///
258    /// let plan = ElementwiseFusionPlan::new(DType::F64, 3, Vec::new(), Vec::new());
259    /// assert_eq!(plan.input_count(), 3);
260    /// ```
261    pub fn input_count(&self) -> usize {
262        self.input_count
263    }
264
265    /// Return the value ids selected as fusion outputs.
266    ///
267    /// # Examples
268    ///
269    /// ```rust
270    /// use tenferro_tensor::backend::ElementwiseFusionPlan;
271    /// use tenferro_tensor::DType;
272    ///
273    /// let plan = ElementwiseFusionPlan::new(DType::F64, 0, vec![0], Vec::new());
274    /// assert_eq!(plan.outputs(), &[0]);
275    /// ```
276    pub fn outputs(&self) -> &[usize] {
277        &self.outputs
278    }
279
280    /// Return the fused elementwise instruction sequence.
281    ///
282    /// # Examples
283    ///
284    /// ```rust
285    /// use tenferro_tensor::backend::{
286    ///     ElementwiseFusionInst, ElementwiseFusionOp, ElementwiseFusionPlan,
287    /// };
288    /// use tenferro_tensor::DType;
289    ///
290    /// let inst = ElementwiseFusionInst::new(ElementwiseFusionOp::Negate, vec![0]);
291    /// let plan = ElementwiseFusionPlan::new(DType::F64, 1, vec![1], vec![inst]);
292    /// assert_eq!(plan.ops().len(), 1);
293    /// ```
294    pub fn ops(&self) -> &[ElementwiseFusionInst] {
295        &self.ops
296    }
297}
298
299impl ElementwiseFusionInst {
300    /// Build a backend elementwise fusion instruction.
301    ///
302    /// # Examples
303    ///
304    /// ```rust
305    /// use tenferro_tensor::backend::{ElementwiseFusionInst, ElementwiseFusionOp};
306    ///
307    /// let inst = ElementwiseFusionInst::new(ElementwiseFusionOp::Add, vec![0, 1]);
308    /// assert_eq!(inst.inputs(), &[0, 1]);
309    /// ```
310    pub fn new(op: ElementwiseFusionOp, inputs: Vec<usize>) -> Self {
311        Self { op, inputs }
312    }
313
314    /// Return the elementwise op executed by this instruction.
315    ///
316    /// # Examples
317    ///
318    /// ```rust
319    /// use tenferro_tensor::backend::{ElementwiseFusionInst, ElementwiseFusionOp};
320    ///
321    /// let inst = ElementwiseFusionInst::new(ElementwiseFusionOp::Negate, vec![0]);
322    /// assert_eq!(inst.op(), ElementwiseFusionOp::Negate);
323    /// ```
324    pub fn op(&self) -> ElementwiseFusionOp {
325        self.op
326    }
327
328    /// Return this instruction's input value ids.
329    ///
330    /// # Examples
331    ///
332    /// ```rust
333    /// use tenferro_tensor::backend::{ElementwiseFusionInst, ElementwiseFusionOp};
334    ///
335    /// let inst = ElementwiseFusionInst::new(ElementwiseFusionOp::Multiply, vec![2, 0]);
336    /// assert_eq!(inst.inputs(), &[2, 0]);
337    /// ```
338    pub fn inputs(&self) -> &[usize] {
339        &self.inputs
340    }
341}
342
343/// Elementwise tensor operations.
344///
345/// # Examples
346///
347/// ```rust
348/// use tenferro_tensor::TensorElementwise;
349///
350/// fn accepts_elementwise<B: TensorElementwise>(_backend: &mut B) {}
351/// ```
352pub trait TensorElementwise {
353    fn add(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
354
355    /// Elementwise addition accepting either owned tensors or borrowed views.
356    ///
357    /// Backends that implement this method must not silently move data across
358    /// devices. A backend that cannot consume views should return an explicit
359    /// backend error rather than materializing or transferring implicitly.
360    ///
361    /// # Examples
362    ///
363    /// ```rust
364    /// use tenferro_tensor::{Tensor, TensorElementwise, TensorRead};
365    ///
366    /// fn add_owned<B: TensorElementwise>(
367    ///     backend: &mut B,
368    ///     lhs: &Tensor,
369    ///     rhs: &Tensor,
370    /// ) -> tenferro_tensor::Result<Tensor> {
371    ///     backend.add_read(TensorRead::from_tensor(lhs), TensorRead::from_tensor(rhs))
372    /// }
373    /// ```
374    fn add_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
375        self.add(read_tensor("add", lhs)?, read_tensor("add", rhs)?)
376    }
377
378    fn mul(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
379    fn mul_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
380        self.mul(read_tensor("mul", lhs)?, read_tensor("mul", rhs)?)
381    }
382
383    fn neg(&mut self, input: &Tensor) -> crate::Result<Tensor>;
384    fn neg_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
385        self.neg(read_tensor("neg", input)?)
386    }
387
388    fn conj(&mut self, input: &Tensor) -> crate::Result<Tensor>;
389    fn conj_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
390        self.conj(read_tensor("conj", input)?)
391    }
392
393    fn div(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
394    fn div_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
395        self.div(read_tensor("div", lhs)?, read_tensor("div", rhs)?)
396    }
397
398    fn abs(&mut self, input: &Tensor) -> crate::Result<Tensor>;
399    fn abs_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
400        self.abs(read_tensor("abs", input)?)
401    }
402
403    fn sign(&mut self, input: &Tensor) -> crate::Result<Tensor>;
404    fn sign_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
405        self.sign(read_tensor("sign", input)?)
406    }
407
408    fn maximum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
409    fn maximum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
410        self.maximum(read_tensor("maximum", lhs)?, read_tensor("maximum", rhs)?)
411    }
412
413    fn minimum(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
414    fn minimum_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
415        self.minimum(read_tensor("minimum", lhs)?, read_tensor("minimum", rhs)?)
416    }
417
418    fn compare(&mut self, lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor>;
419    fn compare_read(
420        &mut self,
421        lhs: TensorRead<'_>,
422        rhs: TensorRead<'_>,
423        dir: &CompareDir,
424    ) -> crate::Result<Tensor> {
425        self.compare(
426            read_tensor("compare", lhs)?,
427            read_tensor("compare", rhs)?,
428            dir,
429        )
430    }
431
432    fn select(
433        &mut self,
434        pred: &Tensor,
435        on_true: &Tensor,
436        on_false: &Tensor,
437    ) -> crate::Result<Tensor>;
438    fn select_read(
439        &mut self,
440        pred: TensorRead<'_>,
441        on_true: TensorRead<'_>,
442        on_false: TensorRead<'_>,
443    ) -> crate::Result<Tensor> {
444        self.select(
445            read_tensor("select", pred)?,
446            read_tensor("select", on_true)?,
447            read_tensor("select", on_false)?,
448        )
449    }
450
451    fn clamp(&mut self, input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor>;
452    fn clamp_read(
453        &mut self,
454        input: TensorRead<'_>,
455        lower: TensorRead<'_>,
456        upper: TensorRead<'_>,
457    ) -> crate::Result<Tensor> {
458        self.clamp(
459            read_tensor("clamp", input)?,
460            read_tensor("clamp", lower)?,
461            read_tensor("clamp", upper)?,
462        )
463    }
464}
465
466/// Analytic unary and binary tensor operations.
467///
468/// # Examples
469///
470/// ```rust
471/// use tenferro_tensor::TensorAnalytic;
472///
473/// fn accepts_analytic<B: TensorAnalytic>(_backend: &mut B) {}
474/// ```
475pub trait TensorAnalytic {
476    fn exp(&mut self, input: &Tensor) -> crate::Result<Tensor>;
477    fn exp_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
478        self.exp(read_tensor("exp", input)?)
479    }
480
481    fn log(&mut self, input: &Tensor) -> crate::Result<Tensor>;
482    fn log_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
483        self.log(read_tensor("log", input)?)
484    }
485
486    fn sin(&mut self, input: &Tensor) -> crate::Result<Tensor>;
487    fn sin_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
488        self.sin(read_tensor("sin", input)?)
489    }
490
491    fn cos(&mut self, input: &Tensor) -> crate::Result<Tensor>;
492    fn cos_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
493        self.cos(read_tensor("cos", input)?)
494    }
495
496    fn tanh(&mut self, input: &Tensor) -> crate::Result<Tensor>;
497    fn tanh_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
498        self.tanh(read_tensor("tanh", input)?)
499    }
500
501    fn sqrt(&mut self, input: &Tensor) -> crate::Result<Tensor>;
502    fn sqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
503        self.sqrt(read_tensor("sqrt", input)?)
504    }
505
506    fn rsqrt(&mut self, input: &Tensor) -> crate::Result<Tensor>;
507    fn rsqrt_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
508        self.rsqrt(read_tensor("rsqrt", input)?)
509    }
510
511    fn pow(&mut self, lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
512    fn pow_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
513        self.pow(read_tensor("pow", lhs)?, read_tensor("pow", rhs)?)
514    }
515
516    fn expm1(&mut self, input: &Tensor) -> crate::Result<Tensor>;
517    fn expm1_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
518        self.expm1(read_tensor("expm1", input)?)
519    }
520
521    fn log1p(&mut self, input: &Tensor) -> crate::Result<Tensor>;
522    fn log1p_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
523        self.log1p(read_tensor("log1p", input)?)
524    }
525}
526
527/// Shape, layout, and dtype transformation operations.
528///
529/// # Examples
530///
531/// ```rust
532/// use tenferro_tensor::TensorStructural;
533///
534/// fn accepts_structural<B: TensorStructural>(_backend: &mut B) {}
535/// ```
536pub trait TensorStructural {
537    fn transpose(&mut self, input: &Tensor, perm: &[usize]) -> crate::Result<Tensor>;
538    fn transpose_read(&mut self, input: TensorRead<'_>, perm: &[usize]) -> crate::Result<Tensor> {
539        self.transpose(read_tensor("transpose", input)?, perm)
540    }
541
542    fn reshape(&mut self, input: &Tensor, shape: &[usize]) -> crate::Result<Tensor>;
543    fn reshape_read(&mut self, input: TensorRead<'_>, shape: &[usize]) -> crate::Result<Tensor> {
544        self.reshape(read_tensor("reshape", input)?, shape)
545    }
546
547    fn broadcast_in_dim(
548        &mut self,
549        input: &Tensor,
550        shape: &[usize],
551        dims: &[usize],
552    ) -> crate::Result<Tensor>;
553    fn broadcast_in_dim_read(
554        &mut self,
555        input: TensorRead<'_>,
556        shape: &[usize],
557        dims: &[usize],
558    ) -> crate::Result<Tensor> {
559        self.broadcast_in_dim(read_tensor("broadcast_in_dim", input)?, shape, dims)
560    }
561
562    /// Cast a tensor to another dtype using explicit dtype projection.
563    ///
564    /// Backends may truncate, narrow precision, project complex values, or use
565    /// boolean truthiness according to their documented cast support.
566    ///
567    /// # Examples
568    ///
569    /// ```rust
570    /// use tenferro_tensor::{DType, Tensor, TensorStructural};
571    ///
572    /// fn cast_to_i32<B: TensorStructural>(
573    ///     backend: &mut B,
574    ///     input: &Tensor,
575    /// ) -> tenferro_tensor::Result<Tensor> {
576    ///     backend.cast(input, DType::I32)
577    /// }
578    /// ```
579    fn cast(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor>;
580
581    /// Convert a tensor to another dtype using checked dtype conversion.
582    ///
583    /// `convert` accepts only conversions allowed by tenferro's dtype-promotion
584    /// lattice. Use [`TensorStructural::cast`] for explicit lossy projection.
585    ///
586    /// # Examples
587    ///
588    /// ```rust
589    /// use tenferro_tensor::{DType, Tensor, TensorStructural};
590    ///
591    /// fn convert_to_f64<B: TensorStructural>(
592    ///     backend: &mut B,
593    ///     input: &Tensor,
594    /// ) -> tenferro_tensor::Result<Tensor> {
595    ///     backend.convert(input, DType::F64)
596    /// }
597    /// ```
598    fn convert(&mut self, input: &Tensor, to: crate::DType) -> crate::Result<Tensor> {
599        validate_convert_dtype("convert", input.dtype(), to)?;
600        self.cast(input, to)
601    }
602
603    fn extract_diagonal(
604        &mut self,
605        input: &Tensor,
606        axis_a: usize,
607        axis_b: usize,
608    ) -> crate::Result<Tensor>;
609    fn embed_diagonal(
610        &mut self,
611        input: &Tensor,
612        axis_a: usize,
613        axis_b: usize,
614    ) -> crate::Result<Tensor>;
615    fn tril(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor>;
616    fn triu(&mut self, input: &Tensor, k: i64) -> crate::Result<Tensor>;
617}
618
619/// Reduction operations.
620///
621/// Reducing over an axis whose extent is zero returns an error for every
622/// reduction operation. Passing an empty `axes` slice is a no-op and returns the
623/// input values unchanged.
624///
625/// # Examples
626///
627/// ```rust
628/// use tenferro_tensor::TensorReduction;
629///
630/// fn accepts_reduction<B: TensorReduction>(_backend: &mut B) {}
631/// ```
632pub trait TensorReduction {
633    fn reduce_sum(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
634
635    /// Sum elements across axes from an owned tensor or borrowed view.
636    ///
637    /// # Examples
638    ///
639    /// ```rust
640    /// use tenferro_tensor::{Tensor, TensorRead, TensorReduction};
641    ///
642    /// fn sum_owned<B: TensorReduction>(
643    ///     backend: &mut B,
644    ///     input: &Tensor,
645    /// ) -> tenferro_tensor::Result<Tensor> {
646    ///     backend.reduce_sum_read(TensorRead::from_tensor(input), &[0])
647    /// }
648    /// ```
649    fn reduce_sum_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
650        match input.as_tensor() {
651            Some(input) => self.reduce_sum(input, axes),
652            None => Err(crate::Error::backend_failure(
653                "reduce_sum",
654                "backend does not accept borrowed tensor views at this execution boundary",
655            )),
656        }
657    }
658
659    fn reduce_prod(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
660
661    /// Multiply elements across axes from an owned tensor or borrowed view.
662    ///
663    /// # Examples
664    ///
665    /// ```rust
666    /// use tenferro_tensor::{Tensor, TensorRead, TensorReduction};
667    ///
668    /// fn prod_owned<B: TensorReduction>(
669    ///     backend: &mut B,
670    ///     input: &Tensor,
671    /// ) -> tenferro_tensor::Result<Tensor> {
672    ///     backend.reduce_prod_read(TensorRead::from_tensor(input), &[0])
673    /// }
674    /// ```
675    fn reduce_prod_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
676        match input.as_tensor() {
677            Some(input) => self.reduce_prod(input, axes),
678            None => Err(crate::Error::backend_failure(
679                "reduce_prod",
680                "backend does not accept borrowed tensor views at this execution boundary",
681            )),
682        }
683    }
684
685    fn reduce_max(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
686
687    /// Take maximum values across axes from an owned tensor or borrowed view.
688    ///
689    /// # Examples
690    ///
691    /// ```rust
692    /// use tenferro_tensor::{Tensor, TensorRead, TensorReduction};
693    ///
694    /// fn max_owned<B: TensorReduction>(
695    ///     backend: &mut B,
696    ///     input: &Tensor,
697    /// ) -> tenferro_tensor::Result<Tensor> {
698    ///     backend.reduce_max_read(TensorRead::from_tensor(input), &[0])
699    /// }
700    /// ```
701    fn reduce_max_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
702        match input.as_tensor() {
703            Some(input) => self.reduce_max(input, axes),
704            None => Err(crate::Error::backend_failure(
705                "reduce_max",
706                "backend does not accept borrowed tensor views at this execution boundary",
707            )),
708        }
709    }
710
711    fn reduce_min(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
712
713    /// Take minimum values across axes from an owned tensor or borrowed view.
714    ///
715    /// # Examples
716    ///
717    /// ```rust
718    /// use tenferro_tensor::{Tensor, TensorRead, TensorReduction};
719    ///
720    /// fn min_owned<B: TensorReduction>(
721    ///     backend: &mut B,
722    ///     input: &Tensor,
723    /// ) -> tenferro_tensor::Result<Tensor> {
724    ///     backend.reduce_min_read(TensorRead::from_tensor(input), &[0])
725    /// }
726    /// ```
727    fn reduce_min_read(&mut self, input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor> {
728        match input.as_tensor() {
729            Some(input) => self.reduce_min(input, axes),
730            None => Err(crate::Error::backend_failure(
731                "reduce_min",
732                "backend does not accept borrowed tensor views at this execution boundary",
733            )),
734        }
735    }
736}
737
738/// Dot-general operations.
739///
740/// # Examples
741///
742/// ```rust
743/// use tenferro_tensor::TensorDot;
744///
745/// fn accepts_dot<B: TensorDot>(_backend: &mut B) {}
746/// ```
747pub trait TensorDot: TensorElementwise {
748    fn dot_general(
749        &mut self,
750        lhs: &Tensor,
751        rhs: &Tensor,
752        config: &DotGeneralConfig,
753    ) -> crate::Result<Tensor>;
754
755    #[doc(hidden)]
756    fn dot_general_read(
757        &mut self,
758        lhs: TensorRead<'_>,
759        rhs: TensorRead<'_>,
760        config: &DotGeneralConfig,
761    ) -> crate::Result<Tensor> {
762        match (lhs.as_tensor(), rhs.as_tensor()) {
763            (Some(lhs), Some(rhs)) => self.dot_general(lhs, rhs, config),
764            _ => {
765                let lhs = lhs.to_tensor()?;
766                let rhs = rhs.to_tensor()?;
767                self.dot_general(&lhs, &rhs, config)
768            }
769        }
770    }
771
772    #[doc(hidden)]
773    fn dot_general_read_into(
774        &mut self,
775        lhs: TensorRead<'_>,
776        rhs: TensorRead<'_>,
777        config: &DotGeneralConfig,
778        mut out: TensorWrite<'_>,
779    ) -> crate::Result<()> {
780        validate_dot_general_read_into(&lhs, &rhs, config, &out, "dot_general")?;
781        let result = self.dot_general_read(lhs, rhs, config)?;
782        out.copy_from_tensor(&result)
783    }
784
785    #[doc(hidden)]
786    fn dot_general_with_conj(
787        &mut self,
788        lhs: &Tensor,
789        rhs: &Tensor,
790        config: &DotGeneralConfig,
791        lhs_conj: bool,
792        rhs_conj: bool,
793    ) -> crate::Result<Tensor> {
794        if !lhs_conj && !rhs_conj {
795            return self.dot_general(lhs, rhs, config);
796        }
797
798        let lhs_tmp;
799        let lhs_ref = if lhs_conj {
800            lhs_tmp = self.conj(lhs)?;
801            &lhs_tmp
802        } else {
803            lhs
804        };
805        let rhs_tmp;
806        let rhs_ref = if rhs_conj {
807            rhs_tmp = self.conj(rhs)?;
808            &rhs_tmp
809        } else {
810            rhs
811        };
812        self.dot_general(lhs_ref, rhs_ref, config)
813    }
814
815    #[allow(clippy::too_many_arguments)]
816    #[doc(hidden)]
817    fn dot_general_with_conj_read(
818        &mut self,
819        lhs: TensorRead<'_>,
820        rhs: TensorRead<'_>,
821        config: &DotGeneralConfig,
822        lhs_conj: bool,
823        rhs_conj: bool,
824    ) -> crate::Result<Tensor> {
825        if !lhs_conj && !rhs_conj {
826            return self.dot_general_read(lhs, rhs, config);
827        }
828
829        let lhs_tmp;
830        let lhs_ref = if let Some(tensor) = lhs.as_tensor() {
831            tensor
832        } else {
833            lhs_tmp = lhs.to_tensor()?;
834            &lhs_tmp
835        };
836        let rhs_tmp;
837        let rhs_ref = if let Some(tensor) = rhs.as_tensor() {
838            tensor
839        } else {
840            rhs_tmp = rhs.to_tensor()?;
841            &rhs_tmp
842        };
843        self.dot_general_with_conj(lhs_ref, rhs_ref, config, lhs_conj, rhs_conj)
844    }
845}
846
847/// Session-scoped cached dot-general operations.
848///
849/// # Examples
850///
851/// ```rust
852/// use tenferro_tensor::BackendSession;
853///
854/// fn accepts_session_dot<S: BackendSession + ?Sized>(_session: &mut S) {}
855/// ```
856pub trait SessionCachedDot: TensorDot {
857    #[doc(hidden)]
858    fn dot_general_cached(
859        &mut self,
860        _cache_slot: Option<usize>,
861        lhs: &Tensor,
862        rhs: &Tensor,
863        config: &DotGeneralConfig,
864    ) -> crate::Result<Tensor> {
865        self.dot_general(lhs, rhs, config)
866    }
867
868    #[doc(hidden)]
869    fn dot_general_read_cached(
870        &mut self,
871        cache_slot: Option<usize>,
872        lhs: TensorRead<'_>,
873        rhs: TensorRead<'_>,
874        config: &DotGeneralConfig,
875    ) -> crate::Result<Tensor> {
876        match (lhs.as_tensor(), rhs.as_tensor()) {
877            (Some(lhs), Some(rhs)) => self.dot_general_cached(cache_slot, lhs, rhs, config),
878            _ => {
879                let lhs = lhs.to_tensor()?;
880                let rhs = rhs.to_tensor()?;
881                self.dot_general_cached(cache_slot, &lhs, &rhs, config)
882            }
883        }
884    }
885
886    // Mirrors the dot-general signature plus runtime-cache metadata.
887    #[allow(clippy::too_many_arguments)]
888    #[doc(hidden)]
889    fn dot_general_with_conj_cached(
890        &mut self,
891        _cache_slot: Option<usize>,
892        lhs: &Tensor,
893        rhs: &Tensor,
894        config: &DotGeneralConfig,
895        lhs_conj: bool,
896        rhs_conj: bool,
897    ) -> crate::Result<Tensor> {
898        self.dot_general_with_conj(lhs, rhs, config, lhs_conj, rhs_conj)
899    }
900
901    // Mirrors the dot-general read signature plus runtime-cache metadata.
902    #[allow(clippy::too_many_arguments)]
903    #[doc(hidden)]
904    fn dot_general_with_conj_read_cached(
905        &mut self,
906        cache_slot: Option<usize>,
907        lhs: TensorRead<'_>,
908        rhs: TensorRead<'_>,
909        config: &DotGeneralConfig,
910        lhs_conj: bool,
911        rhs_conj: bool,
912    ) -> crate::Result<Tensor> {
913        if !lhs_conj && !rhs_conj {
914            return self.dot_general_read_cached(cache_slot, lhs, rhs, config);
915        }
916
917        let lhs_tmp;
918        let lhs_ref = if let Some(tensor) = lhs.as_tensor() {
919            tensor
920        } else {
921            lhs_tmp = lhs.to_tensor()?;
922            &lhs_tmp
923        };
924        let rhs_tmp;
925        let rhs_ref = if let Some(tensor) = rhs.as_tensor() {
926            tensor
927        } else {
928            rhs_tmp = rhs.to_tensor()?;
929            &rhs_tmp
930        };
931        self.dot_general_with_conj_cached(cache_slot, lhs_ref, rhs_ref, config, lhs_conj, rhs_conj)
932    }
933}
934
935/// Indexing, slicing, and padding operations.
936///
937/// # Examples
938///
939/// ```rust
940/// use tenferro_tensor::TensorIndexing;
941///
942/// fn accepts_indexing<B: TensorIndexing>(_backend: &mut B) {}
943/// ```
944pub trait TensorIndexing {
945    fn gather(
946        &mut self,
947        operand: &Tensor,
948        start_indices: &Tensor,
949        config: &GatherConfig,
950    ) -> crate::Result<Tensor>;
951    fn scatter(
952        &mut self,
953        operand: &Tensor,
954        scatter_indices: &Tensor,
955        updates: &Tensor,
956        config: &ScatterConfig,
957    ) -> crate::Result<Tensor>;
958    fn slice(&mut self, input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor>;
959    fn dynamic_slice(
960        &mut self,
961        input: &Tensor,
962        starts: &Tensor,
963        slice_sizes: &[usize],
964    ) -> crate::Result<Tensor>;
965    fn dynamic_update_slice(
966        &mut self,
967        operand: &Tensor,
968        update: &Tensor,
969        starts: &Tensor,
970    ) -> crate::Result<Tensor>;
971    fn pad(&mut self, input: &Tensor, config: &PadConfig) -> crate::Result<Tensor>;
972    fn concatenate(&mut self, inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor>;
973    fn reverse(&mut self, input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
974}
975
976/// Backend-owned canonicalization for typed tensor views.
977///
978/// Implementations must preserve the input placement family. CPU backends
979/// canonicalize host views through explicit host copies and reject backend
980/// buffers with a diagnostic that asks the caller to download first. GPU
981/// backends canonicalize GPU-resident views on the same device and reject host
982/// buffers with an upload hint.
983///
984/// This trait is intentionally separate from [`BackendSession`] so generic
985/// typed methods do not change the object-safety contract of `dyn BackendSession`.
986///
987/// # Examples
988///
989/// ```rust
990/// use tenferro_tensor::{DynRank, TensorViewCanonicalization, TypedTensor};
991///
992/// fn compact_i32<B: TensorViewCanonicalization<i32, DynRank>>(
993///     backend: &mut B,
994///     tensor: &TypedTensor<i32>,
995/// ) -> tenferro_tensor::Result<TypedTensor<i32>> {
996///     backend.to_contiguous(&tensor.as_view())
997/// }
998/// ```
999pub trait TensorViewCanonicalization<T: Clone + 'static, R: TensorRank> {
1000    fn to_contiguous(
1001        &mut self,
1002        view: &TypedTensorView<'_, T, R>,
1003    ) -> crate::Result<TypedTensor<T, R>>;
1004
1005    fn copy_from_contiguous(
1006        &mut self,
1007        src: &TypedTensor<T, R>,
1008        dst: &mut TypedTensorViewMut<'_, T, R>,
1009    ) -> crate::Result<()>;
1010}
1011
1012/// Optional elementwise fusion execution.
1013///
1014/// # Examples
1015///
1016/// ```rust
1017/// use tenferro_tensor::TensorFusion;
1018///
1019/// fn accepts_fusion<B: TensorFusion>(_backend: &mut B) {}
1020/// ```
1021pub trait TensorFusion {
1022    #[doc(hidden)]
1023    fn execute_elementwise_fusion(
1024        &mut self,
1025        _inputs: &[&Tensor],
1026        _plan: &ElementwiseFusionPlan,
1027    ) -> crate::Result<Option<Vec<Tensor>>> {
1028        Ok(None)
1029    }
1030
1031    #[doc(hidden)]
1032    #[allow(clippy::too_many_arguments)]
1033    fn execute_broadcast_multiply(
1034        &mut self,
1035        _lhs: TensorRead<'_>,
1036        _lhs_shape: &[usize],
1037        _lhs_dims: &[usize],
1038        _rhs: TensorRead<'_>,
1039        _rhs_shape: &[usize],
1040        _rhs_dims: &[usize],
1041    ) -> crate::Result<Option<Tensor>> {
1042        Ok(None)
1043    }
1044
1045    #[doc(hidden)]
1046    #[allow(clippy::too_many_arguments)]
1047    fn execute_broadcast_multiply_value(
1048        &mut self,
1049        lhs: TensorRead<'_>,
1050        lhs_shape: &[usize],
1051        lhs_dims: &[usize],
1052        rhs: TensorRead<'_>,
1053        rhs_shape: &[usize],
1054        rhs_dims: &[usize],
1055    ) -> crate::Result<Option<TensorValue>> {
1056        self.execute_broadcast_multiply(lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims)
1057            .map(|tensor| tensor.map(TensorValue::from_tensor))
1058    }
1059}
1060
1061/// Backend buffer lifecycle operations.
1062///
1063/// # Examples
1064///
1065/// ```rust
1066/// use tenferro_tensor::TensorBuffer;
1067///
1068/// fn accepts_buffer<B: TensorBuffer>(_backend: &mut B) {}
1069/// ```
1070pub trait TensorBuffer {
1071    fn reclaim_buffer(&mut self, _tensor: Tensor) {}
1072}
1073
1074/// Device transfer operations on backend boundaries.
1075///
1076/// # Examples
1077///
1078/// ```rust
1079/// use tenferro_tensor::TensorDeviceTransfer;
1080///
1081/// fn accepts_transfer<B: TensorDeviceTransfer>(_backend: &mut B) {}
1082/// ```
1083pub trait TensorDeviceTransfer {
1084    fn download_to_host(&mut self, tensor: &Tensor) -> crate::Result<Tensor> {
1085        Ok(tensor.clone())
1086    }
1087
1088    fn upload_host_tensor(&mut self, tensor: &Tensor) -> crate::Result<Tensor> {
1089        Ok(tensor.clone())
1090    }
1091}
1092
1093/// Runtime cache associated with a backend.
1094///
1095/// # Examples
1096///
1097/// ```rust
1098/// use tenferro_tensor::BackendRuntimeCache;
1099///
1100/// fn accepts_runtime_cache<B: BackendRuntimeCache>(_backend: &B) {}
1101/// ```
1102pub trait BackendRuntimeCache {
1103    #[doc(hidden)]
1104    type RuntimeCache: RuntimeCacheControl + Send + Sync + 'static;
1105}
1106
1107/// Backend-owned cached dot-general operations.
1108///
1109/// # Examples
1110///
1111/// ```rust
1112/// use tenferro_tensor::BackendCachedDot;
1113///
1114/// fn accepts_backend_cached_dot<B: BackendCachedDot>(_backend: &mut B) {}
1115/// ```
1116pub trait BackendCachedDot: BackendRuntimeCache + TensorDot {
1117    #[doc(hidden)]
1118    fn dot_general_cached(
1119        &mut self,
1120        _cache: &mut Self::RuntimeCache,
1121        _cache_slot: Option<usize>,
1122        lhs: &Tensor,
1123        rhs: &Tensor,
1124        config: &DotGeneralConfig,
1125    ) -> crate::Result<Tensor> {
1126        self.dot_general(lhs, rhs, config)
1127    }
1128
1129    #[doc(hidden)]
1130    fn dot_general_read_cached(
1131        &mut self,
1132        cache: &mut Self::RuntimeCache,
1133        cache_slot: Option<usize>,
1134        lhs: TensorRead<'_>,
1135        rhs: TensorRead<'_>,
1136        config: &DotGeneralConfig,
1137    ) -> crate::Result<Tensor> {
1138        match (lhs.as_tensor(), rhs.as_tensor()) {
1139            (Some(lhs), Some(rhs)) => self.dot_general_cached(cache, cache_slot, lhs, rhs, config),
1140            _ => {
1141                let lhs = lhs.to_tensor()?;
1142                let rhs = rhs.to_tensor()?;
1143                self.dot_general_cached(cache, cache_slot, &lhs, &rhs, config)
1144            }
1145        }
1146    }
1147
1148    // Mirrors the dot-general signature plus runtime-cache metadata.
1149    #[allow(clippy::too_many_arguments)]
1150    #[doc(hidden)]
1151    fn dot_general_with_conj_cached(
1152        &mut self,
1153        _cache: &mut Self::RuntimeCache,
1154        _cache_slot: Option<usize>,
1155        lhs: &Tensor,
1156        rhs: &Tensor,
1157        config: &DotGeneralConfig,
1158        lhs_conj: bool,
1159        rhs_conj: bool,
1160    ) -> crate::Result<Tensor> {
1161        self.dot_general_with_conj(lhs, rhs, config, lhs_conj, rhs_conj)
1162    }
1163
1164    // Mirrors the dot-general read signature plus runtime-cache metadata.
1165    #[allow(clippy::too_many_arguments)]
1166    #[doc(hidden)]
1167    fn dot_general_with_conj_read_cached(
1168        &mut self,
1169        cache: &mut Self::RuntimeCache,
1170        cache_slot: Option<usize>,
1171        lhs: TensorRead<'_>,
1172        rhs: TensorRead<'_>,
1173        config: &DotGeneralConfig,
1174        lhs_conj: bool,
1175        rhs_conj: bool,
1176    ) -> crate::Result<Tensor> {
1177        if !lhs_conj && !rhs_conj {
1178            return self.dot_general_read_cached(cache, cache_slot, lhs, rhs, config);
1179        }
1180
1181        let lhs_tmp;
1182        let lhs_ref = if let Some(tensor) = lhs.as_tensor() {
1183            tensor
1184        } else {
1185            lhs_tmp = lhs.to_tensor()?;
1186            &lhs_tmp
1187        };
1188        let rhs_tmp;
1189        let rhs_ref = if let Some(tensor) = rhs.as_tensor() {
1190            tensor
1191        } else {
1192            rhs_tmp = rhs.to_tensor()?;
1193            &rhs_tmp
1194        };
1195        self.dot_general_with_conj_cached(
1196            cache, cache_slot, lhs_ref, rhs_ref, config, lhs_conj, rhs_conj,
1197        )
1198    }
1199}
1200
1201/// Backend execution-session entry points.
1202///
1203/// # Examples
1204///
1205/// ```rust
1206/// use tenferro_tensor::BackendSessionHost;
1207///
1208/// fn accepts_session_host<B: BackendSessionHost>(_backend: &mut B) {}
1209/// ```
1210pub trait BackendSessionHost: BackendRuntimeCache {
1211    fn with_backend_session<R: Send>(
1212        &mut self,
1213        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
1214    ) -> R
1215    where
1216        Self: TensorBackend + Sized,
1217    {
1218        default_backend_session(self, f)
1219    }
1220
1221    #[doc(hidden)]
1222    fn with_backend_session_cached<R: Send>(
1223        &mut self,
1224        _cache: &mut Self::RuntimeCache,
1225        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
1226    ) -> R
1227    where
1228        Self: TensorBackend + Sized,
1229    {
1230        self.with_backend_session(f)
1231    }
1232}
1233
1234/// Operation capabilities shared by backends and backend sessions.
1235#[doc(hidden)]
1236pub trait TensorBackendOps:
1237    TensorElementwise
1238    + TensorAnalytic
1239    + TensorStructural
1240    + TensorReduction
1241    + TensorIndexing
1242    + TensorDot
1243    + TensorFusion
1244    + TensorBuffer
1245{
1246}
1247
1248impl<T> TensorBackendOps for T where
1249    T: TensorElementwise
1250        + TensorAnalytic
1251        + TensorStructural
1252        + TensorReduction
1253        + TensorIndexing
1254        + TensorDot
1255        + TensorFusion
1256        + TensorBuffer
1257        + ?Sized
1258{
1259}
1260
1261/// Execution session surface for dense tensor backends.
1262///
1263/// All operations run within a backend-owned execution scope such as a CPU
1264/// thread policy or a GPU stream. Individual ops must not try to re-enter that
1265/// scope.
1266///
1267/// # Examples
1268///
1269/// ```rust
1270/// use tenferro_tensor::{BackendSessionHost, Tensor, TypedTensor};
1271///
1272/// fn add_in_session<B: BackendSessionHost>(
1273///     backend: &mut B,
1274///     a: &Tensor,
1275///     b: &Tensor,
1276/// ) -> tenferro_tensor::Result<Tensor>
1277/// where
1278///     B: tenferro_tensor::TensorBackend,
1279/// {
1280///     backend.with_backend_session(|exec| exec.add(a, b))
1281/// }
1282/// ```
1283pub trait BackendSession: TensorBackendOps + SessionCachedDot {}
1284
1285impl<T> BackendSession for T where T: TensorBackendOps + SessionCachedDot + ?Sized {}
1286
1287/// Standard runtime backend over dynamic [`Tensor`] values.
1288///
1289/// # Examples
1290///
1291/// ```rust
1292/// use tenferro_tensor::TensorBackend;
1293///
1294/// fn accepts_backend<B: TensorBackend>(_backend: &mut B) {}
1295/// ```
1296pub trait TensorBackend:
1297    BackendRuntimeCache
1298    + TensorBackendOps
1299    + BackendCachedDot
1300    + TensorDeviceTransfer
1301    + BackendSessionHost
1302{
1303}
1304
1305impl<T> SessionCachedDot for T where T: TensorBackend + ?Sized {}
1306
1307/// Run a closure using the backend itself as a default execution session.
1308///
1309/// This is suitable for backends whose individual ops already manage their own
1310/// execution context.
1311///
1312/// # Examples
1313///
1314/// ```rust
1315/// use tenferro_tensor::{default_backend_session, TensorBackend};
1316///
1317/// fn run_with_default_session<B: TensorBackend>(backend: &mut B) -> usize {
1318///     default_backend_session(backend, |_exec| 1usize)
1319/// }
1320/// ```
1321pub fn default_backend_session<B: TensorBackend, R: Send>(
1322    backend: &mut B,
1323    f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
1324) -> R {
1325    f(backend)
1326}
tenferro_tensor/backend.rs

tenferro_tensor/
backend.rs
