Skip to main content

burn_std/config/
fusion.rs

1use cubecl_common::config::logger::{LogLevel, LoggerConfig};
2
3/// Configuration for operation fusion in Burn.
4#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
5pub struct FusionConfig {
6    /// Logger configuration for fusion logs.
7    #[serde(default)]
8    pub logger: LoggerConfig<FusionLogLevel>,
9
10    /// Beam search configuration used when exploring fusion opportunities.
11    #[serde(default)]
12    pub beam_search: BeamSearchConfig,
13
14    /// Maximum number of operations in a single client-cached graph (router graph caching).
15    ///
16    /// The greedy graph fuser keeps accumulating ops into one cached graph; once it reaches this
17    /// many ops the graph is closed and dispatched. `None` (the default) sets no size cap, leaving
18    /// [`growth_patience`](Self::growth_patience) to decide when a graph stops being worth growing;
19    /// set a value to also bound the per-graph memory and the cost of building and replaying any
20    /// single graph.
21    #[serde(default)]
22    pub max_graph_size: Option<usize>,
23
24    /// Close the current graph once its fusion score hasn't reached a new maximum for this many
25    /// consecutive ops (router graph caching).
26    ///
27    /// The fuser scores the accumulated ops as they grow and tracks the best score so far; once
28    /// this many ops have been added without beating it, the graph has stopped getting more worth
29    /// caching and is closed. Higher values keep growing the graph longer in search of a better
30    /// score.
31    #[serde(default = "default_growth_patience")]
32    pub growth_patience: usize,
33}
34
35impl Default for FusionConfig {
36    fn default() -> Self {
37        Self {
38            logger: LoggerConfig::default(),
39            beam_search: BeamSearchConfig::default(),
40            max_graph_size: None,
41            growth_patience: default_growth_patience(),
42        }
43    }
44}
45
46fn default_growth_patience() -> usize {
47    32
48}
49
50/// Beam search configuration controlling how the fusion optimizer explores independent blocks
51/// of operations.
52#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
53pub struct BeamSearchConfig {
54    /// Maximum number of independent blocks explored during the fusion search.
55    ///
56    /// Higher values can find better fusion opportunities at the cost of more cache misses
57    /// in the fusion cache.
58    #[serde(default = "default_max_blocks")]
59    pub max_blocks: usize,
60
61    /// Stop exploring new optimizations after this many explorations (per stream).
62    ///
63    /// Each cache miss runs the (relatively expensive) optimizer to build a new optimization. A
64    /// graph whose relative form changes every step never caches, so it re-explores forever. Once
65    /// this cap is reached, cache-missing segments execute *unfused* (no optimizer work, nothing
66    /// added to the cache) instead. Cache *hits* are unaffected, so already-cached stable graphs
67    /// keep replaying. `None` (the default) never disables exploration.
68    #[serde(default)]
69    pub max_explorations: Option<usize>,
70}
71
72impl Default for BeamSearchConfig {
73    fn default() -> Self {
74        Self {
75            max_blocks: default_max_blocks(),
76            max_explorations: None,
77        }
78    }
79}
80
81fn default_max_blocks() -> usize {
82    5
83}
84
85/// Log levels for fusion logging.
86#[derive(
87    Default,
88    Clone,
89    Copy,
90    Debug,
91    PartialEq,
92    Eq,
93    PartialOrd,
94    Ord,
95    serde::Serialize,
96    serde::Deserialize,
97)]
98pub enum FusionLogLevel {
99    /// Fusion logging is disabled.
100    #[default]
101    #[serde(rename = "disabled")]
102    Disabled,
103
104    /// Log the final execution strategy selected per stream (single vs composed).
105    #[serde(rename = "basic")]
106    Basic,
107
108    /// Log block merge/split decisions and cache hit/miss events.
109    #[serde(rename = "medium")]
110    Medium,
111
112    /// Log every registration, rejection and scoring decision.
113    #[serde(rename = "full")]
114    Full,
115}
116
117impl LogLevel for FusionLogLevel {}