rustytorch/crates/training/rtx-distributed/src/gradient_compression/advanced.rs

//! Advanced compression methods (OneBit, PowerSGD, DGC, Adaptive)

use super::compressor::GradientCompressor;
use super::config::CompressionType;
use super::quantization;
use super::sparsification;
use super::states::{AdaptiveStats, DgcState, PowerSGDState};
use super::types::CompressedGradient;
use crate::error::{DistributedError, Result};

// =============================================================================
// 1-bit SGD
// =============================================================================

pub fn compress_onebit(gradients: &[f32], shape: &[usize]) -> CompressedGradient {
    // Compute mean for threshold
    let mean: f32 = gradients.iter().sum::<f32>() / gradients.len() as f32;

    // Compute positive and negative means for reconstruction
    let (pos_sum, pos_count, neg_sum, neg_count) =
        gradients
            .iter()
            .fold((0.0f32, 0usize, 0.0f32, 0usize), |(ps, pc, ns, nc), &g| {
                if g >= mean {
                    (ps + g, pc + 1, ns, nc)
                } else {
                    (ps, pc, ns + g, nc + 1)
                }
            });

    let pos_mean = if pos_count > 0 {
        pos_sum / pos_count as f32
    } else {
        mean
    };
    let neg_mean = if neg_count > 0 {
        neg_sum / neg_count as f32
    } else {
        mean
    };

    // Pack bits (8 gradients per byte)
    let num_bytes = (gradients.len() + 7) / 8;
    let mut bytes = vec![0u8; num_bytes];

    for (i, &g) in gradients.iter().enumerate() {
        if g >= mean {
            bytes[i / 8] |= 1 << (i % 8);
        }
    }

    CompressedGradient {
        shape: shape.to_vec(),
        compression_type: CompressionType::OneBit,
        data: bytes,
        indices: None,
        scale: pos_mean,      // Store positive mean in scale
        zero_point: neg_mean, // Store negative mean in zero_point
        num_elements: gradients.len(),
        original_size: gradients.len() * 4,
    }
}

pub fn decompress_onebit(compressed: &CompressedGradient) -> Result<Vec<f32>> {
    let pos_mean = compressed.scale;
    let neg_mean = compressed.zero_point;

    let mut result = vec![0.0f32; compressed.num_elements];

    for i in 0..compressed.num_elements {
        let bit = (compressed.data[i / 8] >> (i % 8)) & 1;
        result[i] = if bit == 1 { pos_mean } else { neg_mean };
    }

    Ok(result)
}

// =============================================================================
// PowerSGD (Low-rank approximation)
// =============================================================================

pub fn compress_powersgd(
    compressor: &GradientCompressor,
    name: &str,
    gradients: &[f32],
    shape: &[usize],
) -> Result<CompressedGradient> {
    // PowerSGD requires 2D shape (flatten if needed)
    let (rows, cols) = if shape.len() >= 2 {
        (shape[0], shape[1..].iter().product())
    } else {
        (1, gradients.len())
    };

    let rank = compressor.config.powersgd_rank.min(rows.min(cols));

    // Get or create state
    let mut states = compressor.powersgd_states.write();
    let state = states
        .entry(name.to_string())
        .or_insert_with(|| PowerSGDState::new(rows, cols, rank));

    // Ensure state matches current shape
    if state.shape != (rows, cols) {
        *state = PowerSGDState::new(rows, cols, rank);
    }

    // PowerSGD algorithm:
    // 1. Q = M @ P (project gradients onto P)
    // 2. Orthogonalize Q
    // 3. P = M^T @ Q (update P)
    // 4. Orthogonalize P
    // 5. Send Q and P (low-rank factors)

    // Step 1: Q = M @ P
    for i in 0..rows {
        for j in 0..rank {
            let mut sum = 0.0f32;
            for k in 0..cols {
                sum += gradients[i * cols + k] * state.p_matrix[k * rank + j];
            }
            state.q_matrix[i * rank + j] = sum;
        }
    }

    // Step 2: Orthogonalize Q using Gram-Schmidt
    orthogonalize(&mut state.q_matrix, rows, rank);

    // Step 3: P = M^T @ Q
    let mut new_p = vec![0.0f32; cols * rank];
    for i in 0..cols {
        for j in 0..rank {
            let mut sum = 0.0f32;
            for k in 0..rows {
                sum += gradients[k * cols + i] * state.q_matrix[k * rank + j];
            }
            new_p[i * rank + j] = sum;
        }
    }
    state.p_matrix = new_p;

    // Step 4: Orthogonalize P
    orthogonalize(&mut state.p_matrix, cols, rank);

    // Pack Q and P into data
    let mut data = Vec::with_capacity((rows + cols) * rank * 4);
    for v in &state.q_matrix {
        data.extend_from_slice(&v.to_le_bytes());
    }
    for v in &state.p_matrix {
        data.extend_from_slice(&v.to_le_bytes());
    }

    Ok(CompressedGradient {
        shape: shape.to_vec(),
        compression_type: CompressionType::PowerSGD,
        data,
        indices: None,
        scale: rows as f32, // Store dimensions for reconstruction
        zero_point: cols as f32,
        num_elements: gradients.len(),
        original_size: gradients.len() * 4,
    })
}

