tenferro_runtime/
tensor.rs

1//! Concrete tensor operation extension trait.
2//!
3//! `tenferro-tensor` owns storage and backend traits. This runtime crate
4//! provides backend-parametric session-explicit operation methods through
5//! [`TensorSessionOpsExt`].
6
7use crate::composite;
8use crate::composite::session::{borrowed, run_session_composite};
9use std::borrow::Cow;
10
11use num_complex::Complex64;
12use tenferro_ops::broadcast::{broadcast_error_to_validation, broadcast_shape, broadcast_shapes};
13use tenferro_tensor::validate::matmul_config_for_shapes;
14use tenferro_tensor::{
15    ActivationOp, CompareDir, DType, DotGeneralConfig, Engine, Error, GatherConfig, PadConfig,
16    Result, ScatterConfig, SliceConfig, TensorRead,
17};
18
19use crate::typed_tensor::{broadcast_to_in_read, ReadInput};
20
21use crate::TensorSessionOpsExt;
22use tenferro_tensor::Tensor;
23
24impl TensorSessionOpsExt for Tensor {
25    fn add(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
26        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
27        session.add_read(lhs.tensor_read(), rhs.tensor_read())
28    }
29
30    fn mul(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
31        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
32        session.mul_read(lhs.tensor_read(), rhs.tensor_read())
33    }
34
35    fn exp(&self, session: &mut dyn Engine) -> Result<Tensor> {
36        session.exp_read(TensorRead::from_tensor(self))
37    }
38
39    fn reduce_sum(&self, axes: Option<&[usize]>, session: &mut dyn Engine) -> Result<Tensor> {
40        let axes = all_axes_if_none(self.shape().len(), axes);
41        session.reduce_sum_read(TensorRead::from_tensor(self), &axes)
42    }
43
44    fn convert(&self, to: DType, session: &mut dyn Engine) -> Result<Tensor> {
45        session.convert(self, to)
46    }
47
48    fn cast(&self, to: DType, session: &mut dyn Engine) -> Result<Tensor> {
49        session.cast(self, to)
50    }
51
52    fn sub(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
53        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
54        session.sub_read(lhs.tensor_read(), rhs.tensor_read())
55    }
56
57    fn div(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
58        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
59        session.div_read(lhs.tensor_read(), rhs.tensor_read())
60    }
61
62    fn rem(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
63        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
64        session.rem_read(lhs.tensor_read(), rhs.tensor_read())
65    }
66
67    fn pow(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
68        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
69        session.pow_read(lhs.tensor_read(), rhs.tensor_read())
70    }
71
72    fn maximum(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
73        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
74        session.maximum_read(lhs.tensor_read(), rhs.tensor_read())
75    }
76
77    fn minimum(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
78        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
79        session.minimum_read(lhs.tensor_read(), rhs.tensor_read())
80    }
81
82    fn neg(&self, session: &mut dyn Engine) -> Result<Tensor> {
83        session.neg_read(TensorRead::from_tensor(self))
84    }
85
86    fn abs(&self, session: &mut dyn Engine) -> Result<Tensor> {
87        session.abs_read(TensorRead::from_tensor(self))
88    }
89
90    fn sign(&self, session: &mut dyn Engine) -> Result<Tensor> {
91        session.sign_read(TensorRead::from_tensor(self))
92    }
93
94    fn conj(&self, session: &mut dyn Engine) -> Result<Tensor> {
95        session.conj_read(TensorRead::from_tensor(self))
96    }
97
98    fn log(&self, session: &mut dyn Engine) -> Result<Tensor> {
99        session.log_read(TensorRead::from_tensor(self))
100    }
101
102    fn expm1(&self, session: &mut dyn Engine) -> Result<Tensor> {
103        session.expm1_read(TensorRead::from_tensor(self))
104    }
105
106    fn log1p(&self, session: &mut dyn Engine) -> Result<Tensor> {
107        session.log1p_read(TensorRead::from_tensor(self))
108    }
109
110    fn erf(&self, session: &mut dyn Engine) -> Result<Tensor> {
111        session.erf_read(TensorRead::from_tensor(self))
112    }
113
114    fn sin(&self, session: &mut dyn Engine) -> Result<Tensor> {
115        session.sin_read(TensorRead::from_tensor(self))
116    }
117
118    fn cos(&self, session: &mut dyn Engine) -> Result<Tensor> {
119        session.cos_read(TensorRead::from_tensor(self))
120    }
121
122    fn tanh(&self, session: &mut dyn Engine) -> Result<Tensor> {
123        session.tanh_read(TensorRead::from_tensor(self))
124    }
125
126    fn sqrt(&self, session: &mut dyn Engine) -> Result<Tensor> {
127        session.sqrt_read(TensorRead::from_tensor(self))
128    }
129
130    fn rsqrt(&self, session: &mut dyn Engine) -> Result<Tensor> {
131        session.rsqrt_read(TensorRead::from_tensor(self))
132    }
133
134    fn compare(&self, rhs: &Tensor, dir: CompareDir, session: &mut dyn Engine) -> Result<Tensor> {
135        let (lhs, rhs) = broadcast_binary_in(self, rhs, session)?;
136        session.compare_read(lhs.tensor_read(), rhs.tensor_read(), &dir)
137    }
138
139    fn where_select(
140        &self,
141        on_true: &Tensor,
142        on_false: &Tensor,
143        session: &mut dyn Engine,
144    ) -> Result<Tensor> {
145        let (condition, on_true, on_false) =
146            broadcast_ternary_in(self, on_true, on_false, session)?;
147        session.select_read(
148            condition.tensor_read(),
149            on_true.tensor_read(),
150            on_false.tensor_read(),
151        )
152    }
153
154    fn clamp(&self, lower: &Tensor, upper: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
155        let (input, lower, upper) = broadcast_ternary_in(self, lower, upper, session)?;
156        session.clamp_read(
157            input.tensor_read(),
158            lower.tensor_read(),
159            upper.tensor_read(),
160        )
161    }
162
163    fn matmul(&self, rhs: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
164        let config = matmul_config_for_shapes("matmul", self.shape(), rhs.shape())?;
165        session.dot_general_read(
166            TensorRead::from_tensor(self),
167            TensorRead::from_tensor(rhs),
168            &config,
169        )
170    }
171
172    fn reshape(&self, shape: &[usize], session: &mut dyn Engine) -> Result<Tensor> {
173        session.reshape_read(TensorRead::from_tensor(self), shape)
174    }
175
176    fn transpose(&self, perm: &[usize], session: &mut dyn Engine) -> Result<Tensor> {
177        session.transpose_read(TensorRead::from_tensor(self), perm)
178    }
179
180    fn gather(
181        &self,
182        indices: &Tensor,
183        config: GatherConfig,
184        session: &mut dyn Engine,
185    ) -> Result<Tensor> {
186        session.gather(self, indices, &config)
187    }
188
189    fn scatter(
190        &self,
191        indices: &Tensor,
192        updates: &Tensor,
193        config: ScatterConfig,
194        session: &mut dyn Engine,
195    ) -> Result<Tensor> {
196        session.scatter(self, indices, updates, &config)
197    }
198
199    fn slice(&self, config: SliceConfig, session: &mut dyn Engine) -> Result<Tensor> {
200        session.slice(self, &config)
201    }
202
203    fn dynamic_slice(
204        &self,
205        starts: &Tensor,
206        sizes: &[usize],
207        session: &mut dyn Engine,
208    ) -> Result<Tensor> {
209        session.dynamic_slice(self, starts, sizes)
210    }
211
212    fn dynamic_update_slice(
213        &self,
214        update: &Tensor,
215        starts: &Tensor,
216        session: &mut dyn Engine,
217    ) -> Result<Tensor> {
218        session.dynamic_update_slice(self, update, starts)
219    }
220
221    fn pad(&self, config: PadConfig, session: &mut dyn Engine) -> Result<Tensor> {
222        session.pad(self, &config)
223    }
224
225    fn concatenate(inputs: &[&Tensor], axis: usize, session: &mut dyn Engine) -> Result<Tensor> {
226        session.concatenate(inputs, axis)
227    }
228
229    fn reverse(&self, axes: &[usize], session: &mut dyn Engine) -> Result<Tensor> {
230        session.reverse(self, axes)
231    }
232
233    fn reduce_max(&self, axes: Option<&[usize]>, session: &mut dyn Engine) -> Result<Tensor> {
234        let axes = all_axes_if_none(self.shape().len(), axes);
235        session.reduce_max_read(TensorRead::from_tensor(self), &axes)
236    }
237
238    fn reduce_min(&self, axes: Option<&[usize]>, session: &mut dyn Engine) -> Result<Tensor> {
239        let axes = all_axes_if_none(self.shape().len(), axes);
240        session.reduce_min_read(TensorRead::from_tensor(self), &axes)
241    }
242
243    fn reduce_prod(&self, axes: Option<&[usize]>, session: &mut dyn Engine) -> Result<Tensor> {
244        let axes = all_axes_if_none(self.shape().len(), axes);
245        session.reduce_prod_read(TensorRead::from_tensor(self), &axes)
246    }
247
248    fn reduce_sum_squares(
249        &self,
250        axes: Option<&[usize]>,
251        session: &mut dyn Engine,
252    ) -> Result<Tensor> {
253        let axes = all_axes_if_none(self.shape().len(), axes);
254        session.reduce_sum_squares_read(TensorRead::from_tensor(self), &axes)
255    }
256
257    fn broadcast_in_dim(
258        &self,
259        shape: &[usize],
260        dims: &[usize],
261        session: &mut dyn Engine,
262    ) -> Result<Tensor> {
263        session.broadcast_in_dim_read(TensorRead::from_tensor(self), shape, dims)
264    }
265
266    fn tril(&self, k: i64, session: &mut dyn Engine) -> Result<Tensor> {
267        session.tril(self, k)
268    }
269
270    fn triu(&self, k: i64, session: &mut dyn Engine) -> Result<Tensor> {
271        session.triu(self, k)
272    }
273
274    fn extract_diag(
275        &self,
276        axis_a: usize,
277        axis_b: usize,
278        session: &mut dyn Engine,
279    ) -> Result<Tensor> {
280        session.extract_diagonal(self, axis_a, axis_b)
281    }
282
283    fn embed_diag(&self, axis_a: usize, axis_b: usize, session: &mut dyn Engine) -> Result<Tensor> {
284        session.embed_diagonal(self, axis_a, axis_b)
285    }
286
287    fn dot_general(
288        &self,
289        rhs: &Tensor,
290        config: DotGeneralConfig,
291        session: &mut dyn Engine,
292    ) -> Result<Tensor> {
293        session.dot_general_read(
294            TensorRead::from_tensor(self),
295            TensorRead::from_tensor(rhs),
296            &config,
297        )
298    }
299
300    fn dot_general_with_conj(
301        &self,
302        rhs: &Tensor,
303        config: DotGeneralConfig,
304        lhs_conj: bool,
305        rhs_conj: bool,
306        session: &mut dyn Engine,
307    ) -> Result<Tensor> {
308        session.dot_general_with_conj(self, rhs, &config, lhs_conj, rhs_conj)
309    }
310
311    fn scale_real(&self, factor: f64, session: &mut dyn Engine) -> Result<Tensor> {
312        let scalar = crate::scale::real_scale_scalar(self.dtype(), factor)?;
313        let scalar = session.upload_host_tensor(TensorRead::from_tensor(&scalar))?;
314        TensorSessionOpsExt::mul(self, &scalar, session)
315    }
316
317    fn scale_complex(&self, factor: Complex64, session: &mut dyn Engine) -> Result<Tensor> {
318        let scalar = crate::scale::complex_scale_scalar(self.dtype(), factor)?;
319        let scalar = session.upload_host_tensor(TensorRead::from_tensor(&scalar))?;
320        TensorSessionOpsExt::mul(self, &scalar, session)
321    }
322
323    fn sigmoid(&self, session: &mut dyn Engine) -> Result<Tensor> {
324        if let Some(out) =
325            session.fused_activation_read(ActivationOp::Sigmoid, TensorRead::from_tensor(self))?
326        {
327            return Ok(out);
328        }
329        run_session_composite(session, |ops| composite::sigmoid(ops, &borrowed(self)))
330    }
331
332    fn silu(&self, session: &mut dyn Engine) -> Result<Tensor> {
333        if let Some(out) =
334            session.fused_activation_read(ActivationOp::Silu, TensorRead::from_tensor(self))?
335        {
336            return Ok(out);
337        }
338        run_session_composite(session, |ops| composite::silu(ops, &borrowed(self)))
339    }
340
341    fn softplus(&self, session: &mut dyn Engine) -> Result<Tensor> {
342        if let Some(out) =
343            session.fused_activation_read(ActivationOp::Softplus, TensorRead::from_tensor(self))?
344        {
345            return Ok(out);
346        }
347        run_session_composite(session, |ops| composite::softplus(ops, &borrowed(self)))
348    }
349
350    fn gelu(&self, session: &mut dyn Engine) -> Result<Tensor> {
351        if let Some(out) =
352            session.fused_activation_read(ActivationOp::Gelu, TensorRead::from_tensor(self))?
353        {
354            return Ok(out);
355        }
356        run_session_composite(session, |ops| composite::gelu(ops, &borrowed(self)))
357    }
358
359    fn gelu_tanh(&self, session: &mut dyn Engine) -> Result<Tensor> {
360        if let Some(out) =
361            session.fused_activation_read(ActivationOp::GeluTanh, TensorRead::from_tensor(self))?
362        {
363            return Ok(out);
364        }
365        run_session_composite(session, |ops| composite::gelu_tanh(ops, &borrowed(self)))
366    }
367
368    fn reduce_mean(&self, axes: Option<&[usize]>, session: &mut dyn Engine) -> Result<Tensor> {
369        run_session_composite(session, |ops| {
370            composite::reduce_mean(ops, &borrowed(self), axes)
371        })
372    }
373
374    fn softmax(&self, axis: usize, session: &mut dyn Engine) -> Result<Tensor> {
375        if let Some(output) = session.softmax_read(TensorRead::from_tensor(self), axis)? {
376            return Ok(output);
377        }
378        run_session_composite(session, |ops| {
379            composite::softmax(ops, &borrowed(self), axis)
380        })
381    }
382
383    fn log_softmax(&self, axis: usize, session: &mut dyn Engine) -> Result<Tensor> {
384        run_session_composite(session, |ops| {
385            composite::log_softmax(ops, &borrowed(self), axis)
386        })
387    }
388
389    fn masked_softmax(
390        &self,
391        mask: &Tensor,
392        axis: usize,
393        session: &mut dyn Engine,
394    ) -> Result<Tensor> {
395        let mask = borrowed(mask);
396        run_session_composite(session, |ops| {
397            composite::masked_softmax(ops, &borrowed(self), &mask, axis)
398        })
399    }
400
401    fn masked_log_softmax(
402        &self,
403        mask: &Tensor,
404        axis: usize,
405        session: &mut dyn Engine,
406    ) -> Result<Tensor> {
407        let mask = borrowed(mask);
408        run_session_composite(session, |ops| {
409            composite::masked_log_softmax(ops, &borrowed(self), &mask, axis)
410        })
411    }
412
413    fn layer_norm(
414        &self,
415        axis: usize,
416        weight: Option<&Tensor>,
417        bias: Option<&Tensor>,
418        eps: f64,
419        session: &mut dyn Engine,
420    ) -> Result<Tensor> {
421        let weight = weight.map(borrowed);
422        let bias = bias.map(borrowed);
423        run_session_composite(session, |ops| {
424            composite::layer_norm(
425                ops,
426                &borrowed(self),
427                axis,
428                weight.as_ref(),
429                bias.as_ref(),
430                eps,
431            )
432        })
433    }
434
435    fn rms_norm(
436        &self,
437        axis: usize,
438        weight: Option<&Tensor>,
439        bias: Option<&Tensor>,
440        eps: f64,
441        session: &mut dyn Engine,
442    ) -> Result<Tensor> {
443        let weight = weight.map(borrowed);
444        let bias = bias.map(borrowed);
445        run_session_composite(session, |ops| {
446            composite::rms_norm(
447                ops,
448                &borrowed(self),
449                axis,
450                weight.as_ref(),
451                bias.as_ref(),
452                eps,
453            )
454        })
455    }
456
457    fn take_along_axis(
458        &self,
459        indices: &Tensor,
460        axis: usize,
461        session: &mut dyn Engine,
462    ) -> Result<Tensor> {
463        let indices = borrowed(indices);
464        run_session_composite(session, |ops| {
465            composite::take_along_axis(ops, &borrowed(self), &indices, axis)
466        })
467    }
468}
469
470fn broadcast_binary_in<'a>(
471    lhs: &'a Tensor,
472    rhs: &'a Tensor,
473    session: &mut dyn Engine,
474) -> Result<(ReadInput<'a>, ReadInput<'a>)> {
475    let shape = broadcast_shape(lhs.shape(), rhs.shape()).map_err(broadcast_error)?;
476    Ok((
477        broadcast_to_in_read(TensorRead::from_tensor(lhs), &shape, session)?,
478        broadcast_to_in_read(TensorRead::from_tensor(rhs), &shape, session)?,
479    ))
480}
481
482fn broadcast_ternary_in<'a>(
483    first: &'a Tensor,
484    second: &'a Tensor,
485    third: &'a Tensor,
486    session: &mut dyn Engine,
487) -> Result<(ReadInput<'a>, ReadInput<'a>, ReadInput<'a>)> {
488    let shape = broadcast_shapes([first.shape(), second.shape(), third.shape()])
489        .map_err(broadcast_error)?;
490    Ok((
491        broadcast_to_in_read(TensorRead::from_tensor(first), &shape, session)?,
492        broadcast_to_in_read(TensorRead::from_tensor(second), &shape, session)?,
493        broadcast_to_in_read(TensorRead::from_tensor(third), &shape, session)?,
494    ))
495}
496
497fn broadcast_error(err: tenferro_ops::broadcast::BroadcastError) -> Error {
498    Error::validation("broadcast", broadcast_error_to_validation(err))
499}
500
501/// Resolve a reduction-family axis argument: `None` selects every axis.
502///
503/// An explicit axis list stays borrowed; only `None` builds a list.
504pub(crate) fn all_axes_if_none(rank: usize, axes: Option<&[usize]>) -> Cow<'_, [usize]> {
505    match axes {
506        Some(axes) => Cow::Borrowed(axes),
507        None => Cow::Owned((0..rank).collect()),
508    }
509}
tenferro_runtime/tensor.rs

tenferro_runtime/
tensor.rs
