eegdino_rs/rlx/
encoder.rs

1//! RLX-backed EEG-DINO encoder.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use std::time::Instant;
7
8use crate::config::{ModelConfig, ModelSize};
9use crate::error::{EegDinoError, Result};
10
11use super::graph::{build_encoder_graph, EncoderSpec};
12use super::weights::{
13    apply_params, detect_model_size as detect_size_anyhow, load_safetensors, prepare_params,
14    ParamMap,
15};
16
17/// Per-shape compiled graph plus reusable I/O buffers.
18struct CachedEntry {
19    compiled: rlx::CompiledGraph,
20    x_buf: Vec<f32>,
21    out_buf: Vec<f32>,
22    /// Stable output slots (byte_offset, f32_len); filled on first run.
23    out_slots: Vec<(usize, usize)>,
24    /// Fallback when the backend has no `run_slots` / host arena (e.g. wgpu).
25    host_run: bool,
26}
27
28/// Copy the sole graph output from the arena into `out` after `run_slots`.
29fn read_output_into(compiled: &rlx::CompiledGraph, slots: &[(usize, usize)], out: &mut Vec<f32>) {
30    let (byte_off, len) = slots
31        .first()
32        .copied()
33        .expect("encoder graph has one output");
34    let ptr = unsafe { compiled.arena_ptr().add(byte_off) as *const f32 };
35    // SAFETY: arena valid until the next run; we copy out immediately.
36    let slice = unsafe { std::slice::from_raw_parts(ptr, len) };
37    if out.len() != len {
38        out.resize(len, 0.0);
39    }
40    out.copy_from_slice(slice);
41}
42
43fn run_host(entry: &mut CachedEntry) {
44    let v = entry
45        .compiled
46        .run(&[("x", entry.x_buf.as_slice())])
47        .into_iter()
48        .next()
49        .expect("encoder graph has one output");
50    if entry.out_buf.len() != v.len() {
51        entry.out_buf = v;
52    } else {
53        entry.out_buf.copy_from_slice(&v);
54    }
55}
56
57fn run_forward(entry: &mut CachedEntry) {
58    let x = entry.x_buf.as_slice();
59    if entry.host_run {
60        run_host(entry);
61        return;
62    }
63    if entry.out_slots.is_empty() {
64        if entry.compiled.arena_ptr().is_null() {
65            entry.host_run = true;
66            run_host(entry);
67            return;
68        }
69        entry.out_slots = entry.compiled.run_slots(&[x]).to_vec();
70        if entry.out_slots.is_empty() {
71            entry.host_run = true;
72            run_host(entry);
73            return;
74        }
75    } else {
76        entry.compiled.run_slots(&[x]);
77    }
78    read_output_into(&entry.compiled, &entry.out_slots, &mut entry.out_buf);
79}
80
81/// Result of encoding: per-sample embeddings.
82pub struct EncodingResult {
83    pub embeddings: Vec<f32>,
84    pub shape: Vec<usize>,
85    pub ms_encode: f64,
86}
87
88/// How many distinct `(batch, channels, patches)` compiled graphs to retain.
89/// GPU backends default to `1` so sweeping large batch sizes does not OOM.
90const DEFAULT_MAX_CACHED_CPU: usize = usize::MAX;
91const DEFAULT_MAX_CACHED_GPU: usize = 1;
92
93pub struct EegDinoEncoderBuilder {
94    weights_path: Option<PathBuf>,
95    config: Option<ModelConfig>,
96    normalization: f32,
97    device: Option<rlx::Device>,
98    max_cached_shapes: Option<usize>,
99}
100
101impl Default for EegDinoEncoderBuilder {
102    fn default() -> Self {
103        Self {
104            weights_path: None,
105            config: None,
106            normalization: 100.0,
107            device: None,
108            max_cached_shapes: None,
109        }
110    }
111}
112
113fn default_max_cached_shapes(device: rlx::Device) -> usize {
114    match device {
115        rlx::Device::Cpu => DEFAULT_MAX_CACHED_CPU,
116        rlx::Device::Cuda | rlx::Device::Gpu | rlx::Device::Rocm => DEFAULT_MAX_CACHED_GPU,
117        _ => 8,
118    }
119}
120
121impl EegDinoEncoderBuilder {
122    pub fn weights(mut self, path: impl Into<PathBuf>) -> Self {
123        self.weights_path = Some(path.into());
124        self
125    }
126
127    pub fn size(mut self, size: ModelSize) -> Self {
128        self.config = Some(ModelConfig::from_size(size));
129        self
130    }
131
132    pub fn config(mut self, cfg: ModelConfig) -> Self {
133        self.config = Some(cfg);
134        self
135    }
136
137    pub fn normalization(mut self, n: f32) -> Self {
138        self.normalization = n;
139        self
140    }
141
142    pub fn device(mut self, device: rlx::Device) -> Self {
143        self.device = Some(device);
144        self
145    }
146
147    /// Max compiled graphs kept at once (per distinct batch shape). Default: unlimited
148    /// on CPU, `1` on CUDA/wgpu/ROCm to avoid VRAM exhaustion when batch size changes.
149    pub fn max_cached_shapes(mut self, n: usize) -> Self {
150        self.max_cached_shapes = Some(n.max(1));
151        self
152    }
153
154    pub fn build(self) -> Result<EegDinoEncoder> {
155        let weights_path = self
156            .weights_path
157            .ok_or_else(|| EegDinoError::Builder("weights path is required".into()))?;
158        let device = self
159            .device
160            .ok_or_else(|| EegDinoError::Builder("device is required".into()))?;
161
162        let path_str = weights_path
163            .to_str()
164            .ok_or_else(|| EegDinoError::Builder("weights path is not valid UTF-8".into()))?;
165
166        let cfg = match self.config {
167            Some(c) => c,
168            None => {
169                let size = detect_size_anyhow(path_str)
170                    .map_err(|e| EegDinoError::WeightLoad(e.to_string()))?;
171                ModelConfig::from_size(size)
172            }
173        };
174
175        let raw =
176            load_safetensors(path_str).map_err(|e| EegDinoError::WeightLoad(e.to_string()))?;
177        let params = Arc::new(
178            prepare_params(&cfg, raw).map_err(|e| EegDinoError::WeightLoad(e.to_string()))?,
179        );
180
181        let max_cached_shapes = self
182            .max_cached_shapes
183            .unwrap_or_else(|| default_max_cached_shapes(device));
184
185        Ok(EegDinoEncoder {
186            cfg,
187            device,
188            normalization: self.normalization,
189            session: rlx::Session::new(device),
190            params,
191            cache: HashMap::new(),
192            batch_flat: Vec::new(),
193            max_cached_shapes,
194        })
195    }
196}
197
198/// RLX encoder. Caches one compiled graph per `(b,c,p)` shape.
199pub struct EegDinoEncoder {
200    pub cfg: ModelConfig,
201    pub device: rlx::Device,
202    pub normalization: f32,
203
204    session: rlx::Session,
205    params: Arc<ParamMap>,
206    cache: HashMap<u64, CachedEntry>,
207    /// Reused by [`Self::encode_batch`] to avoid per-call flatten allocations.
208    batch_flat: Vec<f32>,
209    max_cached_shapes: usize,
210}
211
212impl EegDinoEncoder {
213    pub fn builder() -> EegDinoEncoderBuilder {
214        EegDinoEncoderBuilder::default()
215    }
216
217    pub fn load(
218        weights_path: &Path,
219        config: Option<ModelConfig>,
220        device: rlx::Device,
221    ) -> Result<(Self, f64)> {
222        let t0 = Instant::now();
223        let mut b = Self::builder().weights(weights_path).device(device);
224        if let Some(c) = config {
225            b = b.config(c);
226        }
227        let enc = b.build()?;
228        Ok((enc, t0.elapsed().as_secs_f64() * 1000.0))
229    }
230
231    fn cache_key(spec: &EncoderSpec) -> u64 {
232        ((spec.b as u64) << 42) ^ ((spec.c as u64) << 21) ^ (spec.p as u64)
233    }
234
235    fn evict_cache_if_needed(&mut self, incoming: u64) {
236        if self.cache.contains_key(&incoming) {
237            return;
238        }
239        if self.cache.len() < self.max_cached_shapes {
240            return;
241        }
242        if self.max_cached_shapes == 1 {
243            self.cache.clear();
244            return;
245        }
246        // Drop one arbitrary entry (HashMap order); enough for small limits.
247        if let Some(k) = self.cache.keys().next().copied() {
248            self.cache.remove(&k);
249        }
250        if self.cache.len() >= self.max_cached_shapes {
251            self.cache.clear();
252        }
253    }
254
255    /// Drop all compiled graphs (frees GPU memory). Next encode recompiles.
256    pub fn clear_cache(&mut self) {
257        self.cache.clear();
258        // Release large I/O scratch if a prior shape used a bigger batch.
259        self.batch_flat.shrink_to_fit();
260    }
261
262    fn entry_for(&mut self, spec: &EncoderSpec, input_len: usize) -> Result<&mut CachedEntry> {
263        let key = Self::cache_key(spec);
264        if !self.cache.contains_key(&key) {
265            self.evict_cache_if_needed(key);
266            let graph = build_encoder_graph(&self.cfg, spec);
267            let mut compiled = self.session.compile(graph);
268            apply_params(&mut compiled, &self.cfg, spec, &self.params)
269                .map_err(|e| EegDinoError::WeightLoad(e.to_string()))?;
270            let mut x_buf = Vec::with_capacity(input_len);
271            x_buf.resize(input_len, 0.0);
272            let out_buf = Vec::new();
273            let mut entry = CachedEntry {
274                compiled,
275                x_buf,
276                out_buf,
277                out_slots: Vec::new(),
278                host_run: false,
279            };
280            run_forward(&mut entry);
281            self.cache.insert(key, entry);
282        }
283        let entry = self.cache.get_mut(&key).expect("just inserted");
284        if entry.x_buf.len() != input_len {
285            entry.x_buf.resize(input_len, 0.0);
286        }
287        Ok(entry)
288    }
289
290    /// Compile and warm up graphs for the given batch sizes (same `num_channels` / `num_samples`).
291    pub fn prewarm_batch_sizes(
292        &mut self,
293        batch_sizes: &[usize],
294        num_channels: usize,
295        num_samples: usize,
296    ) -> Result<()> {
297        for &b in batch_sizes {
298            if self.max_cached_shapes == 1 {
299                self.cache.clear();
300            }
301            let expected = b * num_channels * num_samples;
302            let spec = EncoderSpec {
303                b,
304                c: num_channels,
305                p: num_samples / self.cfg.patch_size,
306            };
307            self.entry_for(&spec, expected)?;
308        }
309        Ok(())
310    }
311
312    fn validate_encode_input(
313        &self,
314        signal: &[f32],
315        batch_size: usize,
316        num_channels: usize,
317        num_samples: usize,
318    ) -> Result<(usize, EncoderSpec)> {
319        let patch_size = self.cfg.patch_size;
320        if num_channels != self.cfg.num_channels {
321            return Err(EegDinoError::InvalidInput(format!(
322                "num_channels ({num_channels}) must equal model num_channels ({})",
323                self.cfg.num_channels
324            )));
325        }
326        if !num_samples.is_multiple_of(patch_size) {
327            return Err(EegDinoError::InvalidInput(format!(
328                "num_samples ({num_samples}) must be divisible by patch_size ({patch_size})"
329            )));
330        }
331        let expected = batch_size * num_channels * num_samples;
332        if signal.len() != expected {
333            return Err(EegDinoError::InvalidInput(format!(
334                "signal length {} != batch_size({batch_size}) * channels({num_channels}) * samples({num_samples}) = {expected}",
335                signal.len()
336            )));
337        }
338        let spec = EncoderSpec {
339            b: batch_size,
340            c: num_channels,
341            p: num_samples / patch_size,
342        };
343        Ok((expected, spec))
344    }
345
346    fn output_shape(
347        &self,
348        batch_size: usize,
349        num_channels: usize,
350        num_patches: usize,
351    ) -> Vec<usize> {
352        vec![
353            batch_size,
354            self.cfg.num_global_tokens + num_channels * num_patches,
355            self.cfg.feature_size,
356        ]
357    }
358
359    /// Encode from a flat `&[f32]` signal.
360    ///
361    /// The signal is interpreted as `[batch_size, num_channels, num_samples]`
362    /// in row-major order, divided by `normalization`, and reshaped into patches.
363    pub fn encode_raw(
364        &mut self,
365        signal: &[f32],
366        batch_size: usize,
367        num_channels: usize,
368        num_samples: usize,
369    ) -> Result<EncodingResult> {
370        let t0 = Instant::now();
371        let (expected, spec) =
372            self.validate_encode_input(signal, batch_size, num_channels, num_samples)?;
373
374        let inv_norm = 1.0f32 / self.normalization;
375        let entry = self.entry_for(&spec, expected)?;
376        for (dst, &v) in entry.x_buf.iter_mut().zip(signal) {
377            *dst = v * inv_norm;
378        }
379        run_forward(entry);
380        let embeddings = std::mem::take(&mut entry.out_buf);
381
382        Ok(EncodingResult {
383            embeddings,
384            shape: self.output_shape(batch_size, num_channels, spec.p),
385            ms_encode: t0.elapsed().as_secs_f64() * 1000.0,
386        })
387    }
388
389    /// Like [`Self::encode_raw`], but reuses `out` when its length already matches the embedding.
390    pub fn encode_raw_into(
391        &mut self,
392        signal: &[f32],
393        batch_size: usize,
394        num_channels: usize,
395        num_samples: usize,
396        out: &mut Vec<f32>,
397    ) -> Result<(Vec<usize>, f64)> {
398        let t0 = Instant::now();
399        let (expected, spec) =
400            self.validate_encode_input(signal, batch_size, num_channels, num_samples)?;
401
402        let inv_norm = 1.0f32 / self.normalization;
403        let gtok = self.cfg.num_global_tokens;
404        let feat = self.cfg.feature_size;
405        let entry = self.entry_for(&spec, expected)?;
406        for (dst, &v) in entry.x_buf.iter_mut().zip(signal) {
407            *dst = v * inv_norm;
408        }
409        run_forward(entry);
410        let shape = [batch_size, gtok + num_channels * spec.p, feat];
411        let out_len = entry.out_buf.len();
412        if out.len() == out_len {
413            out.copy_from_slice(&entry.out_buf);
414            entry.out_buf.clear();
415        } else {
416            *out = std::mem::take(&mut entry.out_buf);
417        }
418        Ok((shape.to_vec(), t0.elapsed().as_secs_f64() * 1000.0))
419    }
420
421    /// Encode multiple signals as a single batched forward pass.
422    pub fn encode_batch(
423        &mut self,
424        signals: &[Vec<f32>],
425        num_channels: usize,
426        num_samples: usize,
427    ) -> Result<EncodingResult> {
428        let expected_len = num_channels * num_samples;
429        let mut flat = std::mem::take(&mut self.batch_flat);
430        flat.clear();
431        flat.reserve(signals.len() * expected_len);
432        for (i, s) in signals.iter().enumerate() {
433            if s.len() != expected_len {
434                self.batch_flat = flat;
435                return Err(EegDinoError::InvalidInput(format!(
436                    "signal[{i}] length {} != {expected_len}",
437                    s.len()
438                )));
439            }
440            flat.extend_from_slice(s);
441        }
442        let mut embeddings = Vec::new();
443        let (shape, ms_encode) = self.encode_raw_into(
444            &flat,
445            signals.len(),
446            num_channels,
447            num_samples,
448            &mut embeddings,
449        )?;
450        self.batch_flat = flat;
451        Ok(EncodingResult {
452            embeddings,
453            shape,
454            ms_encode,
455        })
456    }
457}
eegdino_rs/rlx/encoder.rs

eegdino_rs/rlx/
encoder.rs
