legume_numeric/matrix/
tensor_util.rs

1use crate::matrix::rand_util::{collect_f32_seeded, entropy_seed};
2use crate::matrix::traits::*;
3use candle_core::{CpuStorage, DType, Device, InplaceOp2, Layout, Tensor};
4use rand_distr::{Gamma, StandardNormal, Uniform};
5use rayon::prelude::*;
6
7impl SampleOps for Tensor {
8    type Mat = Self;
9    type Scalar = f32;
10
11    fn runif(nrow: usize, ncol: usize) -> Self::Mat {
12        Self::runif_seeded(nrow, ncol, entropy_seed())
13    }
14
15    fn rnorm(nrow: usize, ncol: usize) -> Self::Mat {
16        Self::rnorm_seeded(nrow, ncol, entropy_seed())
17    }
18
19    fn rgamma(nrow: usize, ncol: usize, param: (f32, f32)) -> Self::Mat {
20        Self::rgamma_seeded(nrow, ncol, param, entropy_seed())
21    }
22
23    fn runif_seeded(nrow: usize, ncol: usize, seed: u64) -> Self::Mat {
24        let u01 = Uniform::new(0_f32, 1_f32).expect("failed to create uniform distribution");
25        let data = collect_f32_seeded(nrow * ncol, u01, seed);
26        Tensor::from_vec(data, (nrow, ncol), &Device::Cpu)
27            .expect("failed to create Tensor runif_seeded")
28    }
29
30    fn rnorm_seeded(nrow: usize, ncol: usize, seed: u64) -> Self::Mat {
31        // Candle's own `Tensor::randn` on the CPU backend draws from
32        // `rand::rng()` (OS entropy) and its `Device::set_seed` errors out, so
33        // there is no way to seed it through candle. Generate the sample
34        // host-side instead, where the seed is honored.
35        let data = collect_f32_seeded(nrow * ncol, StandardNormal, seed);
36        Tensor::from_vec(data, (nrow, ncol), &Device::Cpu)
37            .expect("failed to create Tensor rnorm_seeded")
38    }
39
40    fn rgamma_seeded(nrow: usize, ncol: usize, param: (f32, f32), seed: u64) -> Self::Mat {
41        let (shape, scale) = param;
42        let pdf = Gamma::new(shape, scale).unwrap();
43        let data = collect_f32_seeded(nrow * ncol, pdf, seed);
44        Tensor::from_vec(data, (nrow, ncol), &Device::Cpu)
45            .expect("failed to create Tensor rgamma_seeded")
46    }
47}
48
49impl MatTriplets for Tensor {
50    type Mat = Self;
51    type Scalar = f32;
52
53    fn from_nonzero_triplets<I>(
54        nrow: usize,
55        ncol: usize,
56        triplets: &[(I, I, Self::Scalar)],
57    ) -> anyhow::Result<Self::Mat>
58    where
59        I: TryInto<usize> + Copy,
60        <I as TryInto<usize>>::Error: std::fmt::Debug,
61    {
62        let mut data = vec![0_f32; ncol * nrow];
63        for &(ii, jj, x_ij) in triplets {
64            let ii: usize = ii.try_into().expect("failed to convert index ii");
65            let jj: usize = jj.try_into().expect("failed to convert index jj");
66            data[ii * ncol + jj] = x_ij;
67        }
68        Ok(Tensor::from_vec(data, (nrow, ncol), &Device::Cpu)?)
69    }
70
71    fn to_nonzero_triplets(&self) -> anyhow::Result<NRowNColTriplets<Self::Scalar>> {
72        if let Ok((nrow, ncol)) = self.dims2() {
73            let eps = 1e-6;
74            let mut ret = vec![];
75            let xx: Vec<Vec<Self::Scalar>> = self.to_vec2()?;
76            for (i, x_i) in xx.iter().enumerate() {
77                for (j, &x_ij) in x_i.iter().enumerate() {
78                    if x_ij.abs() > eps {
79                        ret.push((i, j, x_ij));
80                    }
81                }
82            }
83
84            Ok(NRowNColTriplets {
85                nrow,
86                ncol,
87                triplets: ret,
88            })
89        } else {
90            anyhow::bail!("not a 2D Tensor");
91        }
92    }
93}
94
95// impl CandleDataLoaderOps for Tensor {
96//     fn rows_to_tensor_vec(&self) -> Vec<Tensor> {
97//         let mut idx_data = (0..self.dims()[0])
98//             .map(|i| (i, self.narrow(0, i, 1).expect("").clone()))
99//             .collect::<Vec<_>>();
100
101//         idx_data.sort_by_key(|(i, _)| *i);
102//         idx_data.into_iter().map(|(_, t)| t).collect()
103//     }
104// }
105
106///////////////////////////////////
107// Fused elementwise CPU kernels //
108///////////////////////////////////
109
110impl FusedTensorOps for Tensor {
111    fn clamped_exp_add_inplace(self, offset: &Tensor, ceiling: f64) -> anyhow::Result<Self> {
112        // One eligibility test, not a shape test: contiguity and dtype are this
113        // kernel's requirements, not the math's, so anything it cannot walk flatly
114        // takes the op chain rather than silently reading the wrong elements.
115        let broadcast = match (self.dims(), offset.dims()) {
116            (&[n, f], &[n_off, f_off]) => Broadcast::of(n, f, n_off, f_off),
117            _ => None,
118        };
119        let fused = matches!(self.device(), Device::Cpu)
120            && self.dtype() == DType::F32
121            && offset.dtype() == DType::F32
122            && self.is_contiguous()
123            && offset.is_contiguous();
124
125        let (Some(broadcast), true) = (broadcast, fused) else {
126            return Ok(self.broadcast_add(offset)?.minimum(ceiling)?.exp()?);
127        };
128
129        self.inplace_op2(
130            offset,
131            &ClampedExpAdd {
132                broadcast,
133                // f64 → f32 once here rather than per element, by the same cast
134                // `Tensor::minimum` applies to a scalar bound — so the fused path and
135                // the op chain agree bitwise.
136                ceiling: ceiling as f32,
137            },
138        )?;
139        Ok(self)
140    }
141}
142
143/// How the `[N, F]` receiver reads its offset. The three shapes
144/// `Tensor::broadcast_add` accepts here, resolved once outside the row loop.
145#[derive(Clone, Copy)]
146enum Broadcast {
147    /// `[N, F]` — each row reads its own.
148    Element,
149    /// `[1, F]` — every row reads the same one.
150    Column,
151    /// `[N, 1]` — each row reads a single scalar.
152    Row,
153}
154
155impl Broadcast {
156    fn of(n: usize, f: usize, n_off: usize, f_off: usize) -> Option<Self> {
157        match (n_off, f_off) {
158            _ if n_off == n && f_off == f => Some(Self::Element),
159            (1, _) if f_off == f => Some(Self::Column),
160            (_, 1) if n_off == n => Some(Self::Row),
161            _ => None,
162        }
163    }
164}
165
166/// The fused kernel behind [`FusedTensorOps::clamped_exp_add_inplace`]. CPU only by
167/// construction: the device forwards `InplaceOp2` supplies default to an error, and the
168/// caller never reaches them.
169struct ClampedExpAdd {
170    broadcast: Broadcast,
171    ceiling: f32,
172}
173
174/// Elements a rayon task should carry at minimum. Rows are the natural unit — each is
175/// a contiguous run — but a thin panel (`probe` against a few hundred genes) makes a
176/// row too small to pay for a `join`, so short rows are batched up to this.
177const MIN_FUSED_TASK_ELEMS: usize = 4096;
178
179impl InplaceOp2 for ClampedExpAdd {
180    fn name(&self) -> &'static str {
181        "clamped-exp-add"
182    }
183
184    fn cpu_fwd(
185        &self,
186        s1: &mut CpuStorage,
187        l1: &Layout,
188        s2: &CpuStorage,
189        l2: &Layout,
190    ) -> candle_core::Result<()> {
191        let (CpuStorage::F32(lhs), CpuStorage::F32(rhs)) = (s1, s2) else {
192            candle_core::bail!("clamped-exp-add: expected f32 storage on both operands");
193        };
194        let (n, f) = l1.shape().dims2()?;
195
196        // Both layouts were checked contiguous, so a slice from the start offset is the
197        // whole logical tensor in row-major order.
198        let lhs = &mut lhs[l1.start_offset()..l1.start_offset() + n * f];
199        let rhs = &rhs[l2.start_offset()..l2.start_offset() + l2.shape().elem_count()];
200
201        let (ceiling, broadcast) = (self.ceiling, self.broadcast);
202        let rows_per_task = MIN_FUSED_TASK_ELEMS.div_ceil(f.max(1)).max(1);
203        lhs.par_chunks_mut(rows_per_task * f)
204            .enumerate()
205            .for_each(|(t, block)| {
206                for (r, row) in block.chunks_mut(f).enumerate() {
207                    let i = t * rows_per_task + r;
208                    match broadcast {
209                        Broadcast::Element => apply(row, &rhs[i * f..(i + 1) * f], ceiling),
210                        Broadcast::Column => apply(row, rhs, ceiling),
211                        Broadcast::Row => {
212                            let b = rhs[i];
213                            for x in row.iter_mut() {
214                                *x = (*x + b).min(ceiling).exp();
215                            }
216                        }
217                    }
218                }
219            });
220        Ok(())
221    }
222}
223
224/// `row[j] <- exp(min(row[j] + off[j], ceiling))`.
225#[inline]
226fn apply(row: &mut [f32], off: &[f32], ceiling: f32) {
227    for (x, &b) in row.iter_mut().zip(off) {
228        *x = (*x + b).min(ceiling).exp();
229    }
230}
231
232#[cfg(test)]
233#[path = "tensor_util_tests.rs"]
234mod tests;
legume_numeric/matrix/tensor_util.rs

legume_numeric/matrix/
tensor_util.rs
