burn_cubecl_fusion/optim/reduce/
fuser.rs

1use super::{
2    ReduceSettings,
3    optimization::{FusedReduce, ReduceInstruction, ReduceOptimization},
4};
5use crate::{
6    engine::{
7        codegen::ir::FuseType,
8        fuser::TraceOperationFuser,
9        settings::{FuseSettings, RefLayoutSetting, VectorizationSetting},
10    },
11    optim::CubeOptimization,
12};
13use burn_fusion::{FuserStatus, OperationFuser};
14use burn_ir::{NumericOperationIr, OperationIr, ReduceDimOpIr};
15use burn_std::Shape;
16use cubecl::Runtime;
17
18/// Fuses element wise operations around a reduce operation.
19pub struct ReduceFuser<R: Runtime> {
20    pub(crate) fuser: TraceOperationFuser,
21    pub(crate) fuser_read_fallback: TraceOperationFuser,
22    fuser_write_fallback: TraceOperationFuser,
23    settings_write: FuseSettings,
24    pub(crate) device: R::Device,
25    pub(crate) reduce: Option<FusedReduce>,
26    settings: ReduceSettings,
27}
28
29impl<R: Runtime> Clone for ReduceFuser<R> {
30    fn clone(&self) -> Self {
31        Self {
32            fuser: self.fuser.clone(),
33            fuser_read_fallback: self.fuser_read_fallback.clone(),
34            fuser_write_fallback: self.fuser_write_fallback.clone(),
35            settings_write: self.settings_write,
36            device: self.device.clone(),
37            reduce: self.reduce.clone(),
38            settings: self.settings,
39        }
40    }
41}
42
43#[derive(Debug)]
44pub enum ReduceFuserInfo {
45    FusedReduce { shape_input_id: Shape, axis: usize },
46    FusedElemwise { shape_id: Shape },
47}
48
49impl<R: Runtime> ReduceFuser<R> {
50    pub fn new(device: R::Device, settings: ReduceSettings) -> Self {
51        let client = R::client(&device);
52        let props = client.properties();
53        let max_bindings = props.hardware.max_bindings;
54        let settings_read = FuseSettings {
55            // Inplace would work, but not when we have a concrete output to write too.
56            inplace: true,
57            ref_layout: RefLayoutSetting::OnlyContiguous,
58            broadcast: false,
59            output_shape_updates: true,
60            vectorization: VectorizationSetting::Activated,
61        };
62        let settings_write = FuseSettings {
63            inplace: false,
64            output_shape_updates: false,
65            vectorization: VectorizationSetting::SmallerOrEqualThanPreviousBlock { block_pos: 0 },
66            broadcast: false,
67            ref_layout: RefLayoutSetting::OnlyContiguous,
68        };
69        let settings_fallback = FuseSettings::default();
70
71        Self {
72            fuser: TraceOperationFuser::new(max_bindings, settings_read),
73            fuser_read_fallback: TraceOperationFuser::new(max_bindings, settings_fallback),
74            fuser_write_fallback: TraceOperationFuser::new(max_bindings, settings_fallback),
75            settings_write,
76            device,
77            reduce: None,
78            settings,
79        }
80    }
81
82    pub fn reduce_info(&self) -> ReduceFuserInfo {
83        match &self.reduce {
84            Some(reduce) => {
85                let shape_input_id = reduce.op.input.shape.clone();
86                let axis = reduce.axis;
87
88                ReduceFuserInfo::FusedReduce {
89                    shape_input_id,
90                    axis,
91                }
92            }
93            None => {
94                let shape_id = self.fuser_read_fallback.current_output_shape.clone();
95                ReduceFuserInfo::FusedElemwise { shape_id }
96            }
97        }
98    }
99
100    fn on_reduce(&mut self, op: &ReduceDimOpIr, inst: ReduceInstruction) {
101        // TODO: Fix: we need to have fuse-on-read with an identity block.
102        //
103        // if self.fuser.num_ops == 0 && false {
104        //     self.fuser.current_output_shape = op.input.shape.dims.clone();
105        // } else if self.fuser.current_output_shape != op.input.shape.dims {
106
107        if self.fuser.current_output_shape != op.input.shape {
108            self.fuser.close();
109            self.fuser_read_fallback.close();
110            return;
111        }
112
113        let [input] = self
114            .fuser
115            .next_block([&op.input], self.settings_write, false);
116
117        let output = self.fuser.output_unhandled(&op.out);
118        let axis = op.axis;
119
120        let fuse_on_write_activated = match self.settings {
121            ReduceSettings::Always => true,
122            // We only activate fuse-on-write when the reduction isn't on the last dimension, otherwise
123            // vectorization is impossible. Only [VectorizationMode::Perpendicular] supports vectorization.
124            //
125            // We could still fuse some output operations, but it would probably lead to worse performance.
126            ReduceSettings::OnlyParallel => axis != op.input.shape.rank() - 1,
127            ReduceSettings::Never => false,
128        };
129
130        if !fuse_on_write_activated {
131            self.fuser.close();
132        }
133
134        let acc = match inst {
135            ReduceInstruction::Mean | ReduceInstruction::Prod | ReduceInstruction::Sum => {
136                match input.precision() {
137                    FuseType::F16 | FuseType::BF16 => FuseType::F32,
138                    FuseType::I16 | FuseType::I8 => FuseType::I32,
139                    FuseType::U16 | FuseType::U8 => FuseType::U32,
140                    _ => input.precision(),
141                }
142            }
143            _ => input.precision(),
144        };
145
146        self.reduce = Some(FusedReduce {
147            input,
148            output,
149            acc,
150            axis,
151            op: op.clone(),
152            use_planes: false,
153            shared: false,
154            inst,
155        });
156
157        self.fuser_read_fallback.close();
158    }
159
160    fn on_elemwise_read(&mut self, operation: &OperationIr) {
161        let can_register =
162            self.fuser.can_fuse(operation) && self.fuser_read_fallback.can_fuse(operation);
163
164        match can_register {
165            true => {
166                self.fuser.fuse(operation);
167                self.fuser_read_fallback.fuse(operation);
168            }
169            false => {
170                self.fuser.close();
171                self.fuser_read_fallback.close();
172            }
173        };
174    }
175
176    fn on_elemwise_write(&mut self, operation: &OperationIr) {
177        let can_register =
178            self.fuser.can_fuse(operation) && self.fuser_write_fallback.can_fuse(operation);
179
180        match can_register {
181            true => {
182                self.fuser.fuse(operation);
183                self.fuser_write_fallback.fuse(operation);
184            }
185            false => {
186                self.fuser.close();
187                self.fuser_write_fallback.close();
188            }
189        };
190    }
191}
192
193impl<R: Runtime> OperationFuser<CubeOptimization<R>> for ReduceFuser<R> {
194    fn fuse(&mut self, operation: &OperationIr) {
195        if let FuserStatus::Closed = self.fuser.status() {
196            return;
197        }
198
199        if self.reduce.is_none() {
200            if let OperationIr::NumericFloat(_, op) = operation {
201                match op {
202                    NumericOperationIr::SumDim(op) => {
203                        self.on_reduce(op, ReduceInstruction::Sum);
204                    }
205                    NumericOperationIr::MeanDim(op) => {
206                        self.on_reduce(op, ReduceInstruction::Mean);
207                    }
208                    NumericOperationIr::ProdDim(op) => {
209                        self.on_reduce(op, ReduceInstruction::Prod);
210                    }
211                    NumericOperationIr::ArgMax(op) => {
212                        self.on_reduce(op, ReduceInstruction::ArgMax);
213                    }
214                    NumericOperationIr::ArgMin(op) => {
215                        self.on_reduce(op, ReduceInstruction::ArgMin);
216                    }
217                    NumericOperationIr::MinDim(op) => {
218                        self.on_reduce(op, ReduceInstruction::Min);
219                    }
220                    NumericOperationIr::MaxDim(op) => {
221                        self.on_reduce(op, ReduceInstruction::Max);
222                    }
223                    NumericOperationIr::MaxAbsDim(op) => {
224                        self.on_reduce(op, ReduceInstruction::MaxAbs);
225                    }
226                    _ => {
227                        self.on_elemwise_read(operation);
228                    }
229                };
230            } else if let OperationIr::NumericInt(_, op) = operation {
231                match op {
232                    NumericOperationIr::SumDim(op) => {
233                        self.on_reduce(op, ReduceInstruction::Sum);
234                    }
235                    NumericOperationIr::MeanDim(op) => {
236                        self.on_reduce(op, ReduceInstruction::Mean);
237                    }
238                    NumericOperationIr::ProdDim(op) => {
239                        self.on_reduce(op, ReduceInstruction::Prod);
240                    }
241                    NumericOperationIr::ArgMax(op) => {
242                        self.on_reduce(op, ReduceInstruction::ArgMax);
243                    }
244                    NumericOperationIr::ArgMin(op) => {
245                        self.on_reduce(op, ReduceInstruction::ArgMin);
246                    }
247                    NumericOperationIr::MinDim(op) => {
248                        self.on_reduce(op, ReduceInstruction::Min);
249                    }
250                    NumericOperationIr::MaxDim(op) => {
251                        self.on_reduce(op, ReduceInstruction::Max);
252                    }
253                    NumericOperationIr::MaxAbsDim(op) => {
254                        self.on_reduce(op, ReduceInstruction::MaxAbs);
255                    }
256                    _ => {
257                        self.on_elemwise_read(operation);
258                    }
259                };
260            } else {
261                self.on_elemwise_read(operation);
262            }
263        } else {
264            self.on_elemwise_write(operation);
265        }
266    }
267
268    fn finish(&mut self) -> CubeOptimization<R> {
269        let client = R::client(&self.device);
270        let trace = self.fuser.finish();
271        let trace_read_fallback = self.fuser_read_fallback.finish();
272        let trace_write_fallback = self.fuser_write_fallback.finish();
273        let fuse_reduce = self.reduce.as_ref().unwrap();
274
275        let reduce = ReduceOptimization::new(
276            trace,
277            trace_read_fallback,
278            trace_write_fallback,
279            client,
280            self.device.clone(),
281            self.len(),
282            self.fuser_read_fallback.len(),
283            fuse_reduce.clone(),
284            self.settings,
285        );
286
287        CubeOptimization::Reduce(reduce)
288    }
289
290    fn reset(&mut self) {
291        self.fuser.reset();
292        self.fuser_read_fallback.reset();
293        self.fuser_write_fallback.reset();
294        self.reduce = None;
295    }
296
297    fn status(&self) -> burn_fusion::FuserStatus {
298        self.fuser.status()
299    }
300
301    fn properties(&self) -> burn_fusion::FuserProperties {
302        let mut properties = self.fuser.properties();
303        properties.ready = self.reduce.is_some();
304        properties
305    }
306
307    fn len(&self) -> usize {
308        self.fuser.len() + if self.reduce.is_some() { 1 } else { 0 }
309    }
310
311    fn clone_dyn(&self) -> Box<dyn OperationFuser<CubeOptimization<R>>> {
312        Box::new(self.clone())
313    }
314}
burn_cubecl_fusion/optim/reduce/fuser.rs

burn_cubecl_fusion/optim/reduce/
fuser.rs
