burn_cubecl_fusion/optim/matmul/
fuser.rs

1use super::optimization::{FusedMatmul, MatmulOptimization};
2use crate::{
3    engine::{fuser::TraceOperationFuser, settings::FuseSettings},
4    optim::CubeOptimization,
5    optim::matmul::args::MatmulArg,
6};
7use burn_fusion::{FuserStatus, OperationFuser};
8use burn_ir::{FloatOperationIr, OperationIr};
9use burn_std::DType;
10use cubecl::Runtime;
11
12/// Fused element wise operations that are normally memory bound.
13pub struct MatmulFuser<R: Runtime> {
14    fuser: TraceOperationFuser,
15    fuser_fallback: TraceOperationFuser,
16    device: R::Device,
17    matmul: Option<FusedMatmul>,
18}
19
20impl<R: Runtime> Clone for MatmulFuser<R> {
21    fn clone(&self) -> Self {
22        Self {
23            fuser: self.fuser.clone(),
24            fuser_fallback: self.fuser_fallback.clone(),
25            device: self.device.clone(),
26            matmul: self.matmul.clone(),
27        }
28    }
29}
30
31impl<R: Runtime> MatmulFuser<R> {
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        let settings_matmul = FuseSettings {
37            output_shape_updates: false,
38            ..Default::default()
39        };
40        let settings_fallback = FuseSettings::default();
41
42        Self {
43            fuser: TraceOperationFuser::new(max_bindings, settings_matmul),
44            fuser_fallback: TraceOperationFuser::new(max_bindings, settings_fallback),
45            device,
46            matmul: None,
47        }
48    }
49}
50
51impl<R: Runtime> OperationFuser<CubeOptimization<R>> for MatmulFuser<R> {
52    fn fuse(&mut self, operation: &OperationIr) {
53        if let FuserStatus::Closed = self.fuser.status() {
54            return;
55        }
56
57        if self.matmul.is_none() {
58            if let OperationIr::Float(_, FloatOperationIr::Matmul(op)) = operation {
59                // Precision shouldn't be hardcoded but I don't know how to get float precision of the backend
60                let lhs = match op.lhs.dtype {
61                    DType::QFloat(scheme) => {
62                        let (data, scales) = self.fuser.input_quantized_unhandled(&op.lhs).unwrap();
63                        MatmulArg::Quantized {
64                            data,
65                            scales,
66                            precision: op.out.dtype.into(),
67                            scheme,
68                        }
69                    }
70                    _ => MatmulArg::Normal(self.fuser.input_unhandled(&op.lhs)),
71                };
72                let rhs = match op.rhs.dtype {
73                    DType::QFloat(scheme) => {
74                        let (data, scales) = self.fuser.input_quantized_unhandled(&op.rhs).unwrap();
75                        MatmulArg::Quantized {
76                            data,
77                            scales,
78                            precision: op.out.dtype.into(),
79                            scheme,
80                        }
81                    }
82                    _ => MatmulArg::Normal(self.fuser.input_unhandled(&op.rhs)),
83                };
84
85                let out = self.fuser.output_unhandled(&op.out);
86
87                self.matmul = Some(FusedMatmul::new(
88                    lhs,
89                    rhs,
90                    out,
91                    op.clone().into(),
92                    Default::default(),
93                ));
94            } else {
95                self.fuser.close();
96                self.fuser_fallback.close();
97            }
98        } else {
99            let can_register =
100                self.fuser.can_fuse(operation) && self.fuser_fallback.can_fuse(operation);
101
102            match can_register {
103                true => {
104                    self.fuser.fuse(operation);
105                    self.fuser_fallback.fuse(operation);
106                }
107                false => {
108                    self.fuser.close();
109                    self.fuser_fallback.close();
110                }
111            };
112        }
113    }
114
115    fn finish(&mut self) -> CubeOptimization<R> {
116        let client = R::client(&self.device);
117        let trace = self.fuser.finish();
118        let trace_fallback = self.fuser_fallback.finish();
119
120        let matmul = MatmulOptimization::new(
121            trace,
122            trace_fallback,
123            client,
124            self.device.clone(),
125            self.len(),
126            self.matmul.as_ref().unwrap().clone(),
127        );
128
129        CubeOptimization::Matmul(matmul)
130    }
131
132    fn reset(&mut self) {
133        self.fuser.reset();
134        self.fuser_fallback.reset();
135        self.matmul = None;
136    }
137
138    fn status(&self) -> burn_fusion::FuserStatus {
139        self.fuser.status()
140    }
141
142    fn properties(&self) -> burn_fusion::FuserProperties {
143        self.fuser.properties()
144    }
145
146    fn len(&self) -> usize {
147        // Matmul operation isn't registered in the fuser
148        self.fuser.len() + 1
149    }
150
151    fn clone_dyn(&self) -> Box<dyn OperationFuser<CubeOptimization<R>>> {
152        Box::new(self.clone())
153    }
154}
burn_cubecl_fusion/optim/matmul/fuser.rs

burn_cubecl_fusion/optim/matmul/
fuser.rs