pub fn decompress_powersgd(
    compressor: &GradientCompressor,
    compressed: &CompressedGradient,
) -> Result<Vec<f32>> {
    let rows = compressed.scale as usize;
    let cols = compressed.zero_point as usize;
    let total = rows * cols;

    if total != compressed.num_elements {
        return Err(DistributedError::communication(
            "powersgd",
            format!("Shape mismatch: {} vs {}", total, compressed.num_elements),
        ));
    }

    let rank = compressor.config.powersgd_rank.min(rows.min(cols));

    // Unpack Q and P from data
    let float_data: Vec<f32> = compressed
        .data
        .chunks_exact(4)
        .map(|chunk| {
            let bytes: [u8; 4] = chunk.try_into().unwrap();
            f32::from_le_bytes(bytes)
        })
        .collect();

    let q_size = rows * rank;
    let q_matrix = &float_data[..q_size];
    let p_matrix = &float_data[q_size..];

    // Reconstruct: M = Q @ P^T
    let mut result = vec![0.0f32; total];
    for i in 0..rows {
        for j in 0..cols {
            let mut sum = 0.0f32;
            for k in 0..rank {
                sum += q_matrix[i * rank + k] * p_matrix[j * rank + k];
            }
            result[i * cols + j] = sum;
        }
    }

    Ok(result)
}

/// Gram-Schmidt orthogonalization
fn orthogonalize(matrix: &mut [f32], rows: usize, cols: usize) {
    for j in 0..cols {
        // Subtract projections of previous columns
        for k in 0..j {
            let mut dot = 0.0f32;
            let mut norm_k = 0.0f32;
            for i in 0..rows {
                dot += matrix[i * cols + j] * matrix[i * cols + k];
                norm_k += matrix[i * cols + k] * matrix[i * cols + k];
            }
            if norm_k > 1e-10 {
                let scale = dot / norm_k;
                for i in 0..rows {
                    matrix[i * cols + j] -= scale * matrix[i * cols + k];
                }
            }
        }

        // Normalize
        let mut norm = 0.0f32;
        for i in 0..rows {
            norm += matrix[i * cols + j] * matrix[i * cols + j];
        }
        norm = norm.sqrt();
        if norm > 1e-10 {
            for i in 0..rows {
                matrix[i * cols + j] /= norm;
            }
        }
    }
}

// =============================================================================
// Deep Gradient Compression (DGC)
// =============================================================================

pub fn compress_dgc(
    compressor: &GradientCompressor,
    name: &str,
    gradients: &[f32],
    shape: &[usize],
) -> Result<CompressedGradient> {
    let num_elements = gradients.len();

    // Get or create DGC state
    let mut states = compressor.dgc_states.write();
    let state = states
        .entry(name.to_string())
        .or_insert_with(|| DgcState::new(num_elements, shape.to_vec(), 0.9));

    // Ensure state matches current shape
    if state.shape != shape {
        *state = DgcState::new(num_elements, shape.to_vec(), 0.9);
    }

    // DGC Algorithm:
    // 1. Add accumulated error to gradient
    // 2. Update momentum with gradient
    // 3. Apply momentum correction
    // 4. Select top-k values
    // 5. Update error buffer with unselected values

    // Step 1 & 2: Update momentum
    for (i, &g) in gradients.iter().enumerate() {
        let g_with_error = g + state.accumulated[i];
        state.momentum[i] = state.momentum_factor * state.momentum[i] + g_with_error;
    }

    // Step 3: Apply momentum correction (Nesterov-style)
    let corrected: Vec<f32> = state
        .momentum
        .iter()
        .map(|&m| {
            m / (1.0
                - state
                    .momentum_factor
                    .powi(*compressor.step.read() as i32 + 1))
        })
        .collect();

    // Step 4: Select top-k values
    let k = ((num_elements as f32 * compressor.config.ratio).ceil() as usize).max(1);

    let mut indexed: Vec<(usize, f32)> = corrected.iter().copied().enumerate().collect();
    indexed.sort_by(|a, b| b.1.abs().partial_cmp(&a.1.abs()).unwrap());
    indexed.truncate(k);

    // Create mask
    let mut mask = vec![false; num_elements];
    for (i, _) in &indexed {
        mask[*i] = true;
    }

    // Step 5: Update accumulated error with unselected values
    for i in 0..num_elements {
        if mask[i] {
            state.accumulated[i] = 0.0;
            state.momentum[i] = 0.0; // Reset momentum for selected
        } else {
            state.accumulated[i] = corrected[i];
        }
    }

    // Store mask for next iteration
    state.prev_mask = mask;

    // Pack selected values
    let indices: Vec<u32> = indexed.iter().map(|(i, _)| *i as u32).collect();
    let values: Vec<u8> = indexed.iter().flat_map(|(_, v)| v.to_le_bytes()).collect();

    Ok(CompressedGradient {
        shape: shape.to_vec(),
        compression_type: CompressionType::DeepGradientCompression,
        data: values,
        indices: Some(indices),
        scale: 1.0,
        zero_point: 0.0,
        num_elements,
        original_size: num_elements * 4,
    })
}

