tnc/contractionpath/repartitioning/

simulated_annealing.rs

//! Repartitioning using simulated annealing algorithms.

use std::{
    iter::zip,
    time::{Duration, Instant},
};

use itertools::Itertools;
use ordered_float::NotNan;
use rand::{rngs::StdRng, seq::IteratorRandom, Rng, SeedableRng};
use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
use rustc_hash::FxHashSet;

use crate::{
    contractionpath::{
        communication_schemes::CommunicationScheme,
        contraction_cost::{compute_memory_requirements, contract_size_tensors_bytes},
        paths::{
            cotengrust::{Cotengrust, OptMethod},
            ContractionPathResult, Pathfinder,
        },
        repartitioning::compute_solution,
        SimplePath,
    },
    tensornetwork::{
        partitioning::partition_tensor_network,
        tensor::{CompositeTensor, LeafTensor},
    },
};

type ScoreType = NotNan<f64>;

/// Number of threads to use for processing candidate solutions in parallel. This is
/// a constant (and not hardware-aware) for reproducibility.
const PROCESSING_THREADS: usize = 48;

/// OptModel is a trait that defines requirements to be used with optimization algorithm
pub trait OptModel: Sync + Send {
    /// Type of the Solution
    type SolutionType: Clone + Sync + Send;

    /// Generate a new trial solution from current solution
    fn generate_trial_solution<R: Rng>(
        &self,
        current_solution: Self::SolutionType,
        rng: &mut R,
    ) -> Self::SolutionType;

    /// Evaluate the score of the solution
    fn evaluate<R: Rng>(&self, solution: &Self::SolutionType, rng: &mut R) -> ScoreType;
}

/// Optimizer that implements the simulated annealing algorithm
#[derive(Clone, Copy)]
pub struct SimulatedAnnealingOptimizer {
    /// Number of candidate solutions to generate and evaluate in each iteration.
    n_trials: usize,
    /// Total duration to take for the optimization
    max_time: Duration,
    /// Number of steps to take in each temperature iteration.
    n_steps: usize,
    /// Number of iterations without improvement after which the algorithm should
    /// restart from the best solution found so far.
    restart_iter: usize,
    /// The initial temperature to start the annealing process with.
    initial_temperature: f64,
    /// The final temperature to reach at the end of the annealing process.
    final_temperature: f64,
}

/// Linearly interpolates between two numbers based on parameter `t`.
///
/// Computes `start + (end - start) * t`.
#[inline]
fn linear_interpolation(start: f64, end: f64, t: f64) -> f64 {
    (end - start).mul_add(t, start)
}

impl SimulatedAnnealingOptimizer {
    /// Start optimization with given temperature range
    ///
    /// - `model` : the model to optimize
    /// - `initial_solution` : the initial solution to start optimization.
    /// - `n_iter`: maximum iterations
    #[allow(clippy::too_many_arguments)]
    fn optimize_with_temperature<M, R>(
        &self,
        model: &M,
        initial_solution: M::SolutionType,
        rng: &mut R,
    ) -> (M::SolutionType, ScoreType)
    where
        M: OptModel,
        R: Rng,
    {
        let mut current_score = model.evaluate(&initial_solution, rng);
        let mut current_solution = initial_solution;
        let mut best_solution = current_solution.clone();
        let mut best_score = current_score;
        let mut last_improvement = 0;
        let steps_per_thread = self.n_steps.div_ceil(self.n_trials);

        let log_start = self.initial_temperature.log2();
        let log_end = self.final_temperature.log2();
        let total_seconds = self.max_time.as_secs_f64();
        let mut temperature = self.initial_temperature;
        let mut rngs = (0..self.n_trials)
            .map(|_| StdRng::seed_from_u64(rng.random()))
            .collect_vec();
        let end_time = Instant::now() + self.max_time;
        loop {
            // Generate and evaluate candidate solutions to find the minimum objective
            let (_, trial_solution, trial_score) = rngs
                .par_iter_mut()
                .enumerate()
                .map(|(index, rng)| {
                    let mut trial_score = current_score;
                    let mut trial_solution = current_solution.clone();
                    for _ in 0..steps_per_thread {
                        let solution = model.generate_trial_solution(trial_solution.clone(), rng);
                        let score = model.evaluate(&solution, rng);

                        let diff = (score / trial_score).log2();
                        let acceptance_probability = (-diff / temperature).exp();
                        let random_value = rng.random();

                        if acceptance_probability >= random_value {
                            trial_solution = solution;
                            trial_score = score;
                        }
                    }
                    (index, trial_solution, trial_score)
                })
                .min_by_key(|(index, _, score)| (*score, *index))
                .unwrap();

            current_score = trial_score;
            current_solution = trial_solution;

            // Update the best solution if the current solution is better
            if current_score < best_score {
                best_solution = current_solution.clone();
                best_score = current_score;
                last_improvement = 0;
            }

            last_improvement += 1;

            // Check if we should restart from the best solution
            if last_improvement == self.restart_iter {
                current_solution = best_solution.clone();
                current_score = best_score;
            }

            // Estimate the number of remaining iterations and adapt the temperature
            let now = Instant::now();
            if now > end_time {
                // We've reached the time limit
                break;
            }
            let remaining_time = (end_time - now).as_secs_f64();
            let progress = 1.0 - remaining_time / total_seconds;
            temperature = linear_interpolation(log_start, log_end, progress).exp2();
        }

        (best_solution, best_score)
    }
}

