tenferro-einsum 0.3.0

Subscripts, contraction planning, concrete/traced/eager einsum APIs, extension runtime, and AD rule for tenferro.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

#![cfg(feature = "autodiff")]

use tenferro_einsum::{EinsumOptimize, TraceContextEinsumExt};
use tenferro_ops::dim_expr::DimExpr;
use tenferro_runtime::program::{ProgramInputSpec, SemanticFingerprint};
use tenferro_runtime::{ExtensionCacheSelector, GraphCompiler, TraceContext};
use tenferro_tensor::{DType, Tensor};

use super::support;

fn matrix() -> ProgramInputSpec {
    ProgramInputSpec::new(DType::F64, [DimExpr::Const(2), DimExpr::Const(2)])
}

#[test]
fn trace_context_parse_cache_reuses_exact_notation() {
    let mut trace = TraceContext::new();
    let lhs = trace.input(matrix()).unwrap();
    let rhs = trace.input(matrix()).unwrap();

    trace.einsum(&[lhs, rhs], "ij,jk->ik").unwrap();
    let first = trace
        .extension_caches_mut()
        .stats(ExtensionCacheSelector::All);
    trace.einsum(&[lhs, rhs], "ij,jk->ik").unwrap();
    let second = trace
        .extension_caches_mut()
        .stats(ExtensionCacheSelector::All);

    assert_eq!(first.entries, second.entries);
    assert!(first.entries >= 1);
}

#[test]
fn plan_policy_participates_in_semantic_fingerprint() {
    fn fingerprint(optimize: EinsumOptimize) -> SemanticFingerprint {
        let mut trace = TraceContext::new();
        let lhs = trace.input(matrix()).unwrap();
        let rhs = trace.input(matrix()).unwrap();
        let output = trace
            .einsum_with(&[lhs, rhs], "ij,jk->ik", optimize)
            .unwrap();
        trace
            .finish(&[output])
            .unwrap()
            .program()
            .semantic_fingerprint()
    }

    assert_ne!(
        fingerprint(EinsumOptimize::False),
        fingerprint(EinsumOptimize::Path(vec![(0, 1)]))
    );
}

#[test]
fn compiler_does_not_own_semantic_runtime_staging_cache() {
    let mut trace = TraceContext::new();
    let lhs = trace.input(matrix()).unwrap();
    let rhs = trace.input(matrix()).unwrap();
    let output = trace.einsum(&[lhs, rhs], "ij,jk->ik").unwrap();
    let graph = trace.finish(&[output]).unwrap();
    let mut compiler = GraphCompiler::new();

    compiler.compile_traced_graph(&graph).unwrap();
    let after_first = compiler.cache_stats().entries;
    compiler.compile_traced_graph(&graph).unwrap();

    assert_eq!(compiler.cache_stats().entries, after_first);
    assert_eq!(after_first, 0);
}

#[test]
fn runtime_plan_cache_reuses_identical_shapes() {
    let mut trace = TraceContext::new();
    let lhs = trace.input(matrix()).unwrap();
    let rhs = trace.input(matrix()).unwrap();
    let output = trace.einsum(&[lhs, rhs], "ij,jk->ik").unwrap();
    let graph = trace.finish(&[output]).unwrap();
    let compiled = GraphCompiler::new().compile_traced_graph(&graph).unwrap();
    let lhs = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 3.0, 2.0, 4.0]).unwrap();
    let rhs = Tensor::from_vec_col_major(vec![2, 2], vec![5.0_f64, 7.0, 6.0, 8.0]).unwrap();
    let backend = tenferro_cpu::CpuBackend::new();
    let runtime = support::cpu_runtime_with_einsum(&backend).unwrap();

    runtime.run_compiled(&compiled, &[&lhs, &rhs]).unwrap();
    let after_first = runtime.cache_stats().unwrap().extensions.entries;
    runtime.run_compiled(&compiled, &[&lhs, &rhs]).unwrap();

    assert_eq!(
        runtime.cache_stats().unwrap().extensions.entries,
        after_first
    );
    assert!(after_first >= 1);
}