luma_tensor/device/cpu/ops/
float_ops.rs

1//! `impl FloatOps for Cpu`: dispatch `CpuFloatStorage` variants to generic kernels.
2
3use std::borrow::Cow;
4
5use rand::rng;
6use rand_distr::{Distribution, Normal, Uniform};
7
8use super::kernels::{elementwise as ew, indexing, matmul, nn, reduce};
9use super::{Cpu, CpuBoolStorage, CpuFloatStorage, CpuIntStorage, int_ids_as_usize, usize_to_int_storage};
10use crate::dtype::{BoolDType, FloatDType, IntDType};
11use crate::{
12    BinaryOp, CmpOp, DType, Device, Error, FloatOps, Layout, Result, Shape, dispatch_float, dispatch_float_raw, dispatch_float2,
13    dispatch_float2_raw,
14};
15
16/// Build a float storage of `dtype`, filling `n` elements via `f32`/`f64` closures.
17fn build(n: usize, dtype: FloatDType, f32v: impl Fn() -> f32, f64v: impl Fn() -> f64) -> CpuFloatStorage {
18    match dtype {
19        FloatDType::F32 => CpuFloatStorage::F32((0..n).map(|_| f32v()).collect()),
20        FloatDType::F64 => CpuFloatStorage::F64((0..n).map(|_| f64v()).collect()),
21    }
22}
23
24impl FloatOps<Cpu> for Cpu {
25    fn f_zeros(shape: &Shape, _device: &Cpu, dtype: FloatDType) -> Result<<Cpu as Device>::FloatStorage> {
26        Ok(build(shape.element_count(), dtype, || 0.0, || 0.0))
27    }
28
29    fn f_ones(shape: &Shape, _device: &Cpu, dtype: FloatDType) -> Result<<Cpu as Device>::FloatStorage> {
30        Ok(build(shape.element_count(), dtype, || 1.0, || 1.0))
31    }
32
33    fn f_full(shape: &Shape, value: f64, _device: &Cpu, dtype: FloatDType) -> Result<<Cpu as Device>::FloatStorage> {
34        Ok(build(shape.element_count(), dtype, || value as f32, || value))
35    }
36
37    fn f_from_f64<'a>(data: impl Into<Cow<'a, [f64]>>, _device: &Cpu) -> Result<<Cpu as Device>::FloatStorage> {
38        let data = data.into();
39        Ok(match data {
40            Cow::Owned(v) => CpuFloatStorage::F64(v),
41            Cow::Borrowed(s) => CpuFloatStorage::F64(s.to_vec()),
42        })
43    }
44
45    fn f_from_f32<'a>(data: impl Into<Cow<'a, [f32]>>, _device: &Cpu) -> Result<<Cpu as Device>::FloatStorage> {
46        let data = data.into();
47        Ok(match data {
48            Cow::Owned(v) => CpuFloatStorage::F32(v),
49            Cow::Borrowed(s) => CpuFloatStorage::F32(s.to_vec()),
50        })
51    }
52
53    fn f_from_bytes<'a>(
54        bytes: impl Into<Cow<'a, [u8]>>,
55        _shape: &Shape,
56        _device: &Cpu,
57        dtype: FloatDType,
58    ) -> Result<<Cpu as Device>::FloatStorage> {
59        let bytes = bytes.into();
60        Ok(match dtype {
61            FloatDType::F32 => {
62                let v: Vec<f32> = bytes.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect();
63                CpuFloatStorage::F32(v)
64            }
65            FloatDType::F64 => {
66                let v: Vec<f64> = bytes.chunks_exact(8).map(|c| f64::from_le_bytes(c.try_into().unwrap())).collect();
67                CpuFloatStorage::F64(v)
68            }
69        })
70    }
71
72    fn f_rand_uniform(shape: &Shape, lo: f64, hi: f64, _device: &Cpu, dtype: FloatDType) -> Result<<Cpu as Device>::FloatStorage> {
73        let n = shape.element_count();
74        let mut r = rng();
75        let s = match dtype {
76            FloatDType::F64 => {
77                let u = Uniform::new(lo, hi).map_err(|e| Error::Rand(e.to_string()))?;
78                CpuFloatStorage::F64((0..n).map(|_| u.sample(&mut r)).collect())
79            }
80            FloatDType::F32 => {
81                let u = Uniform::new(lo as f32, hi as f32).map_err(|e| Error::Rand(e.to_string()))?;
82                CpuFloatStorage::F32((0..n).map(|_| u.sample(&mut r)).collect())
83            }
84        };
85        Ok(s)
86    }
87
88    fn f_rand_normal(shape: &Shape, mean: f64, std: f64, _device: &Cpu, dtype: FloatDType) -> Result<<Cpu as Device>::FloatStorage> {
89        let n = shape.element_count();
90        let mut r = rng();
91        let s = match dtype {
92            FloatDType::F64 => {
93                let d = Normal::new(mean, std).map_err(|e| Error::Rand(e.to_string()))?;
94                CpuFloatStorage::F64((0..n).map(|_| d.sample(&mut r)).collect())
95            }
96            FloatDType::F32 => {
97                let d = Normal::new(mean as f32, std as f32).map_err(|e| Error::Rand(e.to_string()))?;
98                CpuFloatStorage::F32((0..n).map(|_| d.sample(&mut r)).collect())
99            }
100        };
101        Ok(s)
102    }
103
104    fn f_contiguous(x: &<Cpu as Device>::FloatStorage, l: &Layout) -> Result<<Cpu as Device>::FloatStorage> {
105        Ok(dispatch_float!(x, |d| super::kernels::iter::gather(d, l)))
106    }
107
108    fn f_cast_float(x: &CpuFloatStorage, layout: &Layout, to: FloatDType) -> Result<CpuFloatStorage> {
109        let s = match to {
110            FloatDType::F32 => CpuFloatStorage::F32(dispatch_float_raw!(x, |d| layout.storage_indices().map(|i| d[i] as f32).collect())),
111            FloatDType::F64 => CpuFloatStorage::F64(dispatch_float_raw!(x, |d| layout.storage_indices().map(|i| d[i] as f64).collect())),
112        };
113        Ok(s)
114    }
115
116    fn f_cast_int(x: &CpuFloatStorage, layout: &Layout, to: IntDType) -> Result<CpuIntStorage> {
117        let s = match to {
118            IntDType::I32 => CpuIntStorage::I32(dispatch_float_raw!(x, |d| layout.storage_indices().map(|i| d[i] as i32).collect())),
119            IntDType::U32 => CpuIntStorage::U32(dispatch_float_raw!(x, |d| layout.storage_indices().map(|i| d[i] as u32).collect())),
120            IntDType::U8 => CpuIntStorage::U8(dispatch_float_raw!(x, |d| layout.storage_indices().map(|i| d[i] as u8).collect())),
121        };
122        Ok(s)
123    }
124
125    fn f_cast_bool(x: &CpuFloatStorage, layout: &Layout, _to: BoolDType) -> Result<CpuBoolStorage> {
126        Ok(CpuBoolStorage(dispatch_float_raw!(x, |d| layout.storage_indices().map(|i| d[i] != 0.).collect())))
127    }
128
129    fn f_to_vec(x: &<Cpu as Device>::FloatStorage, layout: &Layout) -> Result<Vec<f64>> {
130        Ok(match x {
131            CpuFloatStorage::F32(d) => layout.storage_indices().map(|i| d[i] as f64).collect(),
132            CpuFloatStorage::F64(d) => layout.storage_indices().map(|i| d[i]).collect(),
133        })
134    }
135
136    fn f_to_bytes<'a>(x: &'a <Cpu as Device>::FloatStorage, layout: &Layout) -> Result<Cow<'a, [u8]>> {
137        if layout.is_contiguous() {
138            Ok(match x {
139                CpuFloatStorage::F32(d) => Cow::Borrowed(bytemuck::cast_slice(d)),
140                CpuFloatStorage::F64(d) => Cow::Borrowed(bytemuck::cast_slice(d)),
141            })
142        } else {
143            let contig = Self::f_contiguous(x, layout)?;
144            Ok(match contig {
145                CpuFloatStorage::F32(d) => Cow::Owned(bytemuck::cast_slice(&d).to_vec()),
146                CpuFloatStorage::F64(d) => Cow::Owned(bytemuck::cast_slice(&d).to_vec()),
147            })
148        }
149    }
150
151    fn f_binary(
152        lhs: &<Cpu as Device>::FloatStorage,
153        lhs_l: &Layout,
154        rhs: &<Cpu as Device>::FloatStorage,
155        rhs_l: &Layout,
156        op: BinaryOp,
157    ) -> Result<<Cpu as Device>::FloatStorage> {
158        dispatch_float2!(lhs, rhs, "binary", |a, b| ew::num_binary(a, lhs_l, b, rhs_l, op))
159    }
160
161    fn f_binary_scalar(
162        lhs: &<Cpu as Device>::FloatStorage,
163        lhs_l: &Layout,
164        rhs: f64,
165        op: BinaryOp,
166    ) -> Result<<Cpu as Device>::FloatStorage> {
167        Ok(match lhs {
168            CpuFloatStorage::F32(d) => CpuFloatStorage::F32(ew::num_binary_scalar(d, lhs_l, rhs as f32, op)),
169            CpuFloatStorage::F64(d) => CpuFloatStorage::F64(ew::num_binary_scalar(d, lhs_l, rhs, op)),
170        })
171    }
172
173    fn f_binary_scalar_(dst: &mut <Cpu as Device>::FloatStorage, dst_l: &Layout, rhs: f64, op: BinaryOp) -> Result<()> {
174        match dst {
175            CpuFloatStorage::F32(d) => {
176                ew::binary_scalar_(d, dst_l, rhs as f32, binary_fn_f32(op));
177                Ok(())
178            }
179            CpuFloatStorage::F64(d) => {
180                ew::binary_scalar_(d, dst_l, rhs, binary_fn_f64(op));
181                Ok(())
182            }
183        }
184    }
185
186    fn f_binary_scalar_lhs(scalar: f64, rhs: &CpuFloatStorage, rhs_l: &Layout, op: BinaryOp) -> Result<CpuFloatStorage> {
187        Ok(match rhs {
188            CpuFloatStorage::F32(d) => CpuFloatStorage::F32(ew::num_scalar_binary(scalar as f32, d, rhs_l, op)),
189            CpuFloatStorage::F64(d) => CpuFloatStorage::F64(ew::num_scalar_binary(scalar, d, rhs_l, op)),
190        })
191    }
192
193    fn f_unary(x: &<Cpu as Device>::FloatStorage, l: &Layout, op: crate::UnaryOp<f64>) -> Result<<Cpu as Device>::FloatStorage> {
194        match x {
195            CpuFloatStorage::F32(d) => Ok(CpuFloatStorage::F32(match op {
196                crate::UnaryOp::Neg => ew::unary(d, l, |v: f32| -v),
197                crate::UnaryOp::Abs => ew::unary(d, l, |v: f32| v.abs()),
198                crate::UnaryOp::Sign => ew::unary(d, l, |v: f32| v.signum()),
199                crate::UnaryOp::Affine(mul, add) => ew::unary(d, l, |v: f32| v * mul as f32 + add as f32),
200                crate::UnaryOp::Pow(exp) => ew::unary(d, l, |v: f32| v.powf(exp as f32)),
201                crate::UnaryOp::Clamp(min, max) => {
202                    let lo = min.map(|v| v as f32);
203                    let hi = max.map(|v| v as f32);
204                    ew::unary(d, l, |v: f32| {
205                        let mut val = v;
206                        if let Some(lo) = lo {
207                            val = lo.max(val);
208                        }
209                        if let Some(hi) = hi {
210                            val = hi.min(val);
211                        }
212                        val
213                    })
214                }
215            })),
216            CpuFloatStorage::F64(d) => Ok(CpuFloatStorage::F64(match op {
217                crate::UnaryOp::Neg => ew::unary(d, l, |v: f64| -v),
218                crate::UnaryOp::Abs => ew::unary(d, l, |v: f64| v.abs()),
219                crate::UnaryOp::Sign => ew::unary(d, l, |v: f64| v.signum()),
220                crate::UnaryOp::Affine(mul, add) => ew::unary(d, l, |v: f64| v * mul + add),
221                crate::UnaryOp::Pow(exp) => ew::unary(d, l, |v: f64| v.powf(exp)),
222                crate::UnaryOp::Clamp(min, max) => ew::unary(d, l, |v: f64| {
223                    let mut val = v;
224                    if let Some(lo) = min {
225                        val = lo.max(val);
226                    }
227                    if let Some(hi) = max {
228                        val = hi.min(val);
229                    }
230                    val
231                }),
232            })),
233        }
234    }
235
236    fn f_float_unary(x: &<Cpu as Device>::FloatStorage, l: &Layout, op: crate::FloatUnaryOp) -> Result<<Cpu as Device>::FloatStorage> {
237        Ok(dispatch_float!(x, |d| ew::float_unary(d, l, op)))
238    }
239
240    fn f_unary_(dst: &mut <Cpu as Device>::FloatStorage, dst_l: &Layout, op: crate::UnaryOp<f64>) -> Result<()> {
241        match dst {
242            CpuFloatStorage::F32(d) => Ok(match op {
243                crate::UnaryOp::Neg => ew::unary_(d, dst_l, |v: f32| -v),
244                crate::UnaryOp::Abs => ew::unary_(d, dst_l, |v: f32| v.abs()),
245                crate::UnaryOp::Sign => ew::unary_(d, dst_l, |v: f32| v.signum()),
246                crate::UnaryOp::Affine(mul, add) => ew::unary_(d, dst_l, |v: f32| v * mul as f32 + add as f32),
247                crate::UnaryOp::Pow(exp) => ew::unary_(d, dst_l, |v: f32| v.powf(exp as f32)),
248                crate::UnaryOp::Clamp(min, max) => {
249                    let lo = min.map(|v| v as f32);
250                    let hi = max.map(|v| v as f32);
251                    ew::unary_(d, dst_l, |v: f32| {
252                        let mut val = v;
253                        if let Some(lo) = lo {
254                            val = lo.max(val);
255                        }
256                        if let Some(hi) = hi {
257                            val = hi.min(val);
258                        }
259                        val
260                    })
261                }
262            }),
263            CpuFloatStorage::F64(d) => Ok(match op {
264                crate::UnaryOp::Neg => ew::unary_(d, dst_l, |v: f64| -v),
265                crate::UnaryOp::Abs => ew::unary_(d, dst_l, |v: f64| v.abs()),
266                crate::UnaryOp::Sign => ew::unary_(d, dst_l, |v: f64| v.signum()),
267                crate::UnaryOp::Affine(mul, add) => ew::unary_(d, dst_l, |v: f64| v * mul + add),
268                crate::UnaryOp::Pow(exp) => ew::unary_(d, dst_l, |v: f64| v.powf(exp)),
269                crate::UnaryOp::Clamp(min, max) => ew::unary_(d, dst_l, |v: f64| {
270                    let mut val = v;
271                    if let Some(lo) = min {
272                        val = lo.max(val);
273                    }
274                    if let Some(hi) = max {
275                        val = hi.min(val);
276                    }
277                    val
278                }),
279            }),
280        }
281    }
282
283    fn f_float_unary_(dst: &mut <Cpu as Device>::FloatStorage, dst_l: &Layout, op: crate::FloatUnaryOp) -> Result<()> {
284        use super::kernels::element::CpuFloat;
285        match dst {
286            CpuFloatStorage::F32(d) => {
287                match op {
288                    crate::FloatUnaryOp::Exp => ew::unary_(d, dst_l, |v: f32| v.exp()),
289                    crate::FloatUnaryOp::Ln => ew::unary_(d, dst_l, |v: f32| v.ln()),
290                    crate::FloatUnaryOp::Sin => ew::unary_(d, dst_l, |v: f32| v.sin()),
291                    crate::FloatUnaryOp::Cos => ew::unary_(d, dst_l, |v: f32| v.cos()),
292                    crate::FloatUnaryOp::Tanh => ew::unary_(d, dst_l, |v: f32| v.tanh()),
293                    crate::FloatUnaryOp::Sqr => ew::unary_(d, dst_l, |v: f32| v.sqr()),
294                    crate::FloatUnaryOp::Sqrt => ew::unary_(d, dst_l, |v: f32| v.sqrt()),
295                    crate::FloatUnaryOp::Recip => ew::unary_(d, dst_l, |v: f32| v.recip()),
296                    crate::FloatUnaryOp::Gelu => ew::unary_(d, dst_l, |v: f32| v.gelu()),
297                    crate::FloatUnaryOp::GeluErf => ew::unary_(d, dst_l, |v: f32| v.gelu_erf()),
298                    crate::FloatUnaryOp::Erf => ew::unary_(d, dst_l, |v: f32| CpuFloat::erf(v)),
299                    crate::FloatUnaryOp::Relu => ew::unary_(d, dst_l, |v: f32| v.relu()),
300                    crate::FloatUnaryOp::Silu => ew::unary_(d, dst_l, |v: f32| v.silu()),
301                    crate::FloatUnaryOp::Sigmoid => ew::unary_(d, dst_l, |v: f32| v.sigmoid()),
302                    crate::FloatUnaryOp::Floor => ew::unary_(d, dst_l, |v: f32| v.floor()),
303                    crate::FloatUnaryOp::Ceil => ew::unary_(d, dst_l, |v: f32| v.ceil()),
304                    crate::FloatUnaryOp::Round => ew::unary_(d, dst_l, |v: f32| v.round()),
305                    crate::FloatUnaryOp::LeakyRelu(a) => ew::unary_(d, dst_l, |v: f32| v.leaky_relu(a as f32)),
306                }
307                Ok(())
308            }
309            CpuFloatStorage::F64(d) => {
310                match op {
311                    crate::FloatUnaryOp::Exp => ew::unary_(d, dst_l, |v: f64| v.exp()),
312                    crate::FloatUnaryOp::Ln => ew::unary_(d, dst_l, |v: f64| v.ln()),
313                    crate::FloatUnaryOp::Sin => ew::unary_(d, dst_l, |v: f64| v.sin()),
314                    crate::FloatUnaryOp::Cos => ew::unary_(d, dst_l, |v: f64| v.cos()),
315                    crate::FloatUnaryOp::Tanh => ew::unary_(d, dst_l, |v: f64| v.tanh()),
316                    crate::FloatUnaryOp::Sqr => ew::unary_(d, dst_l, |v: f64| v.sqr()),
317                    crate::FloatUnaryOp::Sqrt => ew::unary_(d, dst_l, |v: f64| v.sqrt()),
318                    crate::FloatUnaryOp::Recip => ew::unary_(d, dst_l, |v: f64| v.recip()),
319                    crate::FloatUnaryOp::Gelu => ew::unary_(d, dst_l, |v: f64| v.gelu()),
320                    crate::FloatUnaryOp::GeluErf => ew::unary_(d, dst_l, |v: f64| v.gelu_erf()),
321                    crate::FloatUnaryOp::Erf => ew::unary_(d, dst_l, |v: f64| CpuFloat::erf(v)),
322                    crate::FloatUnaryOp::Relu => ew::unary_(d, dst_l, |v: f64| v.relu()),
323                    crate::FloatUnaryOp::Silu => ew::unary_(d, dst_l, |v: f64| v.silu()),
324                    crate::FloatUnaryOp::Sigmoid => ew::unary_(d, dst_l, |v: f64| v.sigmoid()),
325                    crate::FloatUnaryOp::Floor => ew::unary_(d, dst_l, |v: f64| v.floor()),
326                    crate::FloatUnaryOp::Ceil => ew::unary_(d, dst_l, |v: f64| v.ceil()),
327                    crate::FloatUnaryOp::Round => ew::unary_(d, dst_l, |v: f64| v.round()),
328                    crate::FloatUnaryOp::LeakyRelu(a) => ew::unary_(d, dst_l, |v: f64| v.leaky_relu(a)),
329                }
330                Ok(())
331            }
332        }
333    }
334
335    fn f_cmp(
336        lhs: &<Cpu as Device>::FloatStorage,
337        lhs_l: &Layout,
338        rhs: &<Cpu as Device>::FloatStorage,
339        rhs_l: &Layout,
340        op: CmpOp,
341    ) -> Result<<Cpu as Device>::BoolStorage> {
342        let v = dispatch_float2_raw!(lhs, rhs, "cmp", |a, b| ew::num_cmp(a, lhs_l, b, rhs_l, op))?;
343        Ok(CpuBoolStorage(v))
344    }
345
346    fn f_cmp_scalar(lhs: &<Cpu as Device>::FloatStorage, lhs_l: &Layout, rhs: f64, op: CmpOp) -> Result<<Cpu as Device>::BoolStorage> {
347        match lhs {
348            CpuFloatStorage::F32(d) => Ok(CpuBoolStorage(ew::cmp_scalar(d, lhs_l, rhs as f32, op))),
349            CpuFloatStorage::F64(d) => Ok(CpuBoolStorage(ew::cmp_scalar(d, lhs_l, rhs, op))),
350        }
351    }
352
353    fn f_reduce(
354        x: &<Cpu as Device>::FloatStorage,
355        l: &Layout,
356        dims: &[usize],
357        keepdim: bool,
358        op: crate::ReduceOp,
359    ) -> Result<(<Cpu as Device>::FloatStorage, Shape)> {
360        let reducer = reduce::Reducer::from(op);
361        match x {
362            CpuFloatStorage::F32(d) => {
363                let (v, s) = reduce::reduce_dims(d, l, dims, keepdim, reducer)?;
364                Ok((CpuFloatStorage::F32(v), s))
365            }
366            CpuFloatStorage::F64(d) => {
367                let (v, s) = reduce::reduce_dims(d, l, dims, keepdim, reducer)?;
368                Ok((CpuFloatStorage::F64(v), s))
369            }
370        }
371    }
372
373    fn f_arg_reduce(
374        x: &<Cpu as Device>::FloatStorage,
375        l: &Layout,
376        dim: usize,
377        keepdim: bool,
378        take_max: bool,
379    ) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
380        let (idx, shape) = dispatch_float_raw!(x, |d| reduce::arg_reduce(d, l, dim, keepdim, take_max))?;
381        Ok((usize_to_int_storage(&idx, DType::U32), shape))
382    }
383
384    fn f_matmul(
385        lhs: &<Cpu as Device>::FloatStorage,
386        lhs_l: &Layout,
387        rhs: &<Cpu as Device>::FloatStorage,
388        rhs_l: &Layout,
389    ) -> Result<(<Cpu as Device>::FloatStorage, Shape)> {
390        match (lhs, rhs) {
391            (CpuFloatStorage::F32(a), CpuFloatStorage::F32(b)) => {
392                let (v, s) = matmul::matmul(a, lhs_l, b, rhs_l)?;
393                Ok((CpuFloatStorage::F32(v), s))
394            }
395            (CpuFloatStorage::F64(a), CpuFloatStorage::F64(b)) => {
396                let (v, s) = matmul::matmul(a, lhs_l, b, rhs_l)?;
397                Ok((CpuFloatStorage::F64(v), s))
398            }
399            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "matmul" }),
400        }
401    }
402
403    fn f_add_matmul_(
404        dst: &mut <Cpu as Device>::FloatStorage,
405        dst_l: &Layout,
406        lhs: &<Cpu as Device>::FloatStorage,
407        lhs_l: &Layout,
408        rhs: &<Cpu as Device>::FloatStorage,
409        rhs_l: &Layout,
410    ) -> Result<()> {
411        // dst += lhs @ rhs  (fused, no temporary product buffer)
412        match (dst, lhs, rhs) {
413            (CpuFloatStorage::F32(d), CpuFloatStorage::F32(l), CpuFloatStorage::F32(r)) => matmul::add_matmul(d, dst_l, l, lhs_l, r, rhs_l),
414            (CpuFloatStorage::F64(d), CpuFloatStorage::F64(l), CpuFloatStorage::F64(r)) => matmul::add_matmul(d, dst_l, l, lhs_l, r, rhs_l),
415            (_d, l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "f_add_matmul_" }),
416        }
417    }
418
419    fn f_binary_(
420        dst: &mut <Cpu as Device>::FloatStorage,
421        dst_l: &Layout,
422        src: &<Cpu as Device>::FloatStorage,
423        src_l: &Layout,
424        op: BinaryOp,
425    ) -> Result<()> {
426        match (dst, src) {
427            (CpuFloatStorage::F32(d), CpuFloatStorage::F32(s)) => {
428                ew::binary_(d, dst_l, s, src_l, binary_fn_f32(op));
429                Ok(())
430            }
431            (CpuFloatStorage::F64(d), CpuFloatStorage::F64(s)) => {
432                ew::binary_(d, dst_l, s, src_l, binary_fn_f64(op));
433                Ok(())
434            }
435            (d, s) => Err(Error::DTypeMismatch { lhs: d.dtype(), rhs: s.dtype(), op: "in-place binary" }),
436        }
437    }
438
439    fn f_index_select(
440        x: &<Cpu as Device>::FloatStorage,
441        x_l: &Layout,
442        idx: &<Cpu as Device>::IntStorage,
443        idx_l: &Layout,
444        dim: usize,
445    ) -> Result<(<Cpu as Device>::FloatStorage, Shape)> {
446        let ids = int_ids_as_usize(idx, idx_l);
447        match x {
448            CpuFloatStorage::F32(d) => {
449                let (v, dims) = indexing::index_select(d, x_l, &ids, idx_l, dim)?;
450                Ok((CpuFloatStorage::F32(v), Shape::from(dims)))
451            }
452            CpuFloatStorage::F64(d) => {
453                let (v, dims) = indexing::index_select(d, x_l, &ids, idx_l, dim)?;
454                Ok((CpuFloatStorage::F64(v), Shape::from(dims)))
455            }
456        }
457    }
458
459    fn f_gather(
460        x: &<Cpu as Device>::FloatStorage,
461        x_l: &Layout,
462        idx: &<Cpu as Device>::IntStorage,
463        idx_l: &Layout,
464        dim: usize,
465    ) -> Result<(<Cpu as Device>::FloatStorage, Shape)> {
466        let ids = int_ids_as_usize(idx, idx_l);
467        match x {
468            CpuFloatStorage::F32(d) => {
469                let (v, dims) = indexing::gather(d, x_l, &ids, idx_l, dim)?;
470                Ok((CpuFloatStorage::F32(v), Shape::from(dims)))
471            }
472            CpuFloatStorage::F64(d) => {
473                let (v, dims) = indexing::gather(d, x_l, &ids, idx_l, dim)?;
474                Ok((CpuFloatStorage::F64(v), Shape::from(dims)))
475            }
476        }
477    }
478
479    fn f_index_add(
480        init: &<Cpu as Device>::FloatStorage,
481        init_l: &Layout,
482        idx: &<Cpu as Device>::IntStorage,
483        idx_l: &Layout,
484        src: &<Cpu as Device>::FloatStorage,
485        _src_l: &Layout,
486        dim: usize,
487    ) -> Result<<Cpu as Device>::FloatStorage> {
488        let ids = int_ids_as_usize(idx, idx_l);
489        dispatch_float2!(init, src, "index-add", |a, b| indexing::index_add(a, init_l, &ids, idx_l, b, dim)?)
490    }
491
492    fn f_scatter_add(
493        init: &<Cpu as Device>::FloatStorage,
494        init_l: &Layout,
495        idx: &<Cpu as Device>::IntStorage,
496        idx_l: &Layout,
497        src: &<Cpu as Device>::FloatStorage,
498        _src_l: &Layout,
499        dim: usize,
500    ) -> Result<<Cpu as Device>::FloatStorage> {
501        let ids = int_ids_as_usize(idx, idx_l);
502        dispatch_float2!(init, src, "scatter-add", |a, b| indexing::scatter_add(a, init_l, &ids, idx_l, b, dim)?)
503    }
504
505    fn f_cat(srcs: &[(&<Cpu as Device>::FloatStorage, &Layout)], dim: usize) -> Result<(<Cpu as Device>::FloatStorage, Shape)> {
506        if srcs.is_empty() {
507            return Err(Error::OpRequiresAtLeastOneTensor { op: "cat" });
508        }
509        // all must share dtype (checked against the first)
510        let dt = srcs[0].0.dtype();
511        for (s, _) in srcs {
512            if s.dtype() != dt {
513                return Err(Error::DTypeMismatch { lhs: dt, rhs: s.dtype(), op: "cat" });
514            }
515        }
516        match dt {
517            DType::F32 => {
518                let views: Vec<(&[f32], &Layout)> = srcs.iter().map(|(s, l)| (as_f32(s), *l)).collect();
519                let (v, shape) = super::kernels::shape::cat(&views, dim)?;
520                Ok((CpuFloatStorage::F32(v), shape))
521            }
522            _ => {
523                let views: Vec<(&[f64], &Layout)> = srcs.iter().map(|(s, l)| (as_f64(s), *l)).collect();
524                let (v, shape) = super::kernels::shape::cat(&views, dim)?;
525                Ok((CpuFloatStorage::F64(v), shape))
526            }
527        }
528    }
529
530    fn f_softmax(x: &<Cpu as Device>::FloatStorage, l: &Layout, dim: usize) -> Result<<Cpu as Device>::FloatStorage> {
531        match x {
532            CpuFloatStorage::F32(d) => Ok(CpuFloatStorage::F32(nn::softmax(d, l, dim)?)),
533            CpuFloatStorage::F64(d) => Ok(CpuFloatStorage::F64(nn::softmax(d, l, dim)?)),
534        }
535    }
536
537    fn f_rms_norm(
538        x: &<Cpu as Device>::FloatStorage,
539        x_l: &Layout,
540        weight: &<Cpu as Device>::FloatStorage,
541        weight_l: &Layout,
542        eps: f64,
543    ) -> Result<<Cpu as Device>::FloatStorage> {
544        match (x, weight) {
545            (CpuFloatStorage::F32(d), CpuFloatStorage::F32(w)) => Ok(CpuFloatStorage::F32(nn::rms_norm(d, x_l, w, weight_l, eps as f32)?)),
546            (CpuFloatStorage::F64(d), CpuFloatStorage::F64(w)) => Ok(CpuFloatStorage::F64(nn::rms_norm(d, x_l, w, weight_l, eps)?)),
547            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "rms_norm" }),
548        }
549    }
550
551    fn f_pick(
552        mask: &<Cpu as Device>::BoolStorage,
553        mask_l: &Layout,
554        on_true: &<Cpu as Device>::FloatStorage,
555        true_l: &Layout,
556        on_false: &<Cpu as Device>::FloatStorage,
557        false_l: &Layout,
558    ) -> Result<<Cpu as Device>::FloatStorage> {
559        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
560        match (on_true, on_false) {
561            (CpuFloatStorage::F32(t), CpuFloatStorage::F32(f)) => {
562                let tv = super::kernels::iter::gather(t, true_l);
563                let fv = super::kernels::iter::gather(f, false_l);
564                Ok(CpuFloatStorage::F32(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { fv[i] }).collect()))
565            }
566            (CpuFloatStorage::F64(t), CpuFloatStorage::F64(f)) => {
567                let tv = super::kernels::iter::gather(t, true_l);
568                let fv = super::kernels::iter::gather(f, false_l);
569                Ok(CpuFloatStorage::F64(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { fv[i] }).collect()))
570            }
571            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "pick" }),
572        }
573    }
574
575    fn f_pick_true(
576        mask: &<Cpu as Device>::BoolStorage,
577        mask_l: &Layout,
578        value: f64,
579        on_false: &<Cpu as Device>::FloatStorage,
580        false_l: &Layout,
581    ) -> Result<<Cpu as Device>::FloatStorage> {
582        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
583        match on_false {
584            CpuFloatStorage::F32(f) => {
585                let fv = super::kernels::iter::gather(f, false_l);
586                let val = value as f32;
587                Ok(CpuFloatStorage::F32(m.iter().enumerate().map(|(i, &c)| if c { val } else { fv[i] }).collect()))
588            }
589            CpuFloatStorage::F64(f) => {
590                let fv = super::kernels::iter::gather(f, false_l);
591                Ok(CpuFloatStorage::F64(m.iter().enumerate().map(|(i, &c)| if c { value } else { fv[i] }).collect()))
592            }
593        }
594    }
595
596    fn f_pick_false(
597        mask: &<Cpu as Device>::BoolStorage,
598        mask_l: &Layout,
599        on_true: &<Cpu as Device>::FloatStorage,
600        true_l: &Layout,
601        value: f64,
602    ) -> Result<<Cpu as Device>::FloatStorage> {
603        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
604        match on_true {
605            CpuFloatStorage::F32(t) => {
606                let tv = super::kernels::iter::gather(t, true_l);
607                let val = value as f32;
608                Ok(CpuFloatStorage::F32(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { val }).collect()))
609            }
610            CpuFloatStorage::F64(t) => {
611                let tv = super::kernels::iter::gather(t, true_l);
612                Ok(CpuFloatStorage::F64(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { value }).collect()))
613            }
614        }
615    }
616
617    fn f_allclose(a: &CpuFloatStorage, a_l: &Layout, b: &CpuFloatStorage, b_l: &Layout, rtol: f64, atol: f64) -> Result<bool> {
618        match (a, b) {
619            (CpuFloatStorage::F32(av), CpuFloatStorage::F32(bv)) => {
620                let rtol = rtol as f32;
621                let atol = atol as f32;
622                Ok(a_l.storage_indices().zip(b_l.storage_indices()).all(|(ai, bi)| {
623                    let diff = (av[ai] - bv[bi]).abs();
624                    diff <= atol + rtol * bv[bi].abs()
625                }))
626            }
627            (CpuFloatStorage::F64(av), CpuFloatStorage::F64(bv)) => Ok(a_l.storage_indices().zip(b_l.storage_indices()).all(|(ai, bi)| {
628                let diff = (av[ai] - bv[bi]).abs();
629                diff <= atol + rtol * bv[bi].abs()
630            })),
631            _ => return Err(crate::Error::DTypeMismatch { lhs: a.dtype(), rhs: b.dtype(), op: "allclose" }),
632        }
633    }
634}
635
636fn as_f32(s: &CpuFloatStorage) -> &[f32] {
637    match s {
638        CpuFloatStorage::F32(d) => d,
639        _ => unreachable!("dtype checked by caller"),
640    }
641}
642
643fn as_f64(s: &CpuFloatStorage) -> &[f64] {
644    match s {
645        CpuFloatStorage::F64(d) => d,
646        _ => unreachable!("dtype checked by caller"),
647    }
648}
649
650fn binary_fn_f32(op: BinaryOp) -> fn(f32, f32) -> f32 {
651    use super::kernels::element::CpuNum;
652    match op {
653        BinaryOp::Add => |a, b| a + b,
654        BinaryOp::Sub => |a, b| a - b,
655        BinaryOp::Mul => |a, b| a * b,
656        BinaryOp::Div => |a, b| a / b,
657        BinaryOp::Maximum => CpuNum::maximum,
658        BinaryOp::Minimum => CpuNum::minimum,
659    }
660}
661
662fn binary_fn_f64(op: BinaryOp) -> fn(f64, f64) -> f64 {
663    use super::kernels::element::CpuNum;
664    match op {
665        BinaryOp::Add => |a, b| a + b,
666        BinaryOp::Sub => |a, b| a - b,
667        BinaryOp::Mul => |a, b| a * b,
668        BinaryOp::Div => |a, b| a / b,
669        BinaryOp::Maximum => CpuNum::maximum,
670        BinaryOp::Minimum => CpuNum::minimum,
671    }
672}
luma_tensor/device/cpu/ops/float_ops.rs

luma_tensor/device/cpu/ops/
float_ops.rs