/// Common evaluation method for all simulated annealing methods.
fn evaluate_partitioning<R>(
    tensor: &CompositeTensor,
    partitioning: &[usize],
    communication_scheme: CommunicationScheme,
    memory_limit: Option<f64>,
    rng: &mut R,
) -> NotNan<f64>
where
    R: Rng,
{
    // Construct the tensor network and contraction path from the partitioning
    let (partitioned_tn, path, parallel_cost, _) =
        compute_solution(tensor, partitioning, communication_scheme, Some(rng));

    // If the memory limit is exceeded, return infinity
    if let Some(limit) = memory_limit {
        // Compute memory usage
        let mem = compute_memory_requirements(
            partitioned_tn.tensors(),
            &path,
            contract_size_tensors_bytes,
        );

        if mem > limit {
            return NotNan::new(f64::INFINITY).unwrap();
        }
    }
    NotNan::new(parallel_cost).unwrap()
}

/// A simulated annealing model that moves a random tensor between random partitions.
pub struct NaivePartitioningModel<'a> {
    pub tensor: &'a CompositeTensor,
    pub num_partitions: usize,
    pub communication_scheme: CommunicationScheme,
    pub memory_limit: Option<f64>,
}

impl OptModel for NaivePartitioningModel<'_> {
    type SolutionType = Vec<usize>;

    fn generate_trial_solution<R: Rng>(
        &self,
        mut current_solution: Self::SolutionType,
        rng: &mut R,
    ) -> Self::SolutionType {
        let tensor_index = rng.random_range(0..current_solution.len());
        let current_partition = current_solution[tensor_index];
        let new_partition = loop {
            let b = rng.random_range(0..self.num_partitions);
            if b != current_partition {
                break b;
            }
        };
        current_solution[tensor_index] = new_partition;
        current_solution
    }

    fn evaluate<R: Rng>(&self, solution: &Self::SolutionType, rng: &mut R) -> ScoreType {
        evaluate_partitioning(
            self.tensor,
            solution,
            self.communication_scheme,
            self.memory_limit,
            rng,
        )
    }
}

/// A simulated annealing model that moves a random subtree between random partitions.
pub struct NaiveIntermediatePartitioningModel<'a> {
    pub tensor: &'a CompositeTensor,
    pub num_partitions: usize,
    pub communication_scheme: CommunicationScheme,
    pub memory_limit: Option<f64>,
}

