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 {}