opt_einsum_path/
contract.rs

1//! Contains the primary optimization and contraction routines.
2
3use crate::*;
4
5/* #region PathInfo */
6
7#[derive(Debug, Clone)]
8pub struct PathInfo {
9    pub contraction_list: Vec<ContractionType>,
10    pub input_subscripts: String,
11    pub output_subscript: String,
12    pub indices: ArrayIndexType,
13    pub path: PathType,
14    pub scale_list: Vec<usize>,
15    pub naive_cost: SizeType,
16    pub opt_cost: SizeType,
17    pub speedup: f64,
18    pub size_list: Vec<SizeType>,
19    pub size_dict: SizeDictType,
20    pub shapes: Vec<Vec<usize>>,
21    pub equation: String,
22    pub largest_intermediate: SizeType,
23}
24
25impl PathInfo {
26    pub fn new(
27        contraction_list: Vec<ContractionType>,
28        input_subscripts: String,
29        output_subscript: String,
30        indices: ArrayIndexType,
31        path: PathType,
32        scale_list: Vec<usize>,
33        naive_cost: SizeType,
34        opt_cost: SizeType,
35        size_list: Vec<SizeType>,
36        size_dict: SizeDictType,
37    ) -> Self {
38        let speedup = naive_cost.to_f64().unwrap() / opt_cost.to_f64().unwrap().max(1.0);
39        let shapes = input_subscripts.split(',').map(|ks| ks.chars().map(|k| size_dict[&k]).collect()).collect();
40        let equation = format!("{input_subscripts}->{output_subscript}");
41        let largest_intermediate = size_list.clone().into_iter().reduce(SizeType::max).unwrap_or(SizeType::one());
42
43        Self {
44            contraction_list,
45            input_subscripts,
46            output_subscript,
47            indices,
48            path,
49            scale_list,
50            naive_cost,
51            opt_cost,
52            speedup,
53            size_list,
54            size_dict,
55            shapes,
56            equation,
57            largest_intermediate,
58        }
59    }
60}
61
62impl std::fmt::Display for PathInfo {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        // Header section
65        writeln!(f, "  Complete contraction:  {}", self.equation)?;
66        writeln!(f, "         Naive scaling:  {}", self.indices.len())?;
67        writeln!(f, "     Optimized scaling:  {}", self.scale_list.iter().max().unwrap_or(&0))?;
68        writeln!(f, "      Naive FLOP count:  {:.3e}", self.naive_cost)?;
69        writeln!(f, "  Optimized FLOP count:  {:.3e}", self.opt_cost)?;
70        writeln!(f, "   Theoretical speedup:  {:.3e}", self.speedup)?;
71        writeln!(f, "  Largest intermediate:  {:.3e} elements", self.largest_intermediate)?;
72        writeln!(f, "{}", "-".repeat(80))?;
73        writeln!(f, "{:>6} {:>11} {:>22} {:>37}", "scaling", "BLAS", "current", "remaining")?;
74        writeln!(f, "{}", "-".repeat(80))?;
75
76        // Contraction steps
77        for (n, ContractionType { einsum_str, remaining, do_blas, .. }) in self.contraction_list.iter().enumerate() {
78            let remaining_str = match remaining {
79                Some(remaining) => format!("{}->{}", remaining.join(","), self.output_subscript),
80                None => "...".to_string(),
81            };
82
83            let size_remaining = 56usize.saturating_sub(22.max(einsum_str.len()));
84            let scale = self.scale_list.get(n).unwrap_or(&0);
85            let blas_str = do_blas.unwrap_or("");
86
87            writeln!(f, "{scale:>4} {blas_str:>14} {einsum_str:>22}    {remaining_str:>size_remaining$}")?;
88        }
89
90        Ok(())
91    }
92}
93
94/* #endregion */
95
96#[doc = include_str!("contract_path.md")]
97pub fn contract_path(
98    subscripts: &str,
99    operands: &[TensorShapeType],
100    mut optimize: impl PathOptimizer,
101    memory_limit: impl Into<SizeLimitType>,
102) -> Result<(PathType, PathInfo), String> {
103    // Parse einsum input
104    let (input_subscripts, output_subscript, input_shapes) = parser::parse_einsum_input(subscripts, operands)?;
105
106    // Build useful collections
107    let input_list: Vec<&str> = input_subscripts.split(',').collect_vec();
108    let input_sets: Vec<ArrayIndexType> = input_list.iter().map(|s| s.chars().collect()).collect();
109    let input_sets_ref: Vec<&ArrayIndexType> = input_sets.iter().collect();
110    let output_set: ArrayIndexType = output_subscript.chars().collect();
111    let indices: ArrayIndexType = input_subscripts.chars().filter(|&c| c != ',').collect();
112
113    // Build size dictionary
114    let mut size_dict = SizeDictType::new();
115    for (tnum, term) in input_list.iter().enumerate() {
116        let sh = &input_shapes[tnum];
117        let term_chars = term.chars().collect_vec();
118        if sh.len() != term_chars.len() {
119            return Err(format!(
120                "Einstein sum subscript '{}' does not contain the correct number of indices for operand {tnum}.",
121                input_list[tnum]
122            ));
123        }
124
125        for (cnum, &char) in term_chars.iter().enumerate() {
126            let dim = sh[cnum];
127
128            if let Some(&existing_dim) = size_dict.get(&char) {
129                // For broadcasting cases we always want the largest dim size
130                if existing_dim == 1 {
131                    size_dict.insert(char, dim);
132                } else if dim != 1 && dim != existing_dim {
133                    return Err(format!(
134                        "Size of label '{char}' for operand {tnum} ({existing_dim}) does not match previous terms ({dim})."
135                    ));
136                }
137            } else {
138                size_dict.insert(char, dim);
139            }
140        }
141    }
142
143    // Compute size of each input array plus the output array
144    let size_list: Vec<SizeType> = input_list
145        .iter()
146        .chain([output_subscript.as_str()].iter())
147        .map(|&term| helpers::compute_size_by_dict(term.chars(), &size_dict))
148        .collect();
149    let size_list_max = size_list.iter().cloned().reduce(SizeType::max).unwrap_or(SizeType::zero());
150    let memory_arg = match memory_limit.into() {
151        SizeLimitType::None => None,
152        SizeLimitType::MaxInput => Some(size_list_max),
153        SizeLimitType::Size(size) => Some(size),
154    };
155
156    let num_ops = input_list.len();
157
158    // Compute naive cost
159    let inner_product = (input_sets.iter().map(|s| s.len()).sum::<usize>() - indices.len()) > 0;
160    let naive_cost = helpers::flop_count(indices.iter(), inner_product, num_ops, &size_dict);
161
162    // Compute the path
163    let path_tuple = if num_ops <= 2 {
164        // Nothing to be optimized
165        vec![(0..num_ops).collect()]
166    } else {
167        // Use the optimizer to find the best contraction path
168        optimize.optimize_path(&input_sets_ref, &output_set, &size_dict, memory_arg)?
169    };
170
171    let mut cost_list = Vec::new();
172    let mut scale_list = Vec::new();
173    let mut size_intermediates = Vec::new();
174    let mut contraction_list = Vec::new();
175
176    let mut current_input_list = input_list.iter().map(|s| s.to_string()).collect_vec();
177    let mut current_input_sets = input_sets.clone();
178    let mut current_input_shapes = input_shapes.clone();
179
180    for (cnum, contract_inds) in path_tuple.iter().enumerate() {
181        // Make sure we remove inds from right to left
182        let mut contract_inds = contract_inds.iter().sorted().rev().cloned().collect_vec();
183
184        let (out_inds, new_input_sets, idx_removed, idx_contract) =
185            helpers::find_contraction(&contract_inds, &current_input_sets.iter().collect_vec(), &output_set);
186
187        // Compute cost, scale, and size
188        let cost = helpers::flop_count(idx_contract.iter(), !idx_removed.is_empty(), contract_inds.len(), &size_dict);
189        cost_list.push(cost);
190        scale_list.push(idx_contract.len());
191        size_intermediates.push(helpers::compute_size_by_dict(out_inds.iter(), &size_dict));
192
193        let mut tmp_inputs = contract_inds.iter().map(|&i| current_input_list.remove(i)).collect_vec();
194        let mut tmp_shapes = contract_inds.iter().map(|&i| current_input_shapes.remove(i)).collect_vec();
195
196        let do_blas = blas::can_blas(
197            &tmp_inputs.iter().map(|s| s.as_str()).collect_vec(),
198            &out_inds.iter().join(""),
199            &idx_removed,
200            Some(&tmp_shapes),
201        );
202
203        // Last contraction
204        let idx_result = if cnum == path_tuple.len() - 1 {
205            output_subscript.clone()
206        } else {
207            // use tensordot order to minimize transpositions
208            let all_input_inds = tmp_inputs.join("");
209            let mut sorted_out_inds = out_inds.iter().cloned().collect_vec();
210            sorted_out_inds.sort_by_key(|c| all_input_inds.find(*c));
211            sorted_out_inds.into_iter().collect()
212        };
213
214        let shp_result = parser::find_output_shape(
215            &tmp_inputs.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
216            &tmp_shapes,
217            &idx_result,
218        );
219
220        // Revert the order to normal
221        contract_inds.reverse();
222        tmp_inputs.reverse();
223        tmp_shapes.reverse();
224
225        current_input_list.push(idx_result.clone());
226        current_input_shapes.push(shp_result);
227        current_input_sets = new_input_sets;
228
229        let einsum_str = format!("{}->{}", tmp_inputs.join(","), idx_result);
230
231        // for large expressions saving the remaining terms at each step can
232        // incur a large memory footprint
233        let remaining = if current_input_list.len() <= 20 { Some(current_input_list.clone()) } else { None };
234
235        let contraction =
236            ContractionType { indices: contract_inds, idx_rm: idx_removed, einsum_str, remaining, do_blas };
237        contraction_list.push(contraction);
238    }
239
240    let opt_cost = cost_list.iter().sum::<SizeType>();
241    let speedup = naive_cost.to_f64().unwrap() / opt_cost.to_f64().unwrap().max(1.0);
242    let largest_intermediate = size_intermediates.iter().cloned().reduce(SizeType::max).unwrap_or(SizeType::zero());
243
244    let path_info = PathInfo {
245        contraction_list,
246        input_subscripts: input_subscripts.clone(),
247        output_subscript: output_subscript.clone(),
248        indices,
249        path: path_tuple.clone(),
250        scale_list,
251        naive_cost,
252        opt_cost,
253        speedup,
254        size_list: size_intermediates,
255        size_dict,
256        shapes: input_shapes,
257        equation: format!("{input_subscripts}->{output_subscript}"),
258        largest_intermediate,
259    };
260
261    Ok((path_tuple, path_info))
262}
263
264#[test]
265fn test_contract_path() {
266    let a_shape = vec![4, 4];
267    let b_shape = vec![4, 4];
268    let subscripts = "ij,jk->ik";
269    let (path, path_info) =
270        contract_path(subscripts, &[a_shape, b_shape], OptimizeKind::from("optimal"), None).unwrap();
271
272    assert_eq!(path.len(), 1);
273    assert_eq!(path_info.input_subscripts, "ij,jk");
274    assert_eq!(path_info.output_subscript, "ik");
275    println!("{path:?}");
276    println!("{path_info}");
277    println!("{path_info:?}");
278}
279
280#[test]
281fn test_contract_path_issue_254() {
282    let b = 64;
283    let g = 8;
284    let k = 4096;
285    let d = 128;
286
287    let a_shape = vec![b, g, k];
288    let v_shape = vec![b, k, d];
289    let s_shape = vec![b, k];
290    let subscripts = "bgk,bkd,bk->bgd";
291    let time = std::time::Instant::now();
292    let (path, path_info) =
293        contract_path(subscripts, &[a_shape, v_shape, s_shape], OptimizeKind::from("optimal"), None).unwrap();
294    println!("Time: {:?}", time.elapsed());
295
296    assert_eq!(path.len(), 2);
297    println!("{path:?}");
298    println!("{path_info}");
299}
300
301#[test]
302fn test_contract_path_issue_248() {
303    let no = 300;
304    let naux = 50;
305    let nx = 2000;
306
307    let m = vec![nx, nx];
308    let eta = vec![no, nx];
309    let eta1 = vec![naux, nx];
310    let theta = vec![naux, no];
311    let subscripts = "iP,sP,PQ,jQ,tQ,ti,sj->";
312    let time = std::time::Instant::now();
313    let (path, path_info) = contract_path(
314        subscripts,
315        &[eta.clone(), eta1.clone(), m.clone(), eta.clone(), eta1.clone(), theta.clone(), theta.clone()],
316        OptimizeKind::from("optimal"),
317        None,
318    )
319    .unwrap();
320    println!("Time: {:?}", time.elapsed());
321
322    println!("{path:?}");
323    println!("{path_info}");
324}
325
326#[test]
327fn test_greedy_issue_248() {
328    let no = 300;
329    let naux = 50;
330    let nx = 2000;
331
332    let m = vec![nx, nx];
333    let eta = vec![no, nx];
334    let eta1 = vec![naux, nx];
335    let theta = vec![naux, no];
336    let subscripts = "iP,sP,PQ,jQ,tQ,ti,sj->";
337    let shapes = vec![eta.clone(), eta1.clone(), m.clone(), eta.clone(), eta1.clone(), theta.clone(), theta.clone()];
338
339    let time = std::time::Instant::now();
340    let (path, path_info) = contract_path(subscripts, &shapes, OptimizeKind::from("optimal"), None).unwrap();
341    println!("Time (optimal): {:?}", time.elapsed());
342    println!("{path:?}");
343    println!("{path_info}");
344
345    let time = std::time::Instant::now();
346    let (path, path_info) = contract_path(subscripts, &shapes, OptimizeKind::from("greedy"), None).unwrap();
347    println!("Time (greedy): {:?}", time.elapsed());
348    println!("{path:?}");
349    println!("{path_info}");
350
351    let time = std::time::Instant::now();
352    let (path, path_info) = contract_path(subscripts, &shapes, OptimizeKind::from("dp"), None).unwrap();
353    println!("Time (dp): {:?}", time.elapsed());
354    println!("{path:?}");
355    println!("{path_info}");
356
357    let time = std::time::Instant::now();
358    let (path, path_info) = contract_path(subscripts, &shapes, OptimizeKind::from("random-greedy-128"), None).unwrap();
359    println!("Time (dp): {:?}", time.elapsed());
360    println!("{path:?}");
361    println!("{path_info}");
362}
363
364#[test]
365fn test_greedy_non_optimal() {
366    let m = vec![35, 37, 59];
367    let a = vec![35, 51, 59];
368    let b = vec![37, 51, 51, 59];
369    let c = vec![59, 27];
370    let shapes = vec![m, a, b, c];
371
372    let subscripts = "xyf,xtf,ytpf,fr->tpr";
373    let time = std::time::Instant::now();
374    let (path, path_info) = contract_path(subscripts, &shapes, OptimizeKind::from("greedy"), None).unwrap();
375    println!("Time: {:?}", time.elapsed());
376    println!("{path:?}");
377    println!("{path_info}");
378
379    let subscripts = "xyf,xtf,ytpf,fr->tpr";
380    let time = std::time::Instant::now();
381    let (path, path_info) = contract_path(subscripts, &shapes, OptimizeKind::from("optimal"), None).unwrap();
382    println!("Time: {:?}", time.elapsed());
383    println!("{path:?}");
384    println!("{path_info}");
385}
386
387#[test]
388fn test_pathtype_optimizer() {
389    let subscripts = "qgcf,sotr,klb,jlretia,hpn,nseha,jgoqm,ipkb,cdfm,d->";
390    let shapes = vec![
391        vec![5, 2, 9, 4],
392        vec![4, 9, 5, 9],
393        vec![5, 4, 2],
394        vec![5, 4, 9, 7, 5, 3, 6],
395        vec![5, 2, 8],
396        vec![8, 4, 7, 5, 6],
397        vec![5, 2, 9, 5, 8],
398        vec![3, 2, 5, 2],
399        vec![9, 3, 4, 8],
400        vec![3],
401    ];
402    let path_inp = [[0, 8], [3, 4], [1, 4], [5, 6], [1, 5], [0, 4], [0, 3], [1, 2], [0, 1]];
403    let (path, path_info) = contract_path(subscripts, &shapes, path_inp, None).unwrap();
404    assert_eq!(path, path_inp);
405    println!("{path_info}");
406
407    let path_inp = vec![[0, 8], [3, 4], [1, 4], [5, 6], [1, 5], [0, 4], [0, 3], [1, 3], [0, 1]];
408    let path_inp = path_inp.iter().map(|a| a.to_vec()).collect_vec();
409    assert!(contract_path(subscripts, &shapes, path_inp.clone(), None).is_err());
410}
opt_einsum_path/contract.rs

opt_einsum_path/
contract.rs