impl OptModel for NaiveIntermediatePartitioningModel<'_> {
    type SolutionType = (Vec<usize>, Vec<SimplePath>);

    fn generate_trial_solution<R: Rng>(
        &self,
        current_solution: Self::SolutionType,
        rng: &mut R,
    ) -> Self::SolutionType {
        let (mut partitioning, mut contraction_paths) = current_solution;

        // Select source partition (with more than one tensor)
        let source_partition = contraction_paths
            .iter()
            .enumerate()
            .filter_map(|(contraction_id, contraction)| {
                if contraction.len() >= 3 {
                    Some(contraction_id)
                } else {
                    None
                }
            })
            .choose(rng);

        let Some(source_partition) = source_partition else {
            // No viable partitions, return the current solution
            return (partitioning, contraction_paths);
        };

        // Select random tensor contraction in source partition
        let pair_index = rng.random_range(0..contraction_paths[source_partition].len() - 1);
        let (i, j) = contraction_paths[source_partition][pair_index];
        let mut tensor_leaves = FxHashSet::from_iter([i, j]);

        // Gather all tensors that contribute to the selected contraction
        for (i, j) in contraction_paths[source_partition]
            .iter()
            .take(pair_index)
            .rev()
        {
            if tensor_leaves.contains(i) {
                tensor_leaves.insert(*j);
            }
        }

        let mut shifted_indices = Vec::with_capacity(tensor_leaves.len());
        for (partition_tensor_index, (i, _partition)) in partitioning
            .iter()
            .enumerate()
            .filter(|(_, partition)| *partition == &source_partition)
            .enumerate()
        {
            if tensor_leaves.contains(&partition_tensor_index) {
                shifted_indices.push(i);
            }
        }

        // Select random target partition
        let target_partition = loop {
            let b = rng.random_range(0..self.num_partitions);
            if b != source_partition {
                break b;
            }
        };

        // Change partition
        for index in shifted_indices {
            partitioning[index] = target_partition;
        }

        // Recompute the contraction path for both partitions
        let mut from_tensor = CompositeTensor::default();
        let mut to_tensor = CompositeTensor::default();
        for (partition_index, tensor) in zip(&partitioning, self.tensor.tensors()) {
            if *partition_index == source_partition {
                from_tensor.push_tensor(tensor.clone());
            } else if *partition_index == target_partition {
                to_tensor.push_tensor(tensor.clone());
            }
        }

        let mut from_opt = Cotengrust::new(OptMethod::Greedy);
        let result = from_opt.find_path(&from_tensor);
        let from_path = result.replace_path();
        contraction_paths[source_partition] = from_path.into_simple();

        let mut to_opt = Cotengrust::new(OptMethod::Greedy);
        let result = to_opt.find_path(&to_tensor);
        let to_path = result.replace_path();
        contraction_paths[target_partition] = to_path.into_simple();

        (partitioning, contraction_paths)
    }

    fn evaluate<R: Rng>(&self, solution: &Self::SolutionType, rng: &mut R) -> ScoreType {
        evaluate_partitioning(
            self.tensor,
            &solution.0,
            self.communication_scheme,
            self.memory_limit,
            rng,
        )
    }
}

/// A simulated annealing model that moves a random tensor to the partition that
/// maximizes memory reduction.
pub struct LeafPartitioningModel<'a> {
    pub tensor: &'a CompositeTensor,
    pub communication_scheme: CommunicationScheme,
    pub memory_limit: Option<f64>,
}

impl OptModel for LeafPartitioningModel<'_> {
    type SolutionType = (Vec<usize>, Vec<LeafTensor>);

    fn generate_trial_solution<R: Rng>(
        &self,
        current_solution: Self::SolutionType,
        rng: &mut R,
    ) -> Self::SolutionType {
        let (mut partitioning, mut partition_tensors) = current_solution;
        let tensor_index = rng.random_range(0..partitioning.len());
        let shifted_tensor = self.tensor.tensor(tensor_index).as_leaf().unwrap();
        let source_partition = partitioning[tensor_index];

        let (new_partition, _) = partition_tensors
            .iter()
            .enumerate()
            .filter_map(|(i, partition_tensor)| {
                if i != source_partition {
                    Some((
                        i,
                        (shifted_tensor ^ partition_tensor).size() - partition_tensor.size(),
                    ))
                } else {
                    // Don't consider old partition as move target (would be a NOOP)
                    None
                }
            })
            .min_by(|a, b| a.1.total_cmp(&b.1))
            .unwrap();

        partitioning[tensor_index] = new_partition;
        partition_tensors[source_partition] ^= shifted_tensor;
        partition_tensors[new_partition] ^= shifted_tensor;
        (partitioning, partition_tensors)
    }

    fn evaluate<R: Rng>(&self, solution: &Self::SolutionType, rng: &mut R) -> ScoreType {
        evaluate_partitioning(
            self.tensor,
            &solution.0,
            self.communication_scheme,
            self.memory_limit,
            rng,
        )
    }
}

