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

tnc/contractionpath/paths/
branchbound.rs
