cubecl_runtime/tune/
record.rs

1//! What a tune leaves in the environment beside its answer.
2//!
3//! The table stores a key's winner and every candidate's result, ranked. What
4//! it cannot say is how the tune went: which candidates ran and in what order,
5//! what each cost to compile and benchmark, whether the tune stopped early,
6//! and how long the key took from its miss to its answer. A [`TuneRecord`] is
7//! that account, written to [`cubecl_environment::records`] once per tune.
8
9use crate::tune::{
10    AutotuneKey, AutotuneLogContext, AutotuneLogEvent, PersistentCacheKey, TuneCache,
11};
12use alloc::string::String;
13use alloc::vec::Vec;
14use core::time::Duration;
15use cubecl_environment::records::{Record, RecordEffect, Span};
16use serde::{Deserialize, Serialize};
17
18/// How one autotune key was decided.
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20pub struct TuneRecord<K> {
21    /// The namespace of the table the answer is stored in.
22    pub table: String,
23    /// The table's entry the answer is stored under: the key, and the
24    /// checksum of the candidate list it was tuned under.
25    pub entry: PersistentCacheKey<K>,
26    /// The index of the candidate the key runs.
27    pub winner: usize,
28    /// The candidates that ran, in the order they ran.
29    pub trials: Vec<Trial>,
30    /// The candidate that met the time limit and ended the tune before the
31    /// rest of the plan ran.
32    pub short_circuit: Option<String>,
33    /// From the cache miss to the answer committed.
34    pub wall: Duration,
35    /// Whether the tune ran inside a dry run, where launches compile but do
36    /// not execute.
37    pub dry_run: bool,
38    /// Whether the table took the answer. One that measured nothing, or that
39    /// was tuned with the cache disabled, answers this process alone: the
40    /// table holds another answer to the key, or none.
41    pub stored: bool,
42}
43
44/// One candidate's run within a tune.
45#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
46pub struct Trial {
47    /// The candidate's name.
48    pub name: String,
49    /// Compiling and benchmarking it: from its first launch to its samples
50    /// resolved.
51    pub wall: Duration,
52}
53
54impl<K> Record for TuneRecord<K> {
55    const KIND: &'static str = "autotune";
56}
57
58/// A tune being recorded: stamped when it began, written when it ends. Every
59/// call is a no-op when the environment records nothing.
60#[derive(Debug)]
61pub(crate) struct TuneRecording<K> {
62    open: Option<OpenRecording<K>>,
63}
64
65/// What a [`TuneRecording`] holds while the environment records.
66#[derive(Debug)]
67struct OpenRecording<K> {
68    span: Span,
69    table: String,
70    entry: PersistentCacheKey<K>,
71    dry_run: bool,
72}
73
74/// What [`TuneRecording::finish`] needs to know of the answer.
75pub(crate) struct Answer {
76    pub winner: usize,
77    /// Whether the table took it: see [`TuneRecord::stored`].
78    pub stored: bool,
79}
80
81impl<K: AutotuneKey> TuneRecording<K> {
82    /// Begin recording the tune of `key` in `cache`'s table. The key is
83    /// cloned only when the environment records.
84    pub(crate) fn new(cache: &TuneCache<K>, key: &K, checksum: &str) -> Self {
85        let open = Span::new().map(|span| OpenRecording {
86            span,
87            table: cache.table().into(),
88            entry: PersistentCacheKey {
89                key: key.clone(),
90                checksum: checksum.into(),
91            },
92            dry_run: crate::dry_run::dry_run(),
93        });
94        Self { open }
95    }
96
97    /// Whether the tune is recorded, and so has to track its steps: the
98    /// record's trials are the ones the log context collects.
99    pub(crate) fn is_open(&self) -> bool {
100        self.open.is_some()
101    }
102
103    /// Write the record of the tune that just answered.
104    pub(crate) fn finish(self, answer: Answer, log_context: Option<&AutotuneLogContext>) {
105        let Some(open) = self.open else {
106            return;
107        };
108        // A tune the environment switched away from went with its session.
109        let Some(wall) = open.span.elapsed() else {
110            return;
111        };
112        let mut trials = Vec::new();
113        let mut short_circuit = None;
114        for event in log_context.iter().flat_map(|context| &context.events) {
115            match event {
116                AutotuneLogEvent::TuningStep(name, wall) => trials.push(Trial {
117                    name: name.clone(),
118                    wall: *wall,
119                }),
120                AutotuneLogEvent::ShortCircuit(name) => short_circuit = Some(name.clone()),
121            }
122        }
123        let record = TuneRecord {
124            table: open.table,
125            entry: open.entry,
126            winner: answer.winner,
127            trials,
128            short_circuit,
129            wall,
130            dry_run: open.dry_run,
131            stored: answer.stored,
132        };
133        // A stored winner is the environment changing; an answer kept in
134        // memory is not.
135        let effect = if record.stored {
136            RecordEffect::Changed
137        } else {
138            RecordEffect::Observed
139        };
140        open.span.close(effect, &record);
141    }
142}
cubecl_runtime/tune/record.rs

cubecl_runtime/tune/
record.rs
