tinyquant_core/codec/
prepared.rs

1//! `PreparedCodec`: pre-built session object for the hot compress/decompress path.
2//!
3//! Owns a `RotationMatrix` built once from the supplied [`CodecConfig`], so
4//! subsequent calls to [`Codec::compress_prepared`] and
5//! [`Codec::decompress_prepared_into`] skip the `O(dim²)` QR factorization
6//! that `RotationMatrix::build` performs.  This is the CPU prerequisite for
7//! attaching a GPU backend (Phase 27+): a GPU backend simply calls
8//! [`PreparedCodec::set_gpu_state`] to attach device-resident state without
9//! requiring changes to `tinyquant-core`.
10//!
11//! # Construction cost
12//!
13//! `PreparedCodec::new` runs `RotationMatrix::build` once (O(dim² × dim) QR).
14//! Amortise it by constructing once and reusing across all compress/decompress
15//! calls over the same `(config, codebook)` pair.
16
17use crate::{
18    codec::{codebook::Codebook, codec_config::CodecConfig, rotation_matrix::RotationMatrix},
19    errors::CodecError,
20};
21
22/// Pre-built compress/decompress session.
23///
24/// Construct once with [`PreparedCodec::new`], then pass to
25/// [`Codec::compress_prepared`] / [`Codec::decompress_prepared_into`]
26/// for each vector.  The [`RotationMatrix`] inside is computed exactly once.
27pub struct PreparedCodec {
28    config: CodecConfig,
29    codebook: Codebook,
30    rotation: RotationMatrix,
31    /// Reserved for GPU-resident state attached by a GPU backend (Phase 27+).
32    /// Stored as an erased type so `tinyquant-core` stays free of GPU imports.
33    gpu_state: Option<alloc::boxed::Box<dyn core::any::Any + Send + Sync>>,
34}
35
36impl PreparedCodec {
37    /// Build from a validated config and trained codebook.
38    ///
39    /// The rotation matrix is computed exactly once.
40    ///
41    /// # Errors
42    ///
43    /// [`CodecError::CodebookIncompatible`] if `codebook.bit_width() != config.bit_width()`.
44    pub fn new(config: CodecConfig, codebook: Codebook) -> Result<Self, CodecError> {
45        if codebook.bit_width() != config.bit_width() {
46            return Err(CodecError::CodebookIncompatible {
47                expected: config.bit_width(),
48                got: codebook.bit_width(),
49            });
50        }
51        let rotation = RotationMatrix::build(config.seed(), config.dimension());
52        Ok(Self {
53            config,
54            codebook,
55            rotation,
56            gpu_state: None,
57        })
58    }
59
60    /// The [`CodecConfig`] that defines this session.
61    #[inline]
62    pub const fn config(&self) -> &CodecConfig {
63        &self.config
64    }
65
66    /// The trained [`Codebook`] for this session.
67    #[inline]
68    pub const fn codebook(&self) -> &Codebook {
69        &self.codebook
70    }
71
72    /// The pre-built [`RotationMatrix`] — computed once at construction.
73    #[inline]
74    pub const fn rotation(&self) -> &RotationMatrix {
75        &self.rotation
76    }
77
78    /// Attach opaque GPU state.
79    ///
80    /// Called by GPU backends — do not call from application code.
81    pub fn set_gpu_state(&mut self, state: alloc::boxed::Box<dyn core::any::Any + Send + Sync>) {
82        self.gpu_state = Some(state);
83    }
84
85    /// `true` if a GPU backend has attached device-resident state.
86    #[inline]
87    pub fn has_gpu_state(&self) -> bool {
88        self.gpu_state.is_some()
89    }
90
91    /// Borrow the attached GPU state, if any.
92    ///
93    /// Called by GPU backends to retrieve device-resident buffers.
94    /// Application code should not call this directly.
95    #[inline]
96    pub fn gpu_state(&self) -> Option<&(dyn core::any::Any + Send + Sync)> {
97        self.gpu_state.as_deref()
98    }
99}
tinyquant_core/codec/prepared.rs

tinyquant_core/codec/
prepared.rs
