tenferro_einsum/
optimize.rs

1use std::hash::Hasher;
2
3use omeco::ScoreFunction;
4
5use crate::{
6    ContractionOptimizerOptions, ContractionTree, Error, NestedEinsum, Result, Subscripts,
7};
8
9/// Controls how the contraction path is determined for N-ary einsum.
10///
11/// The traced API resolves this enum into a shape-independent plan
12/// specification and stores that specification in the einsum extension
13/// payload. Concrete traced inputs can resolve a [`ContractionTree`] while the
14/// op is built; symbolic traced inputs carry the plan specification until the
15/// extension runtime sees concrete execution shapes.
16///
17/// Planner options and explicit paths are part of the extension payload
18/// identity. Two otherwise identical traced einsum ops with different
19/// optimizer options, explicit paths, or fixed plan identities are not treated
20/// as the same extension op, and their compile/runtime plan cache entries are
21/// kept separate.
22///
23/// # Examples
24///
25/// ```
26/// use tenferro_einsum::{EinsumOptimize, TraceContextEinsumExt};
27/// use tenferro_ops::dim_expr::DimExpr;
28/// use tenferro_runtime::program::ProgramInputSpec;
29/// use tenferro_runtime::{DType, TraceContext};
30///
31/// let mut trace = TraceContext::new();
32/// let lhs = trace.input(ProgramInputSpec::new(
33///     DType::F64,
34///     DimExpr::from_concrete(&[2, 3]),
35/// )).unwrap();
36/// let rhs = trace.input(ProgramInputSpec::new(
37///     DType::F64,
38///     DimExpr::from_concrete(&[3, 2]),
39/// )).unwrap();
40/// let out = trace.einsum_with(
41///     &[lhs, rhs],
42///     "ij,jk->ik",
43///     EinsumOptimize::False,
44/// )
45/// .unwrap();
46///
47/// let graph = trace.finish(&[out]).unwrap();
48/// assert_eq!(
49///     graph.program().value_metadata(graph.program().outputs()[0]).unwrap().shape().len(),
50///     2,
51/// );
52/// ```
53#[derive(Debug)]
54pub enum EinsumOptimize {
55    /// Automatic optimization via omeco TreeSA.
56    Auto(ContractionOptimizerOptions),
57    /// No optimization: contract operands left-to-right.
58    False,
59    /// Parenthesized notation specifying contraction order.
60    Nested(NestedEinsum),
61    /// JAX-compatible position-based contraction path.
62    ///
63    /// Each pair references positions in a shrinking operand list. After each
64    /// contraction, the two operands are removed and the result is appended.
65    /// Because this representation is independent of concrete dimension
66    /// values, it can be used with symbolic traced inputs.
67    Path(Vec<(usize, usize)>),
68    /// Pre-computed contraction tree.
69    ///
70    /// A tree contains concrete shape-dependent planning results. It is
71    /// accepted when shapes are concrete, then converted into fixed contraction
72    /// pairs for the extension payload. A binary tree with the single pair
73    /// `(0, 1)` or `(1, 0)` may bypass the extension path and lower directly to
74    /// `dot_general`, including with symbolic traced inputs. Use
75    /// [`EinsumOptimize::Path`] instead when building a symbolic traced graph
76    /// for N-ary contraction.
77    Tree(ContractionTree),
78}
79
80impl Default for EinsumOptimize {
81    /// Default: time-optimized automatic planning.
82    fn default() -> Self {
83        Self::Auto(default_auto_options())
84    }
85}
86
87#[derive(Clone, Debug)]
88pub(crate) enum EinsumPlanSpec {
89    Auto(ContractionOptimizerOptions),
90    LeftToRight,
91    Path(Vec<(usize, usize)>),
92    FixedPairs(Vec<(usize, usize)>),
93}
94
95/// Return the default automatic optimizer options.
96#[must_use]
97pub(crate) fn default_auto_options() -> ContractionOptimizerOptions {
98    ContractionOptimizerOptions {
99        score: ScoreFunction::time_optimized(),
100        ..Default::default()
101    }
102}
103
104pub(crate) fn plan_spec_from_optimize(
105    optimize: EinsumOptimize,
106    subscripts: &Subscripts,
107) -> Result<EinsumPlanSpec> {
108    match optimize {
109        EinsumOptimize::Auto(options) => {
110            options.validate()?;
111            Ok(EinsumPlanSpec::Auto(options))
112        }
113        EinsumOptimize::False => Ok(EinsumPlanSpec::LeftToRight),
114        EinsumOptimize::Nested(nested) => {
115            let pairs = nested_to_v1_pairs(&nested, subscripts.inputs.len())?;
116            validate_fixed_pairs(&pairs, subscripts.inputs.len())?;
117            Ok(EinsumPlanSpec::FixedPairs(pairs))
118        }
119        EinsumOptimize::Path(path) => {
120            let _ = jax_path_to_v1_pairs(&path, subscripts.inputs.len())?;
121            Ok(EinsumPlanSpec::Path(path))
122        }
123        EinsumOptimize::Tree(_) => Err(Error::planning(
124            "precomputed contraction tree requires concrete input shapes; use Path or parenthesized notation for symbolic traced einsum",
125        )),
126    }
127}
128
129pub(crate) fn resolve_einsum_strategy_with_spec(
130    optimize: EinsumOptimize,
131    subscripts: &Subscripts,
132    shapes: &[&[usize]],
133) -> Result<(EinsumPlanSpec, ContractionTree)> {
134    match optimize {
135        EinsumOptimize::Tree(tree) => {
136            let pairs = tree_pairs(&tree);
137            let spec = EinsumPlanSpec::FixedPairs(pairs);
138            let tree = resolve_plan_spec(&spec, subscripts, shapes)?;
139            Ok((spec, tree))
140        }
141        optimize => {
142            let spec = plan_spec_from_optimize(optimize, subscripts)?;
143            let tree = resolve_plan_spec(&spec, subscripts, shapes)?;
144            Ok((spec, tree))
145        }
146    }
147}
148
149pub(crate) fn resolve_plan_spec(
150    spec: &EinsumPlanSpec,
151    subscripts: &Subscripts,
152    shapes: &[&[usize]],
153) -> Result<ContractionTree> {
154    match spec {
155        EinsumPlanSpec::Auto(options) => {
156            ContractionTree::optimize_with_options(subscripts, shapes, options)
157        }
158        EinsumPlanSpec::LeftToRight => {
159            let n = subscripts.inputs.len();
160            if n <= 1 {
161                ContractionTree::from_pairs(subscripts, shapes, &[])
162            } else {
163                let path: Vec<(usize, usize)> = (0..n - 1).map(|_| (0, 1)).collect();
164                let pairs = jax_path_to_v1_pairs(&path, n)?;
165                ContractionTree::from_pairs(subscripts, shapes, &pairs)
166            }
167        }
168        EinsumPlanSpec::Path(path) => {
169            let pairs = jax_path_to_v1_pairs(path, subscripts.inputs.len())?;
170            ContractionTree::from_pairs(subscripts, shapes, &pairs)
171        }
172        EinsumPlanSpec::FixedPairs(pairs) => ContractionTree::from_pairs(subscripts, shapes, pairs),
173    }
174}
175
176pub(crate) fn hash_einsum_plan_spec(spec: &EinsumPlanSpec, state: &mut dyn Hasher) {
177    match spec {
178        EinsumPlanSpec::Auto(options) => {
179            state.write_u8(0);
180            hash_optimizer_options(options, state);
181        }
182        EinsumPlanSpec::LeftToRight => state.write_u8(1),
183        EinsumPlanSpec::Path(path) => {
184            state.write_u8(2);
185            hash_pairs(path, state);
186        }
187        EinsumPlanSpec::FixedPairs(pairs) => {
188            state.write_u8(3);
189            hash_pairs(pairs, state);
190        }
191    }
192}
193
194pub(crate) fn plan_specs_equal(lhs: &EinsumPlanSpec, rhs: &EinsumPlanSpec) -> bool {
195    match (lhs, rhs) {
196        (EinsumPlanSpec::Auto(lhs), EinsumPlanSpec::Auto(rhs)) => {
197            optimizer_options_equal_by_bits(lhs, rhs)
198        }
199        (EinsumPlanSpec::LeftToRight, EinsumPlanSpec::LeftToRight) => true,
200        (EinsumPlanSpec::Path(lhs), EinsumPlanSpec::Path(rhs)) => lhs == rhs,
201        (EinsumPlanSpec::FixedPairs(lhs), EinsumPlanSpec::FixedPairs(rhs)) => lhs == rhs,
202        _ => false,
203    }
204}
205
206fn tree_pairs(tree: &ContractionTree) -> Vec<(usize, usize)> {
207    (0..tree.step_count())
208        .filter_map(|step| tree.step_pair(step))
209        .collect()
210}
211
212fn validate_fixed_pairs(pairs: &[(usize, usize)], input_count: usize) -> Result<()> {
213    let required_steps = input_count.saturating_sub(1);
214    if pairs.len() != required_steps {
215        return Err(Error::planning(format!(
216            "explicit contraction path for {input_count} operands must have {required_steps} steps, got {}",
217            pairs.len()
218        )));
219    }
220
221    let mut live = vec![false; input_count + pairs.len()];
222    for slot in live.iter_mut().take(input_count) {
223        *slot = true;
224    }
225
226    for (step_idx, &(left, right)) in pairs.iter().enumerate() {
227        let next_idx = input_count + step_idx;
228        if left == right {
229            return Err(Error::planning(format!(
230                "pair ({left}, {right}) must reference two distinct live operands"
231            )));
232        }
233        if left >= next_idx || right >= next_idx {
234            return Err(Error::planning(format!(
235                "pair ({left}, {right}) references non-existent operand"
236            )));
237        }
238        if !live[left] || !live[right] {
239            return Err(Error::planning(format!(
240                "pair ({left}, {right}) references an operand or intermediate that is no longer live"
241            )));
242        }
243
244        live[left] = false;
245        live[right] = false;
246        live[next_idx] = true;
247    }
248
249    let live_count = live.iter().filter(|&&is_live| is_live).count();
250    if live_count != 1 {
251        return Err(Error::planning(format!(
252            "explicit contraction path must leave exactly one live result, got {live_count}"
253        )));
254    }
255
256    Ok(())
257}
258
259fn hash_pairs(pairs: &[(usize, usize)], state: &mut dyn Hasher) {
260    state.write_usize(pairs.len());
261    for &(left, right) in pairs {
262        state.write_usize(left);
263        state.write_usize(right);
264    }
265}
266
267fn hash_optimizer_options(options: &ContractionOptimizerOptions, state: &mut dyn Hasher) {
268    state.write_usize(options.ntrials);
269    state.write_usize(options.niters);
270    state.write_usize(options.betas.len());
271    for value in &options.betas {
272        state.write_u64(value.to_bits());
273    }
274    state.write_u64(options.score.tc_weight.to_bits());
275    state.write_u64(options.score.sc_weight.to_bits());
276    state.write_u64(options.score.rw_weight.to_bits());
277    state.write_u64(options.score.sc_target.to_bits());
278}
279
280fn optimizer_options_equal_by_bits(
281    lhs: &ContractionOptimizerOptions,
282    rhs: &ContractionOptimizerOptions,
283) -> bool {
284    lhs.ntrials == rhs.ntrials
285        && lhs.niters == rhs.niters
286        && f64_slices_equal_by_bits(&lhs.betas, &rhs.betas)
287        && score_functions_equal_by_bits(&lhs.score, &rhs.score)
288}
289
290/// Convert JAX-style position-based path to fixed-ID pairs.
291///
292/// JAX format: each pair `(i, j)` refers to positions in a shrinking list.
293/// After contraction, the two operands are removed and the result is appended.
294/// Fixed-ID format keeps original operands at `0..input_count` and gives
295/// intermediate at step `k` the ID `input_count + k`.
296///
297/// # Errors
298///
299/// Returns an error if a path step references the same position twice or a
300/// position outside the current shrinking list.
301pub(crate) fn jax_path_to_v1_pairs(
302    jax_path: &[(usize, usize)],
303    input_count: usize,
304) -> Result<Vec<(usize, usize)>> {
305    let required_steps = input_count.saturating_sub(1);
306    if jax_path.len() != required_steps {
307        return Err(Error::planning(format!(
308            "explicit contraction path for {input_count} operands must have {required_steps} steps, got {}",
309            jax_path.len()
310        )));
311    }
312
313    let mut positions: Vec<usize> = (0..input_count).collect();
314    let mut v1_pairs = Vec::with_capacity(jax_path.len());
315
316    for (step, &(pos_a, pos_b)) in jax_path.iter().enumerate() {
317        if pos_a == pos_b {
318            return Err(Error::planning(format!(
319                "path step {step} references the same operand position twice: {pos_a}"
320            )));
321        }
322        let current_len = positions.len();
323        if pos_a >= current_len || pos_b >= current_len {
324            return Err(Error::planning(format!(
325                "path step {step} references operand positions ({pos_a}, {pos_b}) with only {current_len} live operands"
326            )));
327        }
328
329        let (lo, hi) = if pos_a < pos_b {
330            (pos_a, pos_b)
331        } else {
332            (pos_b, pos_a)
333        };
334        let id_a = positions[lo];
335        let id_b = positions[hi];
336        v1_pairs.push((id_a, id_b));
337
338        positions.remove(hi);
339        positions.remove(lo);
340        positions.push(input_count + step);
341    }
342
343    Ok(v1_pairs)
344}
345
346/// Convert a [`NestedEinsum`] tree into fixed-ID pairs.
347///
348/// # Errors
349///
350/// Returns an error if a leaf references an input outside `0..input_count` or if
351/// a node has no children.
352pub(crate) fn nested_to_v1_pairs(
353    nested: &NestedEinsum,
354    input_count: usize,
355) -> Result<Vec<(usize, usize)>> {
356    let mut pairs = Vec::with_capacity(input_count.saturating_sub(1));
357    let mut next_id = input_count;
358    let root_id = walk_nested(nested, input_count, &mut pairs, &mut next_id)?;
359    if input_count == 0 || root_id >= next_id {
360        return Err(Error::planning(
361            "nested einsum did not produce a valid root operand",
362        ));
363    }
364    Ok(pairs)
365}
366
367fn walk_nested(
368    nested: &NestedEinsum,
369    input_count: usize,
370    pairs: &mut Vec<(usize, usize)>,
371    next_id: &mut usize,
372) -> Result<usize> {
373    match nested {
374        NestedEinsum::Leaf(idx) => {
375            if *idx >= input_count {
376                return Err(Error::planning(format!(
377                    "nested einsum leaf {idx} is outside 0..{input_count}"
378                )));
379            }
380            Ok(*idx)
381        }
382        NestedEinsum::Node { children, .. } => {
383            let Some(first) = children.first() else {
384                return Err(Error::planning(
385                    "nested einsum node must have at least one child",
386                ));
387            };
388            let mut result_id = walk_nested(first, input_count, pairs, next_id)?;
389            for child in &children[1..] {
390                let child_id = walk_nested(child, input_count, pairs, next_id)?;
391                pairs.push((result_id, child_id));
392                result_id = *next_id;
393                *next_id += 1;
394            }
395            Ok(result_id)
396        }
397    }
398}
399
400fn f64_slices_equal_by_bits(lhs: &[f64], rhs: &[f64]) -> bool {
401    lhs.len() == rhs.len()
402        && lhs
403            .iter()
404            .zip(rhs)
405            .all(|(lhs, rhs)| lhs.to_bits() == rhs.to_bits())
406}
407
408fn score_functions_equal_by_bits(lhs: &ScoreFunction, rhs: &ScoreFunction) -> bool {
409    lhs.tc_weight.to_bits() == rhs.tc_weight.to_bits()
410        && lhs.sc_weight.to_bits() == rhs.sc_weight.to_bits()
411        && lhs.rw_weight.to_bits() == rhs.rw_weight.to_bits()
412        && lhs.sc_target.to_bits() == rhs.sc_target.to_bits()
413}
414
415#[cfg(test)]
416mod tests;
tenferro_einsum/optimize.rs

tenferro_einsum/
optimize.rs
