tenferro_einsum/planning/
tree.rs

1use std::collections::{HashMap, HashSet};
2use std::fmt;
3use std::mem::{size_of, size_of_val};
4
5use omeco::{
6    CodeOptimizer, EinCode as OmecoEinCode, Initializer, NestedEinsum, ScoreFunction, TreeSA,
7};
8
9use crate::cache::{saturating_sum, vec_of_vec_retained_bytes, vec_retained_bytes};
10use crate::planning::plan::{compile_step_plans, DiagPlan, GemmPlan, ReducePlan, StepPlan};
11use crate::syntax::subscripts::Subscripts;
12use crate::util::{build_size_dict, intermediate_subs};
13use crate::{Error, Result};
14
15/// A single step in the contraction sequence.
16pub(crate) struct ContractionStep {
17    pub(crate) left: usize,
18    pub(crate) right: usize,
19}
20
21/// Public options for automatic contraction-path optimization.
22///
23/// The default planner uses TreeSA with a greedy initializer and zero annealing
24/// iterations. This keeps the public API on a single optimizer family while
25/// making the default behavior effectively "greedy-only".
26#[derive(Debug, Clone)]
27pub struct ContractionOptimizerOptions {
28    /// Inverse-temperature schedule for TreeSA.
29    pub betas: Vec<f64>,
30    /// Number of independent TreeSA trials.
31    pub ntrials: usize,
32    /// Annealing iterations per temperature level.
33    pub niters: usize,
34    /// Score function used by TreeSA.
35    pub score: ScoreFunction,
36}
37
38impl Default for ContractionOptimizerOptions {
39    fn default() -> Self {
40        Self {
41            betas: Vec::new(),
42            ntrials: 1,
43            niters: 0,
44            score: ScoreFunction::default(),
45        }
46    }
47}
48
49impl ContractionOptimizerOptions {
50    fn to_treesa(&self) -> TreeSA {
51        TreeSA::new(
52            self.betas.clone(),
53            self.ntrials,
54            self.niters,
55            Initializer::Greedy,
56            self.score.clone(),
57        )
58    }
59
60    pub(crate) fn validate(&self) -> Result<()> {
61        if self.ntrials == 0 {
62            return Err(Error::planning(
63                "contraction optimizer ntrials must be at least 1",
64            ));
65        }
66        if self.betas.iter().any(|value| value.is_nan()) {
67            return Err(Error::planning(
68                "contraction optimizer betas must not contain NaN",
69            ));
70        }
71        if self.score.tc_weight.is_nan()
72            || self.score.sc_weight.is_nan()
73            || self.score.rw_weight.is_nan()
74            || self.score.sc_target.is_nan()
75        {
76            return Err(Error::planning(
77                "contraction optimizer score fields must not contain NaN",
78            ));
79        }
80        Ok(())
81    }
82}
83
84/// Contraction tree determining pairwise contraction order for N-ary einsum.
85///
86/// When contracting more than two tensors, the order in which pairwise
87/// contractions are performed significantly affects performance.
88/// `ContractionTree` encodes this order as a binary tree.
89///
90/// # Optimization
91///
92/// Use [`ContractionTree::optimize`] for automatic cost-based optimization
93/// (e.g., greedy algorithm based on tensor sizes), or
94/// [`ContractionTree::from_pairs`] for manual specification.
95pub struct ContractionTree {
96    /// Original subscripts.
97    pub(crate) subscripts: Subscripts,
98    /// Steps in the contraction (empty for single-tensor case).
99    pub(crate) steps: Vec<ContractionStep>,
100    /// Label → dimension size mapping.
101    pub(crate) size_dict: HashMap<u32, usize>,
102    /// Subscripts for each operand (0..input_count from input, then intermediates).
103    pub(crate) operand_subs: Vec<Vec<u32>>,
104    /// Pre-compiled step plans (cached to avoid recomputation per execute call).
105    pub(crate) step_plans: Vec<StepPlan>,
106}
107
108impl fmt::Debug for ContractionTree {
109    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110        f.debug_struct("ContractionTree")
111            .field("input_count", &self.subscripts.inputs.len())
112            .field("output_rank", &self.subscripts.output.len())
113            .field("steps_len", &self.steps.len())
114            .field("size_dict_len", &self.size_dict.len())
115            .field("operand_subs_len", &self.operand_subs.len())
116            .field("step_plans_len", &self.step_plans.len())
117            .finish_non_exhaustive()
118    }
119}
120
121impl ContractionTree {
122    /// Automatically compute an optimized contraction order.
123    ///
124    /// Uses a cost-based heuristic (greedy algorithm) to determine
125    /// the pairwise contraction sequence that minimizes total operation count.
126    ///
127    /// # Arguments
128    ///
129    /// * `subscripts` — Einsum subscripts for all tensors
130    /// * `shapes` — Shape of each input tensor
131    ///
132    /// # Errors
133    ///
134    /// Returns [`Error::Validation`] when subscripts and shapes have different
135    /// ranks or incompatible dimensions, or [`Error::Planning`] when no valid
136    /// contraction order can be constructed.
137    pub fn optimize(subscripts: &Subscripts, shapes: &[&[usize]]) -> Result<Self> {
138        Self::optimize_with_options(subscripts, shapes, &ContractionOptimizerOptions::default())
139    }
140
141    /// Automatically compute an optimized contraction order with explicit
142    /// planner options.
143    ///
144    /// For three or more operands, this routes planning through TreeSA using
145    /// the provided configuration. The default is greedy-initialized TreeSA
146    /// with zero annealing iterations. One or two operands need no ordering
147    /// search; their trees are built directly after validating the options.
148    ///
149    /// # Errors
150    ///
151    /// Returns [`Error::Validation`] for rank, shape, or dimension mismatches,
152    /// or [`Error::Planning`] when planner options such as `ntrials` are
153    /// invalid or no contraction order can be constructed.
154    pub fn optimize_with_options(
155        subscripts: &Subscripts,
156        shapes: &[&[usize]],
157        options: &ContractionOptimizerOptions,
158    ) -> Result<Self> {
159        options.validate()?;
160        let input_count = subscripts.inputs.len();
161        if input_count <= 1 {
162            return Self::from_pairs(subscripts, shapes, &[]);
163        }
164        if input_count == 2 {
165            return Self::from_pairs(subscripts, shapes, &[(0, 1)]);
166        }
167
168        let size_dict = build_size_dict(subscripts, shapes, None)?;
169        let pairs =
170            if let Some(omeco_pairs) = optimize_omeco_pairs(subscripts, &size_dict, options)? {
171                omeco_pairs
172            } else {
173                optimize_self_greedy_pairs(subscripts, &size_dict)?
174            };
175        Self::from_pairs(subscripts, shapes, &pairs)
176    }
177
178    /// Manually build a contraction tree from a pairwise contraction sequence.
179    ///
180    /// Each pair `(i, j)` specifies which two tensors (or intermediate results)
181    /// to contract next. Intermediate results are assigned indices starting
182    /// from the number of input tensors.
183    ///
184    /// # Arguments
185    ///
186    /// * `subscripts` — Einsum subscripts for all tensors
187    /// * `shapes` — Shape of each input tensor
188    /// * `pairs` — Ordered list of pairwise contractions
189    ///
190    /// # Examples
191    ///
192    /// ```rust
193    /// use tenferro_einsum::{ContractionTree, Subscripts};
194    ///
195    /// // Three tensors: A[ij] B[jk] C[kl] -> D[il]
196    /// // Contract B and C first, then A with the result:
197    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
198    /// let shapes = [&[3, 4][..], &[4, 5], &[5, 6]];
199    /// let tree = ContractionTree::from_pairs(
200    ///     &subs,
201    ///     &shapes,
202    ///     &[(1, 2), (0, 3)],  // B*C -> T(index=3), then A*T -> D
203    /// ).unwrap();
204    /// ```
205    ///
206    /// # Errors
207    ///
208    /// Returns [`Error::Planning`] when the pair count, operand indices, or
209    /// intermediate sequence is invalid, or [`Error::Validation`] when the
210    /// supplied shapes do not match the subscripts.
211    pub fn from_pairs(
212        subscripts: &Subscripts,
213        shapes: &[&[usize]],
214        pairs: &[(usize, usize)],
215    ) -> Result<Self> {
216        let input_count = subscripts.inputs.len();
217        let required_steps = input_count.saturating_sub(1);
218        if pairs.len() != required_steps {
219            return Err(Error::planning(format!(
220                "explicit contraction path for {input_count} operands must have {required_steps} steps, got {}",
221                pairs.len()
222            )));
223        }
224        let size_dict = build_size_dict(subscripts, shapes, None)?;
225
226        let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
227        let mut live = vec![false; input_count + pairs.len()];
228        for slot in live.iter_mut().take(input_count) {
229            *slot = true;
230        }
231        let mut steps = Vec::new();
232
233        for (step_idx, &(left, right)) in pairs.iter().enumerate() {
234            let next_idx = input_count + step_idx;
235            if left == right {
236                return Err(Error::planning(format!(
237                    "pair ({left}, {right}) must reference two distinct live operands"
238                )));
239            }
240            if left >= next_idx || right >= next_idx {
241                return Err(Error::planning(format!(
242                    "pair ({left}, {right}) references non-existent operand"
243                )));
244            }
245            if !live[left] || !live[right] {
246                return Err(Error::planning(format!(
247                    "pair ({left}, {right}) references an operand or intermediate that is no longer live"
248                )));
249            }
250
251            // Labels needed by other live operands + final output
252            let mut needed: HashSet<u32> = subscripts.output.iter().copied().collect();
253            for (idx, subs) in operand_subs.iter().enumerate() {
254                if idx != left && idx != right && live[idx] {
255                    needed.extend(subs.iter().copied());
256                }
257            }
258
259            let new_subs = intermediate_subs(&operand_subs[left], &operand_subs[right], &needed);
260            operand_subs.push(new_subs);
261            live[left] = false;
262            live[right] = false;
263            live[next_idx] = true;
264            steps.push(ContractionStep { left, right });
265        }
266
267        let live_count = live.iter().filter(|&&is_live| is_live).count();
268        if live_count != 1 {
269            return Err(Error::planning(format!(
270                "explicit contraction path must leave exactly one live result, got {live_count}"
271            )));
272        }
273
274        let mut tree = Self {
275            subscripts: subscripts.clone(),
276            steps,
277            size_dict,
278            operand_subs,
279            step_plans: Vec::new(),
280        };
281        tree.step_plans = compile_step_plans(&tree)?;
282        Ok(tree)
283    }
284
285    /// Return the number of pairwise contraction steps in this tree.
286    ///
287    /// # Examples
288    ///
289    /// ```rust
290    /// use tenferro_einsum::{ContractionTree, Subscripts};
291    ///
292    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
293    /// let tree = ContractionTree::from_pairs(
294    ///     &subs,
295    ///     &[&[2, 2], &[2, 2], &[2, 2]],
296    ///     &[(1, 2), (0, 3)],
297    /// )
298    /// .unwrap();
299    /// assert_eq!(tree.step_count(), 2);
300    /// ```
301    #[must_use]
302    pub fn step_count(&self) -> usize {
303        self.steps.len()
304    }
305
306    pub(crate) fn label_size(&self, label: u32) -> Option<usize> {
307        self.size_dict.get(&label).copied()
308    }
309
310    pub(crate) fn output_shape(&self) -> Vec<usize> {
311        self.subscripts
312            .output
313            .iter()
314            .filter_map(|label| self.label_size(*label))
315            .collect()
316    }
317
318    /// Return the operand indices for a pairwise contraction step.
319    ///
320    /// The returned indices refer to the original inputs (`0..input_count`) and
321    /// then to intermediates (`input_count..`) produced by earlier steps.
322    ///
323    /// # Examples
324    ///
325    /// ```rust
326    /// use tenferro_einsum::{ContractionTree, Subscripts};
327    ///
328    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
329    /// let tree = ContractionTree::from_pairs(
330    ///     &subs,
331    ///     &[&[2, 2], &[2, 2], &[2, 2]],
332    ///     &[(1, 2), (0, 3)],
333    /// )
334    /// .unwrap();
335    /// assert_eq!(tree.step_pair(0), Some((1, 2)));
336    /// ```
337    #[must_use]
338    pub fn step_pair(&self, step_idx: usize) -> Option<(usize, usize)> {
339        self.steps.get(step_idx).map(|step| (step.left, step.right))
340    }
341
342    /// Return the `(lhs, rhs, output)` subscripts for a pairwise step.
343    ///
344    /// The output subscripts are the intermediate labels preserved after the
345    /// contraction, or the final output labels on the last step.
346    ///
347    /// # Examples
348    ///
349    /// ```rust
350    /// use tenferro_einsum::{ContractionTree, Subscripts};
351    ///
352    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
353    /// let tree = ContractionTree::from_pairs(
354    ///     &subs,
355    ///     &[&[2, 2], &[2, 2], &[2, 2]],
356    ///     &[(1, 2), (0, 3)],
357    /// )
358    /// .unwrap();
359    /// let (lhs, rhs, out) = tree.step_subscripts(0).unwrap();
360    /// assert_eq!(lhs, &[1, 2]);
361    /// assert_eq!(rhs, &[2, 3]);
362    /// assert_eq!(out, &[1, 3]);
363    /// ```
364    #[must_use]
365    pub fn step_subscripts(&self, step_idx: usize) -> Option<(&[u32], &[u32], &[u32])> {
366        let input_count = self.subscripts.inputs.len();
367        let step = self.steps.get(step_idx)?;
368        let result_idx = input_count + step_idx;
369        let output_subs = if step_idx + 1 == self.steps.len() {
370            &self.subscripts.output
371        } else {
372            &self.operand_subs[result_idx]
373        };
374        Some((
375            &self.operand_subs[step.left],
376            &self.operand_subs[step.right],
377            output_subs,
378        ))
379    }
380
381    /// Return the precomputed lowering plan for one pairwise contraction step.
382    ///
383    /// # Examples
384    ///
385    /// ```rust
386    /// use tenferro_einsum::{ContractionTree, Subscripts};
387    ///
388    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[0, 2]);
389    /// let tree = ContractionTree::from_pairs(&subs, &[&[2, 3], &[3, 4]], &[(0, 1)]).unwrap();
390    ///
391    /// assert_eq!(tree.step_plan(0).unwrap().gemm().m(), 2);
392    /// ```
393    #[must_use]
394    pub fn step_plan(&self, step_idx: usize) -> Option<crate::lowering::PairwiseStepPlan<'_>> {
395        self.step_plans
396            .get(step_idx)
397            .map(crate::lowering::PairwiseStepPlan::new)
398    }
399
400    #[doc(hidden)]
401    #[must_use]
402    pub(crate) fn retained_bytes_for_cache_stats(&self) -> usize {
403        saturating_sum([
404            size_of::<Self>(),
405            subscripts_retained_bytes(&self.subscripts),
406            self.steps
407                .capacity()
408                .saturating_mul(size_of::<ContractionStep>()),
409            self.size_dict
410                .capacity()
411                .saturating_mul(size_of::<u32>().saturating_add(size_of::<usize>())),
412            vec_of_vec_retained_bytes(&self.operand_subs),
413            self.step_plans
414                .capacity()
415                .saturating_mul(size_of::<StepPlan>()),
416            saturating_sum(self.step_plans.iter().map(step_plan_retained_bytes)),
417        ])
418    }
419}
420
421fn subscripts_retained_bytes(subscripts: &Subscripts) -> usize {
422    saturating_sum([
423        vec_of_vec_retained_bytes(&subscripts.inputs),
424        vec_retained_bytes(&subscripts.output),
425    ])
426}
427
428fn reduce_plan_retained_bytes(plan: &ReducePlan) -> usize {
429    saturating_sum([
430        vec_retained_bytes(&plan.original_subs),
431        vec_retained_bytes(&plan.kept_subs),
432        vec_retained_bytes(&plan.out_shape),
433    ])
434}
435
436fn diag_plan_retained_bytes(plan: &DiagPlan) -> usize {
437    saturating_sum([
438        vec_retained_bytes(&plan.stages),
439        saturating_sum(plan.stages.iter().map(|stage| {
440            saturating_sum([
441                vec_retained_bytes(&stage.axis_pairs),
442                vec_retained_bytes(&stage.result_subs),
443            ])
444        })),
445        vec_retained_bytes(&plan.result_subs),
446    ])
447}
448
449fn gemm_plan_retained_bytes(plan: &GemmPlan) -> usize {
450    saturating_sum([
451        plan.reduce_a.as_ref().map_or(0, reduce_plan_retained_bytes),
452        plan.reduce_b.as_ref().map_or(0, reduce_plan_retained_bytes),
453        vec_retained_bytes(&plan.subs_a),
454        vec_retained_bytes(&plan.subs_b),
455        vec_retained_bytes(&plan.lo_modes),
456        vec_retained_bytes(&plan.ro_modes),
457        vec_retained_bytes(&plan.sum_modes),
458        vec_retained_bytes(&plan.lo_sizes),
459        vec_retained_bytes(&plan.ro_sizes),
460        vec_retained_bytes(&plan.sum_sizes),
461        vec_retained_bytes(&plan.batch_sizes),
462        vec_retained_bytes(&plan.target_a),
463        vec_retained_bytes(&plan.target_b),
464        vec_retained_bytes(&plan.c_gemm_shape),
465        vec_retained_bytes(&plan.expanded_shape),
466        vec_retained_bytes(&plan.canonical_modes),
467        vec_retained_bytes(&plan.a_gemm_shape),
468        vec_retained_bytes(&plan.b_gemm_shape),
469    ])
470}
471
472fn step_plan_retained_bytes(plan: &StepPlan) -> usize {
473    saturating_sum([
474        plan.diag_a.as_ref().map_or(0, diag_plan_retained_bytes),
475        plan.diag_b.as_ref().map_or(0, diag_plan_retained_bytes),
476        plan.strict_binary.as_ref().map_or(0, size_of_val),
477        gemm_plan_retained_bytes(&plan.gemm),
478    ])
479}
480
481fn optimize_omeco_pairs(
482    subscripts: &Subscripts,
483    size_dict: &HashMap<u32, usize>,
484    options: &ContractionOptimizerOptions,
485) -> Result<Option<Vec<(usize, usize)>>> {
486    #[cfg(test)]
487    tests::OMECO_CALLS.with(|count| count.set(count.get() + 1));
488    let code = OmecoEinCode::new(subscripts.inputs.clone(), subscripts.output.clone());
489    let optimizer = options.to_treesa();
490    let Some(nested) = optimizer.optimize(&code, size_dict) else {
491        return Ok(None);
492    };
493
494    let mut next_operand = subscripts.inputs.len();
495    let mut pairs = Vec::with_capacity(subscripts.inputs.len().saturating_sub(1));
496    nested_to_pairs(&nested, &mut next_operand, &mut pairs)?;
497    Ok(Some(pairs))
498}
499
500fn nested_to_pairs(
501    nested: &NestedEinsum<u32>,
502    next_operand: &mut usize,
503    pairs: &mut Vec<(usize, usize)>,
504) -> Result<usize> {
505    match nested {
506        NestedEinsum::Leaf { tensor_index } => Ok(*tensor_index),
507        NestedEinsum::Node { args, .. } => {
508            if args.len() != 2 {
509                return Err(Error::planning(format!(
510                    "omeco returned non-binary contraction node with {} children",
511                    args.len()
512                )));
513            }
514            let left = nested_to_pairs(&args[0], next_operand, pairs)?;
515            let right = nested_to_pairs(&args[1], next_operand, pairs)?;
516            pairs.push((left, right));
517            let result_idx = *next_operand;
518            *next_operand += 1;
519            Ok(result_idx)
520        }
521    }
522}
523
524fn build_operand_label_sets(operand_subs: &[Vec<u32>]) -> Vec<HashSet<u32>> {
525    operand_subs
526        .iter()
527        .map(|subs| subs.iter().copied().collect())
528        .collect()
529}
530
531fn build_needed_label_counts(
532    output_subs: &[u32],
533    available: &[usize],
534    operand_label_sets: &[HashSet<u32>],
535) -> HashMap<u32, usize> {
536    let mut counts = HashMap::new();
537    for &label in output_subs {
538        counts.entry(label).or_insert(1);
539    }
540    for &idx in available {
541        add_labels_to_counts(&mut counts, &operand_label_sets[idx]);
542    }
543    counts
544}
545
546fn add_labels_to_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
547    for &label in labels {
548        *counts.entry(label).or_insert(0) += 1;
549    }
550}
551
552fn remove_labels_from_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
553    for &label in labels {
554        match counts.get(&label).copied() {
555            Some(1) => {
556                counts.remove(&label);
557            }
558            Some(count) => {
559                counts.insert(label, count - 1);
560            }
561            None => {}
562        }
563    }
564}
565
566fn candidate_label_is_needed(
567    label: u32,
568    left: usize,
569    right: usize,
570    operand_label_sets: &[HashSet<u32>],
571    needed_label_counts: &HashMap<u32, usize>,
572) -> bool {
573    let mut selected_count = 0;
574    if operand_label_sets[left].contains(&label) {
575        selected_count += 1;
576    }
577    if operand_label_sets[right].contains(&label) {
578        selected_count += 1;
579    }
580    needed_label_counts.get(&label).copied().unwrap_or(0) > selected_count
581}
582
583fn collect_candidate_intermediate_subs(
584    subs_left: &[u32],
585    subs_right: &[u32],
586    left: usize,
587    right: usize,
588    operand_label_sets: &[HashSet<u32>],
589    needed_label_counts: &HashMap<u32, usize>,
590    output: &mut Vec<u32>,
591) {
592    output.clear();
593    for &label in subs_left.iter().chain(subs_right.iter()) {
594        if candidate_label_is_needed(label, left, right, operand_label_sets, needed_label_counts)
595            && !output.contains(&label)
596        {
597            output.push(label);
598        }
599    }
600}
601
602#[derive(Clone, Copy)]
603struct CandidateCostContext<'a> {
604    operand_label_sets: &'a [HashSet<u32>],
605    needed_label_counts: &'a HashMap<u32, usize>,
606    size_dict: &'a HashMap<u32, usize>,
607}
608
609fn candidate_contraction_cost(
610    subs_left: &[u32],
611    subs_right: &[u32],
612    left: usize,
613    right: usize,
614    context: CandidateCostContext<'_>,
615    candidate_subs: &mut Vec<u32>,
616) -> Result<usize> {
617    collect_candidate_intermediate_subs(
618        subs_left,
619        subs_right,
620        left,
621        right,
622        context.operand_label_sets,
623        context.needed_label_counts,
624        candidate_subs,
625    );
626    let mut cost = 1usize;
627    for &label in candidate_subs.iter() {
628        let size = context.size_dict.get(&label).copied().ok_or_else(|| {
629            Error::planning(format!(
630                "unknown size for label {label} in contraction cost"
631            ))
632        })?;
633        cost = cost.saturating_mul(size);
634    }
635    Ok(cost.max(1))
636}
637
638fn optimize_self_greedy_pairs(
639    subscripts: &Subscripts,
640    size_dict: &HashMap<u32, usize>,
641) -> Result<Vec<(usize, usize)>> {
642    #[cfg(test)]
643    tests::SELF_GREEDY_CALLS.with(|count| count.set(count.get() + 1));
644    let input_count = subscripts.inputs.len();
645    let mut available: Vec<usize> = (0..input_count).collect();
646    let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
647    let mut operand_label_sets = build_operand_label_sets(&operand_subs);
648    let mut needed_label_counts =
649        build_needed_label_counts(&subscripts.output, &available, &operand_label_sets);
650    let mut candidate_subs = Vec::new();
651    let mut pairs: Vec<(usize, usize)> = Vec::new();
652
653    while available.len() > 1 {
654        let mut best_i = 0;
655        let mut best_j = 1;
656        let mut best_cost = usize::MAX;
657
658        for i in 0..available.len() {
659            for j in (i + 1)..available.len() {
660                let li = available[i];
661                let lj = available[j];
662                let cost = candidate_contraction_cost(
663                    &operand_subs[li],
664                    &operand_subs[lj],
665                    li,
666                    lj,
667                    CandidateCostContext {
668                        operand_label_sets: &operand_label_sets,
669                        needed_label_counts: &needed_label_counts,
670                        size_dict,
671                    },
672                    &mut candidate_subs,
673                )?;
674                if cost < best_cost {
675                    best_cost = cost;
676                    best_i = i;
677                    best_j = j;
678                }
679            }
680        }
681
682        let left = available[best_i];
683        let right = available[best_j];
684        pairs.push((left, right));
685
686        let mut new_subs = Vec::new();
687        collect_candidate_intermediate_subs(
688            &operand_subs[left],
689            &operand_subs[right],
690            left,
691            right,
692            &operand_label_sets,
693            &needed_label_counts,
694            &mut new_subs,
695        );
696        let new_idx = operand_subs.len();
697        let new_label_set: HashSet<u32> = new_subs.iter().copied().collect();
698        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[left]);
699        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[right]);
700        add_labels_to_counts(&mut needed_label_counts, &new_label_set);
701        operand_subs.push(new_subs);
702        operand_label_sets.push(new_label_set);
703        available.remove(best_j);
704        available.remove(best_i);
705        available.push(new_idx);
706    }
707
708    Ok(pairs)
709}
710
711#[cfg(test)]
712mod tests;
tenferro_einsum/planning/tree.rs

tenferro_einsum/planning/
tree.rs
