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

tenferro_tensor/
backend.rs
