burn_cubecl_fusion/optim/elemwise/
fuser.rs

1use super::optimization::ElemwiseOptimization;
2use crate::{
3    engine::{
4        fuser::TraceOperationFuser,
5        settings::{FuseSettings, RefLayoutSetting, VectorizationSetting},
6    },
7    optim::CubeOptimization,
8};
9use burn_fusion::OperationFuser;
10use burn_std::Shape;
11use cubecl::Runtime;
12
13/// Fuses element wise operations.
14pub struct ElementWiseFuser<R: Runtime> {
15    fuser: TraceOperationFuser,
16    device: R::Device,
17}
18
19impl<R: Runtime> Clone for ElementWiseFuser<R> {
20    fn clone(&self) -> Self {
21        Self {
22            fuser: self.fuser.clone(),
23            device: self.device.clone(),
24        }
25    }
26}
27
28impl<R: Runtime> ElementWiseFuser<R> {
29    pub fn shape_id(&self) -> Shape {
30        self.fuser.current_output_shape.clone()
31    }
32    pub fn new(device: R::Device) -> Self {
33        let client = R::client(&device);
34        let props = client.properties();
35        let max_bindings = props.hardware.max_bindings;
36
37        Self {
38            fuser: TraceOperationFuser::new(
39                max_bindings,
40                FuseSettings {
41                    broadcast: true,
42                    output_shape_updates: true,
43                    inplace: true,
44                    vectorization: VectorizationSetting::Activated,
45                    ref_layout: RefLayoutSetting::Any,
46                },
47            ),
48            device,
49        }
50    }
51}
52
53impl<R: Runtime> OperationFuser<CubeOptimization<R>> for ElementWiseFuser<R> {
54    fn fuse(&mut self, operation: &burn_ir::OperationIr) {
55        self.fuser.fuse(operation);
56    }
57
58    fn finish(&mut self) -> CubeOptimization<R> {
59        let client = R::client(&self.device);
60        let trace = self.fuser.finish();
61        let elementwise = ElemwiseOptimization::new(trace, client, self.device.clone(), self.len());
62
63        CubeOptimization::ElementWise(elementwise)
64    }
65
66    fn reset(&mut self) {
67        self.fuser.reset()
68    }
69
70    fn status(&self) -> burn_fusion::FuserStatus {
71        self.fuser.status()
72    }
73
74    fn properties(&self) -> burn_fusion::FuserProperties {
75        self.fuser.properties()
76    }
77
78    fn len(&self) -> usize {
79        self.fuser.len()
80    }
81
82    fn clone_dyn(&self) -> Box<dyn OperationFuser<CubeOptimization<R>>> {
83        Box::new(self.clone())
84    }
85}
burn_cubecl_fusion/optim/elemwise/fuser.rs

burn_cubecl_fusion/optim/elemwise/
fuser.rs
