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

use std::time::Duration;

use super::*;

struct ProfileOverrideGuard;

impl ProfileOverrideGuard {
    fn set(profile_enabled: bool, trace_enabled: bool, print_every: Option<usize>) -> Self {
        EAGER_EINSUM_PROFILE_ENABLED_OVERRIDE.with(|state| {
            *state.borrow_mut() = Some(profile_enabled);
        });
        EAGER_EINSUM_TRACE_ENABLED_OVERRIDE.with(|state| {
            *state.borrow_mut() = Some(trace_enabled);
        });
        EAGER_EINSUM_PRINT_EVERY_OVERRIDE.with(|state| {
            *state.borrow_mut() = Some(print_every);
        });
        EAGER_EINSUM_PROFILE_STATE.with(|state| {
            state.borrow_mut().clear();
        });
        Self
    }
}

impl Drop for ProfileOverrideGuard {
    fn drop(&mut self) {
        EAGER_EINSUM_PROFILE_ENABLED_OVERRIDE.with(|state| {
            *state.borrow_mut() = None;
        });
        EAGER_EINSUM_TRACE_ENABLED_OVERRIDE.with(|state| {
            *state.borrow_mut() = None;
        });
        EAGER_EINSUM_PRINT_EVERY_OVERRIDE.with(|state| {
            *state.borrow_mut() = None;
        });
        EAGER_EINSUM_PROFILE_STATE.with(|state| {
            state.borrow_mut().clear();
        });
    }
}

#[test]
fn disabled_profile_bypasses_recording() {
    let _guard = ProfileOverrideGuard::set(false, false, Some(1));

    record_eager_einsum_profile("total", Duration::from_micros(1));
    let value = profile_eager_einsum_section("phase", || 17);
    maybe_print_eager_einsum_profile();

    assert_eq!(value, 17);
    assert!(!eager_einsum_profile_enabled());
    assert!(!eager_einsum_trace_enabled());
    EAGER_EINSUM_PROFILE_STATE.with(|state| {
        assert!(state.borrow().is_empty());
    });
}

#[test]
fn enabled_profile_records_sections_and_print_reset_clears_state() {
    let _guard = ProfileOverrideGuard::set(true, true, Some(1));

    let value = profile_eager_einsum_section("phase", || 23);
    record_eager_einsum_profile("total", Duration::from_micros(2));

    assert_eq!(value, 23);
    assert!(eager_einsum_profile_enabled());
    assert!(eager_einsum_trace_enabled());
    EAGER_EINSUM_PROFILE_STATE.with(|state| {
        let state = state.borrow();
        assert_eq!(state.get("phase").map(|entry| entry.calls), Some(1));
        assert_eq!(state.get("total").map(|entry| entry.calls), Some(1));
    });

    maybe_print_eager_einsum_profile();

    EAGER_EINSUM_PROFILE_STATE.with(|state| {
        assert!(state.borrow().is_empty());
    });
}

#[test]
fn print_every_zero_keeps_recorded_state() {
    let _guard = ProfileOverrideGuard::set(true, false, Some(0));

    record_eager_einsum_profile("total", Duration::from_micros(1));
    maybe_print_eager_einsum_profile();

    EAGER_EINSUM_PROFILE_STATE.with(|state| {
        assert_eq!(
            state.borrow().get("total").map(|entry| entry.calls),
            Some(1)
        );
    });
}