luma_tensor/device/
float_ops.rs

1use std::borrow::Cow;
2
3use crate::Result;
4use crate::dtype::{BoolDType, FloatDType, IntDType};
5use crate::tensor::{Layout, Shape};
6
7/// Operations for floating-point tensors.
8pub trait FloatOps<D: super::Device> {
9    // construction
10    fn f_zeros(shape: &Shape, device: &D, dtype: FloatDType) -> Result<D::FloatStorage>;
11    fn f_ones(shape: &Shape, device: &D, dtype: FloatDType) -> Result<D::FloatStorage>;
12    fn f_full(shape: &Shape, value: f64, device: &D, dtype: FloatDType) -> Result<D::FloatStorage>;
13    fn f_from_f64<'a>(data: impl Into<Cow<'a, [f64]>>, device: &D) -> Result<D::FloatStorage>;
14
15    fn f_from_f32<'a>(data: impl Into<Cow<'a, [f32]>>, device: &D) -> Result<D::FloatStorage>;
16
17    fn f_from_bytes<'a>(bytes: impl Into<Cow<'a, [u8]>>, shape: &Shape, device: &D, dtype: FloatDType) -> Result<D::FloatStorage>;
18    fn f_rand_uniform(shape: &Shape, lo: f64, hi: f64, device: &D, dtype: FloatDType) -> Result<D::FloatStorage>;
19    fn f_rand_normal(shape: &Shape, mean: f64, std: f64, device: &D, dtype: FloatDType) -> Result<D::FloatStorage>;
20
21    // materialization / read-back
22    fn f_contiguous(x: &D::FloatStorage, layout: &Layout) -> Result<D::FloatStorage>;
23    fn f_cast_float(x: &D::FloatStorage, layout: &Layout, to: FloatDType) -> Result<D::FloatStorage>;
24    fn f_cast_int(x: &D::FloatStorage, layout: &Layout, to: IntDType) -> Result<D::IntStorage>;
25    fn f_cast_bool(x: &D::FloatStorage, layout: &Layout, to: BoolDType) -> Result<D::BoolStorage>;
26
27    /// Read all elements into a `Vec<f64>` in logical (layout) order.
28    fn f_to_vec(x: &D::FloatStorage, layout: &Layout) -> Result<Vec<f64>>;
29
30    /// Read raw little-endian bytes in logical (layout) order.
31    /// Returns `Cow::Borrowed` when the underlying storage is already contiguous
32    /// (zero-copy); `Cow::Owned` otherwise.
33    fn f_to_bytes<'a>(x: &'a D::FloatStorage, layout: &Layout) -> Result<Cow<'a, [u8]>>;
34
35    // binary (elementwise, same shape; broadcasting handled above this layer)
36    fn f_binary(
37        lhs: &D::FloatStorage,
38        lhs_l: &Layout,
39        rhs: &D::FloatStorage,
40        rhs_l: &Layout,
41        op: crate::BinaryOp,
42    ) -> Result<D::FloatStorage>;
43    fn f_binary_(dst: &mut D::FloatStorage, dst_l: &Layout, src: &D::FloatStorage, src_l: &Layout, op: crate::BinaryOp) -> Result<()>;
44    fn f_binary_scalar(lhs: &D::FloatStorage, lhs_l: &Layout, rhs: f64, op: crate::BinaryOp) -> Result<D::FloatStorage>;
45    fn f_binary_scalar_(dst: &mut D::FloatStorage, dst_l: &Layout, rhs: f64, op: crate::BinaryOp) -> Result<()>;
46    fn f_binary_scalar_lhs(scalar: f64, rhs: &D::FloatStorage, rhs_l: &Layout, op: crate::BinaryOp) -> Result<D::FloatStorage>;
47
48    // comparison -> bool storage
49    fn f_cmp(lhs: &D::FloatStorage, lhs_l: &Layout, rhs: &D::FloatStorage, rhs_l: &Layout, op: crate::CmpOp) -> Result<D::BoolStorage>;
50    fn f_cmp_scalar(lhs: &D::FloatStorage, lhs_l: &Layout, rhs: f64, op: crate::CmpOp) -> Result<D::BoolStorage>;
51
52    // unary
53    fn f_unary(x: &D::FloatStorage, layout: &Layout, op: crate::UnaryOp<f64>) -> Result<D::FloatStorage>;
54    fn f_unary_(dst: &mut D::FloatStorage, dst_l: &Layout, op: crate::UnaryOp<f64>) -> Result<()>;
55    fn f_float_unary(x: &D::FloatStorage, layout: &Layout, op: crate::FloatUnaryOp) -> Result<D::FloatStorage>;
56    fn f_float_unary_(dst: &mut D::FloatStorage, dst_l: &Layout, op: crate::FloatUnaryOp) -> Result<()>;
57
58    // reduction: returns (storage, resulting shape)
59    fn f_reduce(
60        x: &D::FloatStorage,
61        layout: &Layout,
62        dims: &[usize],
63        keepdim: bool,
64        op: crate::ReduceOp,
65    ) -> Result<(D::FloatStorage, Shape)>;
66    fn f_arg_reduce(x: &D::FloatStorage, layout: &Layout, dim: usize, keepdim: bool, take_max: bool) -> Result<(D::IntStorage, Shape)>;
67
68    // matmul (batched); out shape computed by the caller / this fn
69    fn f_matmul(lhs: &D::FloatStorage, lhs_l: &Layout, rhs: &D::FloatStorage, rhs_l: &Layout) -> Result<(D::FloatStorage, Shape)>;
70    fn f_add_matmul_(
71        dst: &mut D::FloatStorage,
72        dst_l: &Layout,
73        lhs: &D::FloatStorage,
74        lhs_l: &Layout,
75        rhs: &D::FloatStorage,
76        rhs_l: &Layout,
77    ) -> Result<()>;
78
79    // indexing
80    fn f_index_select(
81        x: &D::FloatStorage,
82        x_l: &Layout,
83        idx: &D::IntStorage,
84        idx_l: &Layout,
85        dim: usize,
86    ) -> Result<(D::FloatStorage, Shape)>;
87
88    fn f_gather(x: &D::FloatStorage, x_l: &Layout, idx: &D::IntStorage, idx_l: &Layout, dim: usize) -> Result<(D::FloatStorage, Shape)>;
89
90    fn f_index_add(
91        init: &D::FloatStorage,
92        init_l: &Layout,
93        idx: &D::IntStorage,
94        idx_l: &Layout,
95        src: &D::FloatStorage,
96        src_l: &Layout,
97        dim: usize,
98    ) -> Result<D::FloatStorage>;
99
100    fn f_scatter_add(
101        init: &D::FloatStorage,
102        init_l: &Layout,
103        idx: &D::IntStorage,
104        idx_l: &Layout,
105        src: &D::FloatStorage,
106        src_l: &Layout,
107        dim: usize,
108    ) -> Result<D::FloatStorage>;
109
110    // shape ops that need data movement
111    fn f_cat(srcs: &[(&D::FloatStorage, &Layout)], dim: usize) -> Result<(D::FloatStorage, Shape)>;
112
113    // views (alias on compute devices; a fresh SSA value on tracing devices)
114    fn f_view(_src: &D::FloatStorage, _src_l: &Layout, _dst_l: &Layout, _view: crate::ViewOp) -> Result<Option<D::FloatStorage>> {
115        Ok(None)
116    }
117
118    // nn fused kernels
119    fn f_softmax(x: &D::FloatStorage, layout: &Layout, dim: usize) -> Result<D::FloatStorage>;
120
121    fn f_rms_norm(x: &D::FloatStorage, x_l: &Layout, weight: &D::FloatStorage, weight_l: &Layout, eps: f64) -> Result<D::FloatStorage>;
122
123    // pick via a bool mask
124    fn f_pick(
125        mask: &D::BoolStorage,
126        mask_l: &Layout,
127        on_true: &D::FloatStorage,
128        true_l: &Layout,
129        on_false: &D::FloatStorage,
130        false_l: &Layout,
131    ) -> Result<D::FloatStorage>;
132
133    fn f_pick_true(
134        mask: &D::BoolStorage,
135        mask_l: &Layout,
136        value: f64,
137        on_false: &D::FloatStorage,
138        false_l: &Layout,
139    ) -> Result<D::FloatStorage>;
140
141    fn f_pick_false(
142        mask: &D::BoolStorage,
143        mask_l: &Layout,
144        on_true: &D::FloatStorage,
145        true_l: &Layout,
146        value: f64,
147    ) -> Result<D::FloatStorage>;
148
149    // allclose
150    fn f_allclose(a: &D::FloatStorage, a_l: &Layout, b: &D::FloatStorage, b_l: &Layout, rtol: f64, atol: f64) -> Result<bool>;
151}
luma_tensor/device/float_ops.rs

luma_tensor/device/
float_ops.rs