/// A simulated annealing model that moves a random subtree to the partition that
/// maximizes memory reduction.
pub struct IntermediatePartitioningModel<'a> {
    pub tensor: &'a CompositeTensor,
    pub communication_scheme: CommunicationScheme,
    pub memory_limit: Option<f64>,
}

impl IntermediatePartitioningModel<'_> {
    pub fn compute_initial_solution(
        &self,
        initial_partitioning: &[usize],
        initial_contraction_paths: Option<Vec<SimplePath>>,
    ) -> <IntermediatePartitioningModel<'_> as OptModel>::SolutionType {
        let partitioned_tn = partition_tensor_network(self.tensor.clone(), initial_partitioning);

        // Get the partition tensors
        let partition_tensors = partitioned_tn
            .tensors()
            .iter()
            .map(|t| t.as_composite().unwrap().external_tensor())
            .collect_vec();

        let contraction_paths = initial_contraction_paths.unwrap_or_else(|| {
            // Find contraction paths using Greedy
            partitioned_tn
                .tensors()
                .iter()
                .map(|t| {
                    // Find path for this partition
                    let mut opt = Cotengrust::new(OptMethod::Greedy);
                    let result = opt.find_path(t.as_composite().unwrap());
                    let path = result.replace_path();
                    path.into_simple()
                })
                .collect()
        });

        // Find a contraction path for the partitioned tensor network
        (
            initial_partitioning.to_vec(),
            partition_tensors,
            contraction_paths,
        )
    }
}

impl OptModel for IntermediatePartitioningModel<'_> {
    type SolutionType = (Vec<usize>, Vec<LeafTensor>, Vec<SimplePath>);

    fn generate_trial_solution<R: Rng>(
        &self,
        current_solution: Self::SolutionType,
        rng: &mut R,
    ) -> Self::SolutionType {
        let (mut partitioning, mut partition_tensors, mut contraction_paths) = current_solution;

        // Select source partition (with more than one tensor)
        let source_partition = contraction_paths
            .iter()
            .enumerate()
            .filter_map(|(contraction_id, contraction)| {
                if contraction.len() >= 3 {
                    Some(contraction_id)
                } else {
                    None
                }
            })
            .choose(rng);

        let Some(source_partition) = source_partition else {
            // No viable partitions, return the current solution
            return (partitioning, partition_tensors, contraction_paths);
        };

        // Select random tensor contraction in source partition
        let pair_index = rng.random_range(0..contraction_paths[source_partition].len() - 1);
        let (i, j) = contraction_paths[source_partition][pair_index];
        let mut tensor_leaves = FxHashSet::from_iter([i, j]);

        // Gather all tensors that contribute to the selected contraction
        for (i, j) in contraction_paths[source_partition]
            .iter()
            .take(pair_index)
            .rev()
        {
            if tensor_leaves.contains(i) {
                tensor_leaves.insert(*j);
            }
        }

        let mut shifted_tensor = LeafTensor::default();
        let mut shifted_indices = Vec::with_capacity(tensor_leaves.len());
        for (partition_tensor_index, (i, _partition)) in partitioning
            .iter()
            .enumerate()
            .filter(|(_, partition)| *partition == &source_partition)
            .enumerate()
        {
            if tensor_leaves.contains(&partition_tensor_index) {
                shifted_tensor ^= self.tensor.tensor(i).as_leaf().unwrap();
                shifted_indices.push(i);
            }
        }

        // Find best target partition
        // Cost function is actually quite important!!
        let (target_partition, _) = partition_tensors
            .iter()
            .enumerate()
            .filter_map(|(i, partition_tensor)| {
                if i != source_partition {
                    Some((
                        i,
                        (&shifted_tensor ^ partition_tensor).size() - partition_tensor.size(),
                    ))
                } else {
                    // Don't consider old partition as move target (would be a NOOP)
                    None
                }
            })
            .min_by(|a, b| a.1.total_cmp(&b.1))
            .unwrap();

        // Change partition
        for index in shifted_indices {
            partitioning[index] = target_partition;
        }

        // Recompute the tensors for both partitions
        partition_tensors[source_partition] ^= &shifted_tensor;
        partition_tensors[target_partition] ^= &shifted_tensor;

        // Recompute the contraction path for both partitions
        let mut from_tensor = CompositeTensor::default();
        let mut to_tensor = CompositeTensor::default();
        for (partition_index, tensor) in zip(&partitioning, self.tensor.tensors()) {
            if *partition_index == source_partition {
                from_tensor.push_tensor(tensor.clone());
            } else if *partition_index == target_partition {
                to_tensor.push_tensor(tensor.clone());
            }
        }

        let mut from_opt = Cotengrust::new(OptMethod::Greedy);
        let result = from_opt.find_path(&from_tensor);
        let from_path = result.replace_path();
        contraction_paths[source_partition] = from_path.into_simple();

        let mut to_opt = Cotengrust::new(OptMethod::Greedy);
        let result = to_opt.find_path(&to_tensor);
        let to_path = result.replace_path();
        contraction_paths[target_partition] = to_path.into_simple();

        (partitioning, partition_tensors, contraction_paths)
    }

    fn evaluate<R: Rng>(&self, solution: &Self::SolutionType, rng: &mut R) -> ScoreType {
        evaluate_partitioning(
            self.tensor,
            &solution.0,
            self.communication_scheme,
            self.memory_limit,
            rng,
        )
    }
}