// =============================================================================
// Adaptive Compression
// =============================================================================

pub fn compress_adaptive(
    compressor: &GradientCompressor,
    name: &str,
    gradients: &[f32],
    shape: &[usize],
) -> Result<CompressedGradient> {
    // Compute gradient statistics
    let stats = compute_gradient_stats(gradients);

    // Update adaptive statistics
    {
        let mut adaptive = compressor.adaptive_stats.write();
        adaptive.insert(name.to_string(), stats.clone());
    }

    // Select compression strategy based on statistics
    let compression_type = select_adaptive_strategy(&stats);

    // Apply selected compression
    let result = match compression_type {
        CompressionType::TopK => {
            sparsification::compress_topk(gradients, shape, compressor.config.ratio)
        }
        CompressionType::Int8 => quantization::compress_int8(gradients, shape),
        CompressionType::OneBit => compress_onebit(gradients, shape),
        CompressionType::Fp16 => quantization::compress_fp16(gradients, shape),
        _ => compressor.no_compression(gradients, shape),
    };

    // Mark as adaptive (store actual type in indices as first element)
    let mut adaptive_result = result;
    adaptive_result.compression_type = CompressionType::Adaptive;
    // Store the actual compression type used in the first index
    let actual_type = compression_type as u8;
    if let Some(ref mut indices) = adaptive_result.indices {
        indices.insert(0, actual_type as u32);
    } else {
        adaptive_result.indices = Some(vec![actual_type as u32]);
    }

    Ok(adaptive_result)
}

fn compute_gradient_stats(gradients: &[f32]) -> AdaptiveStats {
    let n = gradients.len() as f32;
    if n == 0.0 {
        return AdaptiveStats::default();
    }

    let mean: f32 = gradients.iter().sum::<f32>() / n;
    let variance: f32 = gradients.iter().map(|g| (g - mean).powi(2)).sum::<f32>() / n;

    let threshold = 1e-6;
    let near_zero_count = gradients.iter().filter(|&&g| g.abs() < threshold).count();
    let sparsity = near_zero_count as f32 / n;

    let l2_norm: f32 = gradients.iter().map(|g| g.powi(2)).sum::<f32>().sqrt();
    let max_abs = gradients.iter().map(|g| g.abs()).fold(0.0f32, f32::max);

    AdaptiveStats {
        variance,
        sparsity,
        l2_norm,
        max_abs,
        num_samples: gradients.len(),
    }
}

fn select_adaptive_strategy(stats: &AdaptiveStats) -> CompressionType {
    // Decision tree for adaptive compression:
    // 1. High sparsity (>80%) -> TopK
    // 2. Low variance, uniform distribution -> 1-bit
    // 3. Medium variance -> INT8
    // 4. Otherwise -> FP16

    if stats.sparsity > 0.8 {
        CompressionType::TopK
    } else if stats.variance < 1e-4 && stats.max_abs < 0.1 {
        CompressionType::OneBit
    } else if stats.variance < 0.01 {
        CompressionType::Int8
    } else {
        CompressionType::Fp16
    }
}

pub fn decompress_adaptive(
    compressor: &GradientCompressor,
    compressed: &CompressedGradient,
) -> Result<Vec<f32>> {
    // Extract actual compression type from indices
    let actual_type = compressed
        .indices
        .as_ref()
        .and_then(|i| i.first())
        .copied()
        .unwrap_or(0) as u8;

    // Create a modified compressed gradient with the actual type
    let mut inner = compressed.clone();
    inner.compression_type = match actual_type {
        1 => CompressionType::TopK,
        2 => CompressionType::Int8,
        3 => CompressionType::OneBit,
        4 => CompressionType::Fp16,
        _ => CompressionType::None,
    };

    // Remove the type marker from indices if present
    if let Some(ref mut indices) = inner.indices {
        if !indices.is_empty() {
            indices.remove(0);
        }
        if indices.is_empty() {
            inner.indices = None;
        }
    }

    compressor.decompress(&inner)
}