foxhunt/crates/ml/tests/per_channel_quantization_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Per-Channel Quantization Tests
//!
//! Validates that per-channel quantization reduces quantization error from 2.5% to 1.5%
//! on attention weights and linear layer weights.

use candle_core::{DType, Device, Tensor};
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::MLError;
use tracing::info;

/// Test 1: Per-channel quantization reduces error on attention weights
#[cfg_attr(not(feature = "cuda"), ignore)]
#[test]
fn test_per_channel_attention_weight_accuracy() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");

    // Create attention weight matrix [256, 256] (out_channels, in_channels)
    let weight_data: Vec<f32> = (0..256 * 256)
        .map(|i| ((i as f32) * 0.01).sin() * 0.5)
        .collect();
    let weight = Tensor::from_slice(&weight_data, (256, 256), &device)?;

    // Test per-tensor quantization (per_channel = false)
    let config_per_tensor = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer_per_tensor = Quantizer::new(config_per_tensor, device.clone());
    let quantized_per_tensor = quantizer_per_tensor.quantize_tensor(&weight, "q_weight")?;
    let dequantized_per_tensor = quantizer_per_tensor.dequantize_tensor(&quantized_per_tensor)?;

    // Calculate per-tensor error
    let error_per_tensor = calculate_relative_error(&weight, &dequantized_per_tensor)?;

    // Test per-channel quantization (per_channel = true)
    let config_per_channel = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer_per_channel = Quantizer::new(config_per_channel, device.clone());
    let quantized_per_channel = quantizer_per_channel.quantize_tensor(&weight, "q_weight")?;

    // Verify per-channel params were stored
    assert!(
        quantizer_per_channel.has_per_channel_params("q_weight"),
        "Per-channel params should be stored"
    );

    // Dequantize using per-channel method
    let dequantized_per_channel =
        quantizer_per_channel.dequantize_tensor_per_channel(&quantized_per_channel, "q_weight")?;

    // Calculate per-channel error
    let error_per_channel = calculate_relative_error(&weight, &dequantized_per_channel)?;

    info!(error_per_tensor_pct = error_per_tensor * 100.0, error_per_channel_pct = error_per_channel * 100.0, "Quantization error: per-tensor vs per-channel");

    // Validate error reduction: per-channel should be lower than per-tensor
    assert!(
        error_per_channel < error_per_tensor,
        "Per-channel error {:.4}% should be lower than per-tensor error {:.4}%",
        error_per_channel * 100.0,
        error_per_tensor * 100.0
    );

    // Target validation: per-channel error should be < 1.5%
    assert!(
        error_per_channel < 0.015,
        "Per-channel error {:.4}% should be < 1.5%",
        error_per_channel * 100.0
    );

    Ok(())
}

/// Test 2: Per-channel quantization on linear layer weights
#[cfg_attr(not(feature = "cuda"), ignore)]
#[test]
fn test_per_channel_linear_layer_accuracy() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");

    // Create linear layer weight [128, 256] (out_features, in_features)
    let weight_data: Vec<f32> = (0..128 * 256)
        .map(|i| ((i as f32) * 0.02).cos() * 0.3)
        .collect();
    let weight = Tensor::from_slice(&weight_data, (128, 256), &device)?;

    // Per-channel quantization
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize
    let quantized = quantizer.quantize_tensor(&weight, "linear_weight")?;

    // Verify quantized dtype is U8
    assert_eq!(
        quantized.data.dtype(),
        DType::U8,
        "Quantized data should be U8"
    );

    // Verify shape is preserved
    assert_eq!(
        quantized.data.dims(),
        &[128, 256],
        "Shape should be preserved"
    );

    // Dequantize
    let dequantized = quantizer.dequantize_tensor_per_channel(&quantized, "linear_weight")?;

    // Calculate error
    let error = calculate_relative_error(&weight, &dequantized)?;

    info!(error_pct = error * 100.0, "Linear layer per-channel quantization error");

    // Validate error < 1.5%
    assert!(
        error < 0.015,
        "Per-channel error {:.4}% should be < 1.5%",
        error * 100.0
    );

    Ok(())
}

/// Test 3: Per-channel parameters storage and retrieval
#[cfg_attr(not(feature = "cuda"), ignore)]
#[test]
fn test_per_channel_params_storage() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");

    // Create weight [64, 128]
    let weight_data: Vec<f32> = (0..64 * 128).map(|i| (i as f32) * 0.01).collect();
    let weight = Tensor::from_slice(&weight_data, (64, 128), &device)?;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize
    let _quantized = quantizer.quantize_tensor(&weight, "test_weight")?;

    // Check params were stored
    assert!(
        quantizer.has_per_channel_params("test_weight"),
        "Per-channel params should be stored"
    );

    // Retrieve params
    let params = quantizer
        .get_per_channel_params("test_weight")
        .expect("Params should exist");

    // Verify params shape
    assert_eq!(
        params.scales.len(),
        64,
        "Should have 64 scales (one per output channel)"
    );
    assert_eq!(params.zero_points.len(), 64, "Should have 64 zero points");
    assert_eq!(params.min_vals.len(), 64, "Should have 64 min values");
    assert_eq!(params.max_vals.len(), 64, "Should have 64 max values");

    // Verify scales are positive
    for (idx, scale) in params.scales.iter().enumerate() {
        assert!(
            *scale > 0.0,
            "Scale {} should be positive, got {}",
            idx,
            scale
        );
    }

    // Verify zero points are within valid range for symmetric quantization
    for (idx, zp) in params.zero_points.iter().enumerate() {
        assert_eq!(
            *zp, 127,
            "Zero point {} should be 127 (symmetric), got {}",
            idx, zp
        );
    }

    Ok(())
}