/// Runs simulated annealing to find a better partitioning.
pub fn balance_partitions<R, M>(
    model: M,
    initial_solution: M::SolutionType,
    rng: &mut R,
    max_time: Duration,
) -> (M::SolutionType, ScoreType)
where
    R: Rng,
    M: OptModel,
{
    let optimizer = SimulatedAnnealingOptimizer {
        n_trials: PROCESSING_THREADS,
        max_time,
        n_steps: PROCESSING_THREADS * 10,
        restart_iter: 50,
        initial_temperature: 2.0,
        final_temperature: 0.05,
    };
    optimizer.optimize_with_temperature(&model, initial_solution, rng)
}

#[cfg(test)]
mod tests {
    use approx::assert_abs_diff_eq;

    use super::*;

    #[test]
    fn simple_linear_interpolation() {
        assert_abs_diff_eq!(linear_interpolation(0., 6., 0.5), 3.0);
        assert_abs_diff_eq!(linear_interpolation(-1.0, 4.0, 0.2), 0.0);
        assert_abs_diff_eq!(linear_interpolation(-7.0, -6.0, 0.0), -7.0);
        assert_abs_diff_eq!(linear_interpolation(3.0, 5.0, 1.0), 5.0);
    }

    #[test]
    fn small_leaf_partitioning() {
        let t1 = LeafTensor::new_from_const(vec![0, 1], 2);
        let t2 = LeafTensor::new_from_const(vec![2, 3], 2);
        let t3 = LeafTensor::new_from_const(vec![0, 1, 4], 2);
        let t4 = LeafTensor::new_from_const(vec![2, 3, 4], 2);
        let tn = CompositeTensor::new(vec![t1.clone(), t2.clone(), t3.clone(), t4.clone()]);
        let tn1 = CompositeTensor::new(vec![t1, t2]);
        let tn2 = CompositeTensor::new(vec![t3, t4]);
        let initial_partitioning = vec![0, 0, 1, 1];
        let initial_partitions = vec![tn1.external_tensor(), tn2.external_tensor()];
        let mut rng = StdRng::seed_from_u64(42);

        let ((partitioning, _partitions), _) = balance_partitions(
            LeafPartitioningModel {
                tensor: &tn,
                communication_scheme: CommunicationScheme::Greedy,
                memory_limit: None,
            },
            (initial_partitioning, initial_partitions),
            &mut rng,
            Duration::from_secs(2),
        );
        // Normalize for comparability
        let ref_partitioning = if partitioning[0] == 0 {
            [0, 1, 0, 1]
        } else {
            [1, 0, 1, 0]
        };
        assert_eq!(partitioning, ref_partitioning);
    }
}