tnc/contractionpath/paths/
weighted_branchbound.rs

1use std::collections::BinaryHeap;
2
3use itertools::Itertools;
4use rustc_hash::FxHashMap;
5
6use crate::{
7    contractionpath::{
8        candidates::Candidate,
9        contraction_cost::{contract_op_cost_tensors, contract_size_tensors},
10        paths::{BasicContractionPathResult, ContractionPathResult, CostType, Pathfinder},
11        ssa_ordering, ContractionPath,
12    },
13    tensornetwork::tensor::{CompositeTensor, LeafTensor},
14    utils::traits::HashMapInsertNew,
15};
16
17/// A struct with an [`Pathfinder`] implementation that explores possible pair contractions in a depth-first manner.
18pub struct WeightedBranchBound {
19    nbranch: Option<usize>,
20    cutoff_flops_factor: f64,
21    minimize: CostType,
22    best_flops: f64,
23    best_size: f64,
24    best_path: ContractionPath,
25    best_progress: FxHashMap<usize, f64>,
26    largest_latency: f64,
27    result_cache: FxHashMap<(usize, usize), (usize, f64, f64)>,
28    comm_cache: FxHashMap<usize, f64>,
29    tensor_cache: FxHashMap<usize, LeafTensor>,
30}
31
32impl WeightedBranchBound {
33    pub fn new(
34        nbranch: Option<usize>,
35        cutoff_flops_factor: f64,
36        latency_map: FxHashMap<usize, f64>,
37        minimize: CostType,
38    ) -> Self {
39        Self {
40            nbranch,
41            cutoff_flops_factor,
42            minimize,
43            best_flops: f64::INFINITY,
44            best_size: f64::INFINITY,
45            best_path: ContractionPath::default(),
46            best_progress: FxHashMap::default(),
47            largest_latency: Default::default(),
48            result_cache: FxHashMap::default(),
49            comm_cache: latency_map,
50            tensor_cache: FxHashMap::default(),
51        }
52    }
53
54    fn assess_candidate(
55        &mut self,
56        mut i: usize,
57        mut j: usize,
58        size: f64,
59        remaining_len: usize,
60    ) -> Option<Candidate> {
61        if self.tensor_cache[&j].size() > self.tensor_cache[&i].size() {
62            (i, j) = (j, i);
63        }
64
65        let &mut (k12, flops_12, size_12) = self.result_cache.entry((i, j)).or_insert_with(|| {
66            let k12 = self.tensor_cache.len();
67            let flops_12 = contract_op_cost_tensors(&self.tensor_cache[&i], &self.tensor_cache[&j]);
68            let size_12 = contract_size_tensors(&self.tensor_cache[&i], &self.tensor_cache[&j]);
69            let k12_tensor = &self.tensor_cache[&i] ^ &self.tensor_cache[&j];
70            self.tensor_cache.insert_new(k12, k12_tensor);
71            (k12, flops_12, size_12)
72        });
73
74        let current_flops = if let Some(total_flops) = self.comm_cache.get(&k12) {
75            *total_flops
76        } else {
77            let total_flops = flops_12 + self.comm_cache[&i].max(self.comm_cache[&j]);
78            self.comm_cache.insert(k12, total_flops);
79            total_flops
80        };
81        let current_size = size.max(size_12);
82
83        if current_flops > self.best_flops && current_size > self.best_size {
84            return None;
85        }
86        let best_flops = *self
87            .best_progress
88            .entry(remaining_len)
89            .or_insert(current_flops);
90
91        if current_flops < best_flops {
92            self.best_progress.insert(remaining_len, current_flops);
93        } else if current_flops > (self.cutoff_flops_factor * best_flops + self.largest_latency) {
94            return None;
95        }
96
97        Some(Candidate {
98            flop_cost: current_flops,
99            size_cost: current_size,
100            parent_ids: (i, j),
101            child_id: k12,
102        })
103    }
104
105    /// Explores possible pair contractions in a depth-first
106    /// recursive manner like the `optimal` approach, but with extra heuristic early pruning of branches
107    /// as well sieving by `memory_limit` and the best path found so far. A rust implementation of
108    /// the Python based `opt_einsum` implementation. Found at <https://github.com/dgasmith/opt_einsum>.
109    fn branch_iterate(
110        &mut self,
111        tensor: &CompositeTensor,
112        path: &[(usize, usize, usize)],
113        remaining: &[usize],
114        flops: f64,
115        size: f64,
116    ) {
117        if remaining.len() == 1 {
118            match self.minimize {
119                CostType::Flops => {
120                    if self.best_flops > flops {
121                        self.best_flops = flops;
122                        self.best_size = size;
123                        self.best_path = ssa_ordering(path, tensor.tensors().len());
124                    }
125                }
126                CostType::Size => {
127                    if self.best_size > size {
128                        self.best_flops = flops;
129                        self.best_size = size;
130                        self.best_path = ssa_ordering(path, tensor.tensors().len());
131                    }
132                }
133            }
134            return;
135        }
136
137        let mut candidates = BinaryHeap::with_capacity(remaining.len() * (remaining.len() - 1) / 2);
138        for pair in remaining.iter().copied().combinations(2) {
139            let candidate = self.assess_candidate(pair[0], pair[1], size, remaining.len());
140            if let Some(new_candidate) = candidate {
141                candidates.push(new_candidate);
142            }
143        }
144        let mut candidates = candidates.into_sorted_vec();
145        if let Some(limit) = self.nbranch {
146            candidates.truncate(limit);
147        }
148
149        let mut new_path = Vec::with_capacity(path.len() + 1);
150        new_path.extend_from_slice(path);
151
152        for candidate in candidates.into_iter().rev() {
153            let Candidate {
154                flop_cost,
155                size_cost,
156                parent_ids,
157                child_id,
158            } = candidate;
159            let mut new_remaining = remaining.to_vec();
160            new_remaining.retain(|e| *e != parent_ids.0 && *e != parent_ids.1);
161            new_remaining.push(child_id);
162            new_path.push((parent_ids.0, parent_ids.1, child_id));
163            self.branch_iterate(tensor, &new_path, &new_remaining, flop_cost, size_cost);
164            new_path.pop();
165        }
166    }
167}
168
169impl Pathfinder for WeightedBranchBound {
170    type Result = BasicContractionPathResult;
171
172    fn find_path(&mut self, tensor: &CompositeTensor) -> BasicContractionPathResult {
173        let tensors = tensor.tensors().clone();
174        self.result_cache.clear();
175        self.tensor_cache.clear();
176        self.largest_latency = *self
177            .comm_cache
178            .iter()
179            .max_by(|a, b| a.1.partial_cmp(b.1).expect("Tried to compare NaN"))
180            .unwrap()
181            .1;
182        let mut nested_paths = FxHashMap::default();
183        // Get the initial space requirements for uncontracted tensors
184        for (index, mut tensor) in tensors.into_iter().enumerate() {
185            // Check that tensor has sub-tensors and doesn't have external legs set
186            if let Some(composite) = tensor.as_composite() {
187                let mut bb = WeightedBranchBound::new(
188                    self.nbranch,
189                    self.cutoff_flops_factor,
190                    self.comm_cache.clone(),
191                    self.minimize,
192                );
193                let result = bb.find_path(composite);
194                nested_paths.insert(index, result.ssa_path().clone());
195                tensor = composite.external_tensor().into();
196            }
197            let tensor = tensor.into_leaf().unwrap();
198            self.tensor_cache.insert_new(index, tensor);
199        }
200        let remaining = (0..tensor.tensors().len()).collect_vec();
201        self.branch_iterate(tensor, &[], &remaining, 0f64, 0f64);
202        let best_path = ContractionPath {
203            nested: nested_paths,
204            toplevel: std::mem::take(&mut self.best_path).into_simple(),
205        };
206        BasicContractionPathResult {
207            ssa_path: best_path,
208            flops: self.best_flops,
209            size: self.best_size,
210        }
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    use rustc_hash::FxHashMap;
219
220    use crate::contractionpath::paths::CostType;
221    use crate::contractionpath::paths::Pathfinder;
222    use crate::path;
223
224    fn setup_simple() -> (CompositeTensor, FxHashMap<usize, f64>) {
225        let bond_dims =
226            FxHashMap::from_iter([(0, 5), (1, 2), (2, 6), (3, 8), (4, 1), (5, 3), (6, 4)]);
227        (
228            CompositeTensor::new(vec![
229                LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
230                LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
231                LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
232            ]),
233            FxHashMap::from_iter([(0, 20.), (1, 40.), (2, 85.)]),
234        )
235    }
236
237    fn setup_complex() -> (CompositeTensor, FxHashMap<usize, f64>) {
238        let bond_dims = FxHashMap::from_iter([
239            (0, 27),
240            (1, 18),
241            (2, 12),
242            (3, 15),
243            (4, 5),
244            (5, 3),
245            (6, 18),
246            (7, 22),
247            (8, 45),
248            (9, 65),
249            (10, 5),
250            (11, 17),
251        ]);
252        (
253            CompositeTensor::new(vec![
254                LeafTensor::new_from_map(vec![4, 3, 2], &bond_dims),
255                LeafTensor::new_from_map(vec![0, 1, 3, 2], &bond_dims),
256                LeafTensor::new_from_map(vec![4, 5, 6], &bond_dims),
257                LeafTensor::new_from_map(vec![6, 8, 9], &bond_dims),
258                LeafTensor::new_from_map(vec![10, 8, 9], &bond_dims),
259                LeafTensor::new_from_map(vec![5, 1, 0], &bond_dims),
260            ]),
261            FxHashMap::from_iter([(0, 120.), (1, 0.), (2, 15.), (3, 15.), (4, 85.), (5, 15.)]),
262        )
263    }
264
265    #[test]
266    fn test_contract_order_simple() {
267        let (tn, latency_costs) = setup_simple();
268        let mut opt = WeightedBranchBound::new(None, 20., latency_costs, CostType::Flops);
269        let result = opt.find_path(&tn);
270
271        assert_eq!(
272            result,
273            BasicContractionPathResult {
274                ssa_path: path![(1, 0), (2, 3)],
275                flops: 640.,
276                size: 538.
277            }
278        );
279    }
280
281    #[test]
282    fn test_contract_order_complex() {
283        let (tn, latency_costs) = setup_complex();
284        let mut opt = WeightedBranchBound::new(None, 20., latency_costs, CostType::Flops);
285        let result = opt.find_path(&tn);
286
287        assert_eq!(
288            result,
289            BasicContractionPathResult {
290                ssa_path: path![(3, 4), (2, 6), (1, 5), (0, 8), (7, 9)],
291                flops: 265230.,
292                size: 89478.
293            }
294        );
295    }
296}
tnc/contractionpath/paths/weighted_branchbound.rs

tnc/contractionpath/paths/
weighted_branchbound.rs