/// Test 4: Per-channel quantization comparison with per-tensor
#[cfg_attr(not(feature = "cuda"), ignore)]
#[test]
fn test_per_channel_vs_per_tensor_comparison() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");

    // Create weight with varying scales across channels
    // Each channel has different magnitude to amplify per-channel benefit
    let mut weight_data = Vec::with_capacity(32 * 64);
    for channel in 0..32 {
        let scale = (channel + 1) as f32 * 0.1;
        for i in 0..64 {
            weight_data.push((i as f32 * 0.01).sin() * scale);
        }
    }
    let weight = Tensor::from_slice(&weight_data, (32, 64), &device)?;

    // Per-tensor quantization
    let config_per_tensor = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer_per_tensor = Quantizer::new(config_per_tensor, device.clone());
    let quantized_pt = quantizer_per_tensor.quantize_tensor(&weight, "w1")?;
    let dequantized_pt = quantizer_per_tensor.dequantize_tensor(&quantized_pt)?;
    let error_pt = calculate_relative_error(&weight, &dequantized_pt)?;

    // Per-channel quantization
    let config_per_channel = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer_per_channel = Quantizer::new(config_per_channel, device.clone());
    let quantized_pc = quantizer_per_channel.quantize_tensor(&weight, "w1")?;
    let dequantized_pc =
        quantizer_per_channel.dequantize_tensor_per_channel(&quantized_pc, "w1")?;
    let error_pc = calculate_relative_error(&weight, &dequantized_pc)?;

    info!(error_per_tensor_pct = error_pt * 100.0, error_per_channel_pct = error_pc * 100.0, "Variable-scale weights: per-tensor vs per-channel error");

    // Per-channel should be significantly better for variable-scale weights
    let improvement_ratio = error_pt / error_pc;
    assert!(
        improvement_ratio > 1.0,
        "Per-channel should be better than per-tensor (improvement ratio: {:.2}x)",
        improvement_ratio
    );

    // Per-channel should achieve target <1.5% error
    assert!(
        error_pc < 0.015,
        "Per-channel error {:.4}% should be < 1.5%",
        error_pc * 100.0
    );

    Ok(())
}

/// Test 5: Matmul with per-channel dequantized weights
#[cfg_attr(not(feature = "cuda"), ignore)]
#[test]
fn test_per_channel_matmul_integration() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");

    // Create weight [64, 128]
    let weight_data: Vec<f32> = (0..64 * 128)
        .map(|i| ((i as f32) * 0.01).sin() * 0.2)
        .collect();
    let weight = Tensor::from_slice(&weight_data, (64, 128), &device)?;

    // Create input [2, 128] (batch_size=2, in_features=128)
    let input_data: Vec<f32> = (0..2 * 128).map(|i| (i as f32) * 0.1).collect();
    let input = Tensor::from_slice(&input_data, (2, 128), &device)?;

    // F32 matmul (ground truth)
    let output_f32 = input.matmul(&weight.t()?)?;

    // Per-channel quantization
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize weight
    let quantized_weight = quantizer.quantize_tensor(&weight, "weight")?;

    // Dequantize for inference
    let dequantized_weight =
        quantizer.dequantize_tensor_per_channel(&quantized_weight, "weight")?;

    // INT8 matmul (with per-channel dequantization)
    let output_int8 = input.matmul(&dequantized_weight.t()?)?;

    // Calculate output error
    let error = calculate_relative_error(&output_f32, &output_int8)?;

    info!(error_pct = error * 100.0, "Matmul output error with per-channel dequantized weights");

    // Output error should be low
    assert!(
        error < 0.02,
        "Matmul output error {:.4}% should be < 2.0%",
        error * 100.0
    );

    Ok(())
}

/// Test 6: Per-channel quantization rejects non-2D tensors
#[cfg_attr(not(feature = "cuda"), ignore)]
#[test]
fn test_per_channel_rejects_non_2d() -> Result<(), MLError> {
    let device = Device::new_cuda(0).expect("CUDA required");

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Test 1D tensor
    let tensor_1d = Tensor::zeros((128,), DType::F32, &device)?;
    let result_1d = quantizer.quantize_tensor_per_channel(&tensor_1d, "test");
    assert!(
        result_1d.is_err(),
        "1D tensor should be rejected for per-channel quantization"
    );

    // Test 3D tensor
    let tensor_3d = Tensor::zeros((2, 64, 128), DType::F32, &device)?;
    let result_3d = quantizer.quantize_tensor_per_channel(&tensor_3d, "test");
    assert!(
        result_3d.is_err(),
        "3D tensor should be rejected for per-channel quantization"
    );

    Ok(())
}

/// Helper: Calculate relative error between two tensors
fn calculate_relative_error(original: &Tensor, reconstructed: &Tensor) -> Result<f32, MLError> {
    // Flatten both tensors
    let orig_vec = original.flatten_all()?.to_vec1::<f32>()?;
    let recon_vec = reconstructed.flatten_all()?.to_vec1::<f32>()?;

    assert_eq!(orig_vec.len(), recon_vec.len());

    // Calculate mean absolute error
    let mae: f32 = orig_vec
        .iter()
        .zip(recon_vec.iter())
        .map(|(o, r)| (o - r).abs())
        .sum::<f32>()
        / orig_vec.len() as f32;

    // Calculate mean of original values
    let orig_mean = orig_vec.iter().sum::<f32>().abs() / orig_vec.len() as f32;

    // Relative error
    let relative_error = if orig_mean > 1e-8 {
        mae / orig_mean
    } else {
        mae // If original is near zero, use absolute error
    };

    Ok(relative_error)
}