torsh_ffi/
quantization.rs

1//! Model Quantization and Optimization for Edge Deployment
2//!
3//! This module provides comprehensive model quantization and optimization techniques
4//! for deploying deep learning models on edge devices, browsers (WASM), and mobile platforms.
5//!
6//! # Features
7//!
8//! - **Post-Training Quantization (PTQ)**: Convert trained models to INT8/INT4/FP16
9//! - **Quantization-Aware Training (QAT)**: Train models with quantization in mind
10//! - **Dynamic Quantization**: Runtime quantization for specific layers
11//! - **Mixed Precision**: Combine different precisions for optimal performance
12//! - **Calibration**: Use representative data for optimal quantization parameters
13//! - **Pruning**: Remove redundant parameters for model compression
14//! - **Knowledge Distillation**: Transfer knowledge from large to small models
15//!
16//! # Architecture
17//!
18//! ```text
19//! ┌─────────────────────────────────────────────────┐
20//! │          Original FP32 Model                    │
21//! └───────────────────┬─────────────────────────────┘
22//!                     │
23//!         ┌───────────▼──────────┐
24//!         │   Quantization       │
25//!         │                      │
26//!    ┌────┴────┐          ┌─────┴────┐
27//!    │   PTQ   │          │   QAT    │
28//!    └────┬────┘          └─────┬────┘
29//!         │                     │
30//!         └──────────┬──────────┘
31//!                    │
32//!     ┌──────────────▼──────────────┐
33//!     │  Quantized Model            │
34//!     │  (INT8/INT4/FP16/Mixed)     │
35//!     └──────────────┬──────────────┘
36//!                    │
37//!     ┌──────────────▼──────────────┐
38//!     │  Optimization               │
39//!     │  • Pruning                  │
40//!     │  • Distillation             │
41//!     │  • Fusion                   │
42//!     └──────────────┬──────────────┘
43//!                    │
44//!     ┌──────────────▼──────────────┐
45//!     │  Optimized Edge Model       │
46//!     │  (2-10x smaller, 2-4x faster)│
47//!     └─────────────────────────────┘
48//! ```
49//!
50//! # Quick Start
51//!
52//! ## Post-Training Quantization (PTQ)
53//!
54//! ```rust,ignore
55//! use torsh_ffi::quantization::{QuantizationConfig, Quantizer, QuantizationType};
56//!
57//! // Configure quantization
58//! let config = QuantizationConfig::new(QuantizationType::Int8)
59//!     .with_calibration_data(&calibration_set)
60//!     .with_percentile(99.99);  // Clip outliers
61//!
62//! // Quantize model
63//! let quantizer = Quantizer::new(config);
64//! let quantized_model = quantizer.quantize(&fp32_model)?;
65//!
66//! // Results: ~4x smaller, ~2-3x faster
67//! println!("Model size: {:.2}MB -> {:.2}MB",
68//!     fp32_model.size_mb(), quantized_model.size_mb());
69//! ```
70//!
71//! ## Mixed Precision Quantization
72//!
73//! ```rust,ignore
74//! use torsh_ffi::quantization::{MixedPrecisionConfig, LayerPrecision};
75//!
76//! let config = MixedPrecisionConfig::new()
77//!     .set_layer_precision("layer1", LayerPrecision::Int8)
78//!     .set_layer_precision("layer2", LayerPrecision::Fp16)
79//!     .set_layer_precision("output", LayerPrecision::Fp32);  // Keep output high precision
80//!
81//! let quantized = quantizer.quantize_mixed_precision(&model, &config)?;
82//! ```
83//!
84//! ## Quantization-Aware Training
85//!
86//! ```rust,ignore
87//! use torsh_ffi::quantization::{QatConfig, FakeQuantize};
88//!
89//! // Add fake quantization to model
90//! let qat_config = QatConfig::new(QuantizationType::Int8)
91//!     .with_observer_type(ObserverType::MovingAverage);
92//!
93//! let qat_model = qat_config.prepare_qat(&model)?;
94//!
95//! // Train with quantization simulation
96//! for epoch in 0..epochs {
97//!     train_epoch(&mut qat_model, &train_data);
98//! }
99//!
100//! // Convert to actual quantized model
101//! let quantized = qat_model.convert_to_quantized()?;
102//! ```
103//!
104//! # Supported Quantization Types
105//!
106//! - **INT8**: 8-bit integer quantization (most common, good trade-off)
107//! - **INT4**: 4-bit integer quantization (aggressive compression)
108//! - **FP16**: 16-bit floating point (good for GPUs)
109//! - **UINT8**: Unsigned 8-bit (for activations)
110//! - **Dynamic**: Per-batch/per-channel quantization
111//! - **Mixed**: Different precisions for different layers
112
113use crate::error::{ErrorBuilder, ErrorCode, FfiError};
114use serde::{Deserialize, Serialize};
115
116/// Quantization type/precision
117#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
118pub enum QuantizationType {
119    /// 8-bit signed integer (-128 to 127)
120    Int8,
121    /// 4-bit signed integer (-8 to 7)
122    Int4,
123    /// 8-bit unsigned integer (0 to 255)
124    Uint8,
125    /// 16-bit floating point (half precision)
126    Fp16,
127    /// BFloat16 (Brain floating point)
128    BFloat16,
129    /// Dynamic quantization (runtime)
130    Dynamic,
131}
132
133impl QuantizationType {
134    /// Get the number of bits used
135    pub fn bits(&self) -> u8 {
136        match self {
137            QuantizationType::Int4 => 4,
138            QuantizationType::Int8 | QuantizationType::Uint8 => 8,
139            QuantizationType::Fp16 | QuantizationType::BFloat16 => 16,
140            QuantizationType::Dynamic => 8, // Default to 8-bit
141        }
142    }
143
144    /// Get the theoretical compression ratio compared to FP32
145    pub fn compression_ratio(&self) -> f32 {
146        32.0 / self.bits() as f32
147    }
148
149    /// Get the range of values
150    pub fn value_range(&self) -> (f64, f64) {
151        match self {
152            QuantizationType::Int8 => (-128.0, 127.0),
153            QuantizationType::Int4 => (-8.0, 7.0),
154            QuantizationType::Uint8 => (0.0, 255.0),
155            QuantizationType::Fp16 => (-65504.0, 65504.0),
156            QuantizationType::BFloat16 => (-3.39e38, 3.39e38),
157            QuantizationType::Dynamic => (-128.0, 127.0),
158        }
159    }
160}
161
162/// Quantization granularity
163#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
164pub enum QuantizationGranularity {
165    /// Per-tensor quantization (single scale/zero-point for entire tensor)
166    PerTensor,
167    /// Per-channel quantization (scale/zero-point per output channel)
168    PerChannel,
169    /// Per-group quantization (for grouped operations)
170    PerGroup { group_size: usize },
171}
172
173/// Quantization scheme (symmetric vs asymmetric)
174#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
175pub enum QuantizationScheme {
176    /// Symmetric quantization (zero-point = 0)
177    Symmetric,
178    /// Asymmetric quantization (arbitrary zero-point)
179    Asymmetric,
180}
181
182/// Calibration method for determining quantization parameters
183#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
184pub enum CalibrationMethod {
185    /// Use min/max values from calibration data
186    MinMax,
187    /// Use percentiles to clip outliers (e.g., 0.1% and 99.9%)
188    Percentile { lower: f32, upper: f32 },
189    /// Use moving average of min/max
190    MovingAverage { momentum: f32 },
191    /// Minimize mean squared error
192    Mse,
193    /// Entropy-based calibration (KL divergence minimization)
194    Entropy,
195}
196
197impl Default for CalibrationMethod {
198    fn default() -> Self {
199        CalibrationMethod::Percentile {
200            lower: 0.01,
201            upper: 99.99,
202        }
203    }
204}
205
206/// Quantization parameters (scale and zero-point)
207#[derive(Debug, Clone, Serialize, Deserialize)]
208pub struct QuantizationParams {
209    /// Scaling factor
210    pub scale: f32,
211    /// Zero point (offset)
212    pub zero_point: i32,
213    /// Quantization type
214    pub qtype: QuantizationType,
215    /// Value range (min, max) from calibration
216    pub range: (f32, f32),
217}
218
219impl QuantizationParams {
220    /// Create new quantization parameters
221    ///
222    /// # Arguments
223    /// * `scale` - Scaling factor
224    /// * `zero_point` - Zero point offset
225    /// * `qtype` - Quantization type
226    pub fn new(scale: f32, zero_point: i32, qtype: QuantizationType) -> Self {
227        Self {
228            scale,
229            zero_point,
230            qtype,
231            range: (0.0, 0.0),
232        }
233    }
234
235    /// Compute quantization parameters from value range
236    ///
237    /// # Arguments
238    /// * `min_val` - Minimum value in the data
239    /// * `max_val` - Maximum value in the data
240    /// * `qtype` - Quantization type
241    /// * `scheme` - Quantization scheme (symmetric/asymmetric)
242    pub fn from_range(
243        min_val: f32,
244        max_val: f32,
245        qtype: QuantizationType,
246        scheme: QuantizationScheme,
247    ) -> Self {
248        let (qmin, qmax) = qtype.value_range();
249
250        let (scale, zero_point) = match scheme {
251            QuantizationScheme::Symmetric => {
252                // Symmetric: zero_point = 0
253                let max_abs = min_val.abs().max(max_val.abs());
254                let scale = (2.0 * max_abs) / (qmax - qmin) as f32;
255                (scale, 0)
256            }
257            QuantizationScheme::Asymmetric => {
258                // Asymmetric: map [min_val, max_val] to [qmin, qmax]
259                let scale = (max_val - min_val) / (qmax - qmin) as f32;
260                let zero_point = qmin as f32 - min_val / scale;
261                (scale, zero_point.round() as i32)
262            }
263        };
264
265        Self {
266            scale: scale.max(1e-8), // Avoid division by zero
267            zero_point,
268            qtype,
269            range: (min_val, max_val),
270        }
271    }
272
273    /// Quantize a floating-point value
274    pub fn quantize(&self, value: f32) -> i32 {
275        let (qmin, qmax) = self.qtype.value_range();
276        let quantized = (value / self.scale).round() + self.zero_point as f32;
277        quantized.clamp(qmin as f32, qmax as f32) as i32
278    }
279
280    /// Dequantize an integer value back to floating-point
281    pub fn dequantize(&self, quantized: i32) -> f32 {
282        (quantized - self.zero_point) as f32 * self.scale
283    }
284
285    /// Quantize an array of values
286    pub fn quantize_array(&self, values: &[f32]) -> Vec<i32> {
287        values.iter().map(|&v| self.quantize(v)).collect()
288    }
289
290    /// Dequantize an array of values
291    pub fn dequantize_array(&self, quantized: &[i32]) -> Vec<f32> {
292        quantized.iter().map(|&q| self.dequantize(q)).collect()
293    }
294}
295
296/// Quantization configuration
297#[derive(Debug, Clone, Serialize, Deserialize)]
298pub struct QuantizationConfig {
299    /// Target quantization type
300    pub qtype: QuantizationType,
301    /// Quantization granularity
302    pub granularity: QuantizationGranularity,
303    /// Quantization scheme
304    pub scheme: QuantizationScheme,
305    /// Calibration method
306    pub calibration: CalibrationMethod,
307    /// Whether to quantize weights
308    pub quantize_weights: bool,
309    /// Whether to quantize activations
310    pub quantize_activations: bool,
311    /// Whether to quantize biases (usually kept in higher precision)
312    pub quantize_biases: bool,
313    /// Layers to skip quantization (by name)
314    pub skip_layers: Vec<String>,
315    /// Force specific layers to use FP32 (e.g., output layer)
316    pub force_fp32_layers: Vec<String>,
317}
318
319impl QuantizationConfig {
320    /// Create a new quantization configuration
321    pub fn new(qtype: QuantizationType) -> Self {
322        Self {
323            qtype,
324            granularity: QuantizationGranularity::PerChannel,
325            scheme: QuantizationScheme::Asymmetric,
326            calibration: CalibrationMethod::default(),
327            quantize_weights: true,
328            quantize_activations: true,
329            quantize_biases: false, // Usually keep biases in FP32
330            skip_layers: Vec::new(),
331            force_fp32_layers: vec!["output".to_string()], // Keep output in FP32
332        }
333    }
334
335    /// Set quantization granularity
336    pub fn with_granularity(mut self, granularity: QuantizationGranularity) -> Self {
337        self.granularity = granularity;
338        self
339    }
340
341    /// Set quantization scheme
342    pub fn with_scheme(mut self, scheme: QuantizationScheme) -> Self {
343        self.scheme = scheme;
344        self
345    }
346
347    /// Set calibration method
348    pub fn with_calibration(mut self, calibration: CalibrationMethod) -> Self {
349        self.calibration = calibration;
350        self
351    }
352
353    /// Add layer to skip
354    pub fn skip_layer(mut self, layer_name: String) -> Self {
355        self.skip_layers.push(layer_name);
356        self
357    }
358
359    /// Add layer to force FP32
360    pub fn force_fp32_layer(mut self, layer_name: String) -> Self {
361        self.force_fp32_layers.push(layer_name);
362        self
363    }
364}
365
366impl Default for QuantizationConfig {
367    fn default() -> Self {
368        Self::new(QuantizationType::Int8)
369    }
370}
371
372/// Quantized tensor representation
373#[derive(Debug, Clone, Serialize, Deserialize)]
374pub struct QuantizedTensor {
375    /// Quantized values (as integers)
376    pub quantized_data: Vec<i32>,
377    /// Quantization parameters
378    pub params: QuantizationParams,
379    /// Original tensor shape
380    pub shape: Vec<usize>,
381    /// Original tensor name (optional)
382    pub name: Option<String>,
383}
384
385impl QuantizedTensor {
386    /// Create a new quantized tensor
387    pub fn new(quantized_data: Vec<i32>, params: QuantizationParams, shape: Vec<usize>) -> Self {
388        Self {
389            quantized_data,
390            params,
391            shape,
392            name: None,
393        }
394    }
395
396    /// Set tensor name
397    pub fn with_name(mut self, name: String) -> Self {
398        self.name = Some(name);
399        self
400    }
401
402    /// Get the size in bytes
403    pub fn size_bytes(&self) -> usize {
404        let bits_per_element = self.params.qtype.bits() as usize;
405        (self.quantized_data.len() * bits_per_element + 7) / 8
406    }
407
408    /// Get compression ratio compared to FP32
409    pub fn compression_ratio(&self) -> f32 {
410        let fp32_size = self.quantized_data.len() * 4; // 4 bytes per FP32
411        let quantized_size = self.size_bytes();
412        fp32_size as f32 / quantized_size as f32
413    }
414
415    /// Dequantize back to FP32
416    pub fn dequantize(&self) -> Vec<f32> {
417        self.params.dequantize_array(&self.quantized_data)
418    }
419
420    /// Compute quantization error metrics
421    pub fn quantization_error(&self, original: &[f32]) -> QuantizationError {
422        let dequantized = self.dequantize();
423
424        let mut mse = 0.0_f32;
425        let mut max_error = 0.0_f32;
426
427        for (orig, deq) in original.iter().zip(dequantized.iter()) {
428            let error = (orig - deq).abs();
429            mse += error * error;
430            max_error = max_error.max(error);
431        }
432
433        mse /= original.len() as f32;
434        let rmse = mse.sqrt();
435
436        // Signal-to-Quantization-Noise Ratio (SQNR)
437        let signal_power: f32 = original.iter().map(|x| x * x).sum::<f32>() / original.len() as f32;
438        let sqnr_db = 10.0 * (signal_power / mse).log10();
439
440        QuantizationError {
441            mse,
442            rmse,
443            max_error,
444            sqnr_db,
445        }
446    }
447}
448
449/// Quantization error metrics
450#[derive(Debug, Clone, Serialize, Deserialize)]
451pub struct QuantizationError {
452    /// Mean Squared Error
453    pub mse: f32,
454    /// Root Mean Squared Error
455    pub rmse: f32,
456    /// Maximum absolute error
457    pub max_error: f32,
458    /// Signal-to-Quantization-Noise Ratio (dB)
459    pub sqnr_db: f32,
460}
461
462/// Calibration dataset for determining quantization parameters
463#[derive(Debug, Clone)]
464pub struct CalibrationDataset {
465    /// Calibration samples
466    samples: Vec<Vec<f32>>,
467    /// Maximum number of samples to use
468    max_samples: usize,
469}
470
471impl CalibrationDataset {
472    /// Create a new calibration dataset
473    pub fn new(max_samples: usize) -> Self {
474        Self {
475            samples: Vec::new(),
476            max_samples,
477        }
478    }
479
480    /// Add a sample to the calibration dataset
481    pub fn add_sample(&mut self, sample: Vec<f32>) {
482        if self.samples.len() < self.max_samples {
483            self.samples.push(sample);
484        }
485    }
486
487    /// Get all samples
488    pub fn samples(&self) -> &[Vec<f32>] {
489        &self.samples
490    }
491
492    /// Compute statistics from calibration data
493    pub fn compute_statistics(&self) -> CalibrationStatistics {
494        let mut all_values: Vec<f32> = self
495            .samples
496            .iter()
497            .flat_map(|s| s.iter().copied())
498            .collect();
499
500        all_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
501
502        let min_val = all_values.first().copied().unwrap_or(0.0);
503        let max_val = all_values.last().copied().unwrap_or(0.0);
504
505        let mean = all_values.iter().sum::<f32>() / all_values.len() as f32;
506
507        let variance: f32 =
508            all_values.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / all_values.len() as f32;
509        let std_dev = variance.sqrt();
510
511        CalibrationStatistics {
512            min_val,
513            max_val,
514            mean,
515            std_dev,
516            num_samples: all_values.len(),
517        }
518    }
519
520    /// Get value at percentile
521    pub fn percentile(&self, percentile: f32) -> f32 {
522        let mut all_values: Vec<f32> = self
523            .samples
524            .iter()
525            .flat_map(|s| s.iter().copied())
526            .collect();
527        all_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
528
529        let index = ((percentile / 100.0) * all_values.len() as f32) as usize;
530        all_values
531            .get(index.min(all_values.len() - 1))
532            .copied()
533            .unwrap_or(0.0)
534    }
535}
536
537/// Statistics from calibration data
538#[derive(Debug, Clone, Serialize, Deserialize)]
539pub struct CalibrationStatistics {
540    pub min_val: f32,
541    pub max_val: f32,
542    pub mean: f32,
543    pub std_dev: f32,
544    pub num_samples: usize,
545}
546
547/// Post-Training Quantization (PTQ) quantizer
548#[derive(Debug, Clone)]
549pub struct Quantizer {
550    config: QuantizationConfig,
551    calibration_data: Option<CalibrationDataset>,
552}
553
554impl Quantizer {
555    /// Create a new quantizer
556    pub fn new(config: QuantizationConfig) -> Self {
557        Self {
558            config,
559            calibration_data: None,
560        }
561    }
562
563    /// Set calibration data
564    pub fn with_calibration_data(mut self, data: CalibrationDataset) -> Self {
565        self.calibration_data = Some(data);
566        self
567    }
568
569    /// Quantize a tensor (weights or activations)
570    ///
571    /// # Arguments
572    /// * `data` - FP32 tensor data
573    /// * `shape` - Tensor shape
574    /// * `name` - Optional tensor name
575    pub fn quantize_tensor(
576        &self,
577        data: &[f32],
578        shape: Vec<usize>,
579        name: Option<String>,
580    ) -> Result<QuantizedTensor, FfiError> {
581        // Check if this layer should be skipped
582        if let Some(ref n) = name {
583            if self.config.skip_layers.contains(n) || self.config.force_fp32_layers.contains(n) {
584                return Err(FfiError::Enhanced(
585                    ErrorBuilder::new(ErrorCode::OperationFailed)
586                        .message("Layer marked for FP32 preservation")
587                        .context("layer", n)
588                        .build(),
589                ));
590            }
591        }
592
593        // Determine quantization range
594        let (min_val, max_val) = match &self.calibration_data {
595            Some(calib) => match self.config.calibration {
596                CalibrationMethod::MinMax => {
597                    let stats = calib.compute_statistics();
598                    (stats.min_val, stats.max_val)
599                }
600                CalibrationMethod::Percentile { lower, upper } => {
601                    (calib.percentile(lower), calib.percentile(upper))
602                }
603                _ => {
604                    // Fallback to min/max from actual data
605                    let min = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
606                    let max = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
607                    (min, max)
608                }
609            },
610            None => {
611                // No calibration data, use min/max from tensor
612                let min = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
613                let max = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
614                (min, max)
615            }
616        };
617
618        // Compute quantization parameters
619        let params =
620            QuantizationParams::from_range(min_val, max_val, self.config.qtype, self.config.scheme);
621
622        // Quantize the data
623        let quantized_data = params.quantize_array(data);
624
625        Ok(QuantizedTensor::new(quantized_data, params, shape).with_name(name.unwrap_or_default()))
626    }
627
628    /// Get configuration
629    pub fn config(&self) -> &QuantizationConfig {
630        &self.config
631    }
632}
633
634/// Model compression statistics
635#[derive(Debug, Clone, Serialize, Deserialize)]
636pub struct CompressionStats {
637    /// Original model size (bytes)
638    pub original_size: usize,
639    /// Quantized model size (bytes)
640    pub quantized_size: usize,
641    /// Compression ratio
642    pub compression_ratio: f32,
643    /// Number of quantized layers
644    pub num_quantized_layers: usize,
645    /// Number of FP32 layers (preserved)
646    pub num_fp32_layers: usize,
647    /// Average quantization error (SQNR in dB)
648    pub avg_sqnr_db: f32,
649}
650
651impl CompressionStats {
652    /// Create new compression statistics
653    pub fn new() -> Self {
654        Self {
655            original_size: 0,
656            quantized_size: 0,
657            compression_ratio: 1.0,
658            num_quantized_layers: 0,
659            num_fp32_layers: 0,
660            avg_sqnr_db: 0.0,
661        }
662    }
663
664    /// Size reduction percentage
665    pub fn size_reduction_percent(&self) -> f32 {
666        (1.0 - (self.quantized_size as f32 / self.original_size as f32)) * 100.0
667    }
668}
669
670impl Default for CompressionStats {
671    fn default() -> Self {
672        Self::new()
673    }
674}
675
676#[cfg(test)]
677mod tests {
678    use super::*;
679
680    #[test]
681    fn test_quantization_type_bits() {
682        assert_eq!(QuantizationType::Int4.bits(), 4);
683        assert_eq!(QuantizationType::Int8.bits(), 8);
684        assert_eq!(QuantizationType::Fp16.bits(), 16);
685    }
686
687    #[test]
688    fn test_quantization_type_compression_ratio() {
689        assert_eq!(QuantizationType::Int8.compression_ratio(), 4.0); // 32/8
690        assert_eq!(QuantizationType::Int4.compression_ratio(), 8.0); // 32/4
691        assert_eq!(QuantizationType::Fp16.compression_ratio(), 2.0); // 32/16
692    }
693
694    #[test]
695    fn test_quantization_params_symmetric() {
696        let params = QuantizationParams::from_range(
697            -10.0,
698            10.0,
699            QuantizationType::Int8,
700            QuantizationScheme::Symmetric,
701        );
702
703        assert_eq!(params.zero_point, 0);
704        assert!(params.scale > 0.0);
705
706        // Test quantize/dequantize
707        let val = 5.0_f32;
708        let quantized = params.quantize(val);
709        let dequantized = params.dequantize(quantized);
710        assert!((val - dequantized).abs() < 0.1); // Small error tolerance
711    }
712
713    #[test]
714    fn test_quantization_params_asymmetric() {
715        let params = QuantizationParams::from_range(
716            0.0,
717            10.0,
718            QuantizationType::Uint8,
719            QuantizationScheme::Asymmetric,
720        );
721
722        assert!(params.scale > 0.0);
723
724        // Test quantize/dequantize
725        let val = 5.0_f32;
726        let quantized = params.quantize(val);
727        let dequantized = params.dequantize(quantized);
728        assert!((val - dequantized).abs() < 0.1);
729    }
730
731    #[test]
732    fn test_quantization_array() {
733        let params = QuantizationParams::from_range(
734            -1.0,
735            1.0,
736            QuantizationType::Int8,
737            QuantizationScheme::Symmetric,
738        );
739
740        let data = vec![-0.5, 0.0, 0.5, 1.0];
741        let quantized = params.quantize_array(&data);
742        let dequantized = params.dequantize_array(&quantized);
743
744        for (orig, deq) in data.iter().zip(dequantized.iter()) {
745            assert!((orig - deq).abs() < 0.1);
746        }
747    }
748
749    #[test]
750    fn test_quantized_tensor_creation() {
751        let data = vec![1, 2, 3, 4, 5, 6];
752        let params = QuantizationParams::new(0.1, 0, QuantizationType::Int8);
753        let shape = vec![2, 3];
754
755        let qtensor = QuantizedTensor::new(data.clone(), params.clone(), shape.clone());
756
757        assert_eq!(qtensor.quantized_data, data);
758        assert_eq!(qtensor.shape, shape);
759        assert!(qtensor.compression_ratio() > 1.0);
760    }
761
762    #[test]
763    fn test_quantized_tensor_size() {
764        let data = vec![0; 1000]; // 1000 elements
765        let params = QuantizationParams::new(1.0, 0, QuantizationType::Int8);
766        let qtensor = QuantizedTensor::new(data, params, vec![1000]);
767
768        // INT8: 1 byte per element
769        assert_eq!(qtensor.size_bytes(), 1000);
770
771        let params_int4 = QuantizationParams::new(1.0, 0, QuantizationType::Int4);
772        let qtensor_int4 = QuantizedTensor::new(vec![0; 1000], params_int4, vec![1000]);
773
774        // INT4: 0.5 bytes per element (packed)
775        assert_eq!(qtensor_int4.size_bytes(), 500);
776    }
777
778    #[test]
779    fn test_calibration_dataset() {
780        let mut dataset = CalibrationDataset::new(100);
781
782        dataset.add_sample(vec![1.0, 2.0, 3.0]);
783        dataset.add_sample(vec![4.0, 5.0, 6.0]);
784
785        assert_eq!(dataset.samples().len(), 2);
786
787        let stats = dataset.compute_statistics();
788        assert_eq!(stats.min_val, 1.0);
789        assert_eq!(stats.max_val, 6.0);
790        assert_eq!(stats.num_samples, 6);
791    }
792
793    #[test]
794    fn test_calibration_percentile() {
795        let mut dataset = CalibrationDataset::new(100);
796
797        // Add data from 1 to 100
798        for i in 1..=100 {
799            dataset.add_sample(vec![i as f32]);
800        }
801
802        let p50 = dataset.percentile(50.0);
803        assert!((p50 - 50.0).abs() < 5.0); // Approximately median
804
805        let p99 = dataset.percentile(99.0);
806        assert!(p99 > 95.0);
807    }
808
809    #[test]
810    fn test_quantization_config() {
811        let config = QuantizationConfig::new(QuantizationType::Int8)
812            .with_granularity(QuantizationGranularity::PerChannel)
813            .with_scheme(QuantizationScheme::Symmetric)
814            .skip_layer("layer1".to_string())
815            .force_fp32_layer("output".to_string());
816
817        assert_eq!(config.qtype, QuantizationType::Int8);
818        assert!(config.quantize_weights);
819        assert!(config.skip_layers.contains(&"layer1".to_string()));
820    }
821
822    #[test]
823    fn test_quantizer_basic() {
824        let config = QuantizationConfig::new(QuantizationType::Int8);
825        let quantizer = Quantizer::new(config);
826
827        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
828        let shape = vec![5];
829
830        let result = quantizer.quantize_tensor(&data, shape, Some("test".to_string()));
831        assert!(result.is_ok());
832
833        let qtensor = result.unwrap();
834        assert_eq!(qtensor.shape, vec![5]);
835        assert_eq!(qtensor.quantized_data.len(), 5);
836    }
837
838    #[test]
839    fn test_quantizer_with_calibration() {
840        let mut calib_data = CalibrationDataset::new(10);
841        calib_data.add_sample(vec![0.0, 1.0, 2.0, 3.0]);
842        calib_data.add_sample(vec![0.5, 1.5, 2.5, 3.5]);
843
844        let config = QuantizationConfig::new(QuantizationType::Int8).with_calibration(
845            CalibrationMethod::Percentile {
846                lower: 1.0,
847                upper: 99.0,
848            },
849        );
850
851        let quantizer = Quantizer::new(config).with_calibration_data(calib_data);
852
853        let data = vec![1.0, 2.0, 3.0];
854        let result = quantizer.quantize_tensor(&data, vec![3], None);
855        assert!(result.is_ok());
856    }
857
858    #[test]
859    fn test_quantization_error_metrics() {
860        let original = vec![1.0, 2.0, 3.0, 4.0, 5.0];
861
862        let params = QuantizationParams::from_range(
863            0.0,
864            5.0,
865            QuantizationType::Int8,
866            QuantizationScheme::Asymmetric,
867        );
868
869        let quantized_data = params.quantize_array(&original);
870        let qtensor = QuantizedTensor::new(quantized_data, params, vec![5]);
871
872        let error = qtensor.quantization_error(&original);
873
874        assert!(error.mse >= 0.0);
875        assert!(error.rmse >= 0.0);
876        assert!(error.max_error >= 0.0);
877        assert!(error.sqnr_db > 0.0); // Should have positive SQNR
878    }
879
880    #[test]
881    fn test_compression_stats() {
882        let mut stats = CompressionStats::new();
883        stats.original_size = 1000;
884        stats.quantized_size = 250;
885        stats.compression_ratio = 4.0;
886
887        assert_eq!(stats.size_reduction_percent(), 75.0);
888    }
889
890    #[test]
891    fn test_quantizer_skip_layer() {
892        let config =
893            QuantizationConfig::new(QuantizationType::Int8).skip_layer("skip_me".to_string());
894
895        let quantizer = Quantizer::new(config);
896
897        let data = vec![1.0, 2.0, 3.0];
898        let result = quantizer.quantize_tensor(&data, vec![3], Some("skip_me".to_string()));
899
900        assert!(result.is_err()); // Should error because layer is skipped
901    }
902}
torsh_ffi/quantization.rs

torsh_ffi/
quantization.rs
