tnc/contractionpath/repartitioning/
genetic.rs

1//! Repartitioning using a genetic algorithm.
2
3use genetic_algorithm::{
4    crossover::CrossoverUniform,
5    fitness::{Fitness, FitnessChromosome, FitnessOrdering, FitnessValue},
6    genotype::{Genotype, RangeGenotype},
7    mutate::MutateSingleGene,
8    select::SelectTournament,
9    strategy::{evolve::Evolve, Strategy},
10};
11use ordered_float::NotNan;
12use rand::rngs::StdRng;
13
14use crate::{
15    contractionpath::{
16        communication_schemes::CommunicationScheme,
17        contraction_cost::{compute_memory_requirements, contract_size_tensors_bytes},
18        repartitioning::compute_solution,
19    },
20    tensornetwork::tensor::CompositeTensor,
21};
22
23#[derive(Clone, Debug)]
24struct PartitioningFitness<'a> {
25    tensor: &'a CompositeTensor,
26    communication_scheme: CommunicationScheme,
27    memory_limit: Option<f64>,
28}
29
30impl PartitioningFitness<'_> {
31    fn calculate_fitness(&self, partitioning: &[usize]) -> NotNan<f64> {
32        // Construct the tensor network and contraction path from the partitioning
33        let (partitioned_tn, path, cost, _) =
34            compute_solution::<StdRng>(self.tensor, partitioning, self.communication_scheme, None);
35
36        // Compute memory usage
37        let mem = compute_memory_requirements(
38            partitioned_tn.tensors(),
39            &path,
40            contract_size_tensors_bytes,
41        );
42
43        // If the memory limit is exceeded, return infinity
44        let score = if self.memory_limit.is_some_and(|limit| mem > limit) {
45            f64::INFINITY
46        } else {
47            cost
48        };
49        NotNan::new(score).unwrap()
50    }
51}
52
53impl Fitness for PartitioningFitness<'_> {
54    type Genotype = RangeGenotype<usize>;
55
56    fn calculate_for_chromosome(
57        &mut self,
58        chromosome: &FitnessChromosome<Self>,
59        _genotype: &Self::Genotype,
60    ) -> Option<FitnessValue> {
61        Some(self.calculate_fitness(&chromosome.genes))
62    }
63}
64
65/// Balances partitions using a genetic algorithm. Finds the partitioning that reduces
66/// the total contraction cost.
67pub fn balance_partitions(
68    tensor: &CompositeTensor,
69    num_partitions: usize,
70    initial_partitioning: &[usize],
71    communication_scheme: CommunicationScheme,
72    memory_limit: Option<f64>,
73) -> (Vec<usize>, f64) {
74    // Chromosomes: Possible partitions, e.g. [0, 1, 0, 2, 2, 1, 0, 0, 1, 1]
75    // Genes: tensor (in vector)
76    // Alleles: partition id
77
78    let num_tensors = initial_partitioning.len();
79
80    let genotype = RangeGenotype::builder()
81        .with_genes_size(num_tensors)
82        .with_allele_range(0..=num_partitions - 1)
83        .with_seed_genes_list(vec![initial_partitioning.to_vec()])
84        .build()
85        .unwrap();
86
87    let fitness = PartitioningFitness {
88        tensor,
89        communication_scheme,
90        memory_limit,
91    };
92
93    let evolve = Evolve::builder()
94        .with_genotype(genotype)
95        .with_target_population_size(100)
96        .with_max_stale_generations(100)
97        .with_fitness(fitness)
98        .with_fitness_ordering(FitnessOrdering::Minimize)
99        .with_mutate(MutateSingleGene::new(0.2))
100        .with_crossover(CrossoverUniform::new(1.0, 1.0))
101        .with_select(SelectTournament::new(1.0, 0.02, 4))
102        // .with_reporter(EvolveReporterDuration::new())
103        .with_par_fitness(true)
104        .with_rng_seed_from_u64(0)
105        .call()
106        .unwrap();
107
108    evolve
109        .best_genes_and_fitness_score()
110        .map(|(partitioning, score)| (partitioning, score.into_inner()))
111        .unwrap()
112}
113
114#[cfg(test)]
115mod tests {
116    use crate::tensornetwork::tensor::LeafTensor;
117
118    use super::*;
119
120    #[test]
121    fn small_partitioning() {
122        let t1 = LeafTensor::new_from_const(vec![0, 1], 2);
123        let t2 = LeafTensor::new_from_const(vec![2, 3], 2);
124        let t3 = LeafTensor::new_from_const(vec![0, 1, 4], 2);
125        let t4 = LeafTensor::new_from_const(vec![2, 3, 4], 2);
126        let tn = CompositeTensor::new(vec![t1, t2, t3, t4]);
127        let initial_partitioning = vec![0, 0, 1, 1];
128
129        let (partitioning, _) = balance_partitions(
130            &tn,
131            2,
132            &initial_partitioning,
133            CommunicationScheme::RandomGreedy,
134            None,
135        );
136        // Normalize for comparability
137        let ref_partitioning = if partitioning[0] == 0 {
138            [0, 1, 0, 1]
139        } else {
140            [1, 0, 1, 0]
141        };
142        assert_eq!(partitioning, ref_partitioning);
143    }
144}
tnc/contractionpath/repartitioning/genetic.rs

tnc/contractionpath/repartitioning/
genetic.rs
