1use burn_backend::{
2 ExecutionError, FloatDType, Shape, Slice, TensorData, TensorPrimitive,
3 ops::QTensorOps,
4 quantization::{QuantScheme, QuantizationParametersPrimitive},
5 tensor::{FloatTensor, IntTensor, QuantizedTensor},
6};
7use burn_backend_extension::backend_dispatch;
8
9use crate::{
10 BackendTensor, Dispatch, DispatchAutodiffContext, DispatchDevice, DispatchTensor,
11 DispatchTensorKind,
12};
13
14macro_rules! wrap_q_matmul_concrete {
15 ($Backend:ident, $output:expr, $autodiff:expr) => {{
16 match $output {
17 TensorPrimitive::QFloat(output) => TensorPrimitive::QFloat(DispatchTensor {
18 kind: DispatchTensorKind::$Backend(BackendTensor::Quantized(output)),
19 autodiff: $autodiff,
20 }),
21 TensorPrimitive::Float(output) => match $autodiff {
22 DispatchAutodiffContext::Disabled => TensorPrimitive::Float(DispatchTensor {
23 kind: DispatchTensorKind::$Backend(BackendTensor::Float(output)),
24 autodiff: DispatchAutodiffContext::Disabled,
25 }),
26 #[cfg(feature = "autodiff")]
27 DispatchAutodiffContext::Enabled(strategy) => {
28 with_autodiff_backend!($Backend, strategy, |AD| {
29 TensorPrimitive::Float(DispatchTensor {
30 kind: DispatchTensorKind::Autodiff(alloc::boxed::Box::new(
31 DispatchTensorKind::$Backend(BackendTensor::Autodiff(
32 <AD as burn_backend::AutodiffBackend>::from_inner(output),
33 )),
34 )),
35 autodiff: DispatchAutodiffContext::Enabled(strategy),
36 })
37 })
38 }
39 #[cfg(not(feature = "autodiff"))]
40 DispatchAutodiffContext::Enabled(_) => {
41 panic!("autodiff context requires the `autodiff` feature")
42 }
43 },
44 }
45 }};
46}
47
48#[cfg(feature = "autodiff")]
49macro_rules! wrap_q_matmul_autodiff {
50 ($Backend:ident, $output:expr, $strategy:expr) => {{
51 match $output {
52 TensorPrimitive::QFloat(output) => TensorPrimitive::QFloat(DispatchTensor {
53 kind: DispatchTensorKind::$Backend(BackendTensor::Quantized(output)),
54 autodiff: DispatchAutodiffContext::Enabled($strategy),
55 }),
56 TensorPrimitive::Float(output) => TensorPrimitive::Float(DispatchTensor {
57 kind: DispatchTensorKind::Autodiff(alloc::boxed::Box::new(
58 DispatchTensorKind::$Backend(BackendTensor::Autodiff(output)),
59 )),
60 autodiff: DispatchAutodiffContext::Enabled($strategy),
61 }),
62 }
63 }};
64}
65
66macro_rules! q_matmul_qq_arms {
67 ($lhs:expr, $rhs:expr, $autodiff:expr; $([$Backend:ident, $cfg:meta]),*) => {{
68 match ($lhs.kind, $rhs.kind) {
69 $(
70 #[cfg($cfg)]
71 (DispatchTensorKind::$Backend(lhs), DispatchTensorKind::$Backend(rhs)) => {
72 type B = crate::backends::$Backend;
73 let output = B::q_matmul(
74 TensorPrimitive::QFloat(lhs.quantized()),
75 TensorPrimitive::QFloat(rhs.quantized()),
76 );
77 wrap_q_matmul_concrete!($Backend, output, $autodiff)
78 }
79 )*
80 #[allow(unreachable_patterns)]
81 _ => panic!("q_matmul inputs are on different backends"),
82 }
83 }};
84}
85
86macro_rules! q_matmul_fq_arms {
87 ($lhs:expr, $rhs:expr, $autodiff:expr; $([$Backend:ident, $cfg:meta]),*) => {{
88 match ($lhs.kind, $rhs.kind) {
89 $(
90 #[cfg($cfg)]
91 (DispatchTensorKind::$Backend(lhs), DispatchTensorKind::$Backend(rhs)) => {
92 match $autodiff {
93 DispatchAutodiffContext::Disabled => {
94 type B = crate::backends::$Backend;
95 let output = B::q_matmul(
96 TensorPrimitive::Float(lhs.float()),
97 TensorPrimitive::QFloat(rhs.quantized()),
98 );
99 wrap_q_matmul_concrete!($Backend, output, DispatchAutodiffContext::Disabled)
100 }
101 #[cfg(feature = "autodiff")]
102 DispatchAutodiffContext::Enabled(strategy) => {
103 with_autodiff_backend!($Backend, strategy, |B| {
104 let lhs = <B as burn_backend::AutodiffBackend>::from_inner(lhs.float());
105 let output = B::q_matmul(
106 TensorPrimitive::Float(lhs),
107 TensorPrimitive::QFloat(rhs.quantized()),
108 );
109 wrap_q_matmul_autodiff!($Backend, output, strategy)
110 })
111 }
112 #[cfg(not(feature = "autodiff"))]
113 DispatchAutodiffContext::Enabled(_) => {
114 panic!("autodiff context requires the `autodiff` feature")
115 }
116 }
117 }
118 )*
119 #[cfg(feature = "autodiff")]
120 (DispatchTensorKind::Autodiff(lhs), rhs) => match (*lhs, rhs) {
121 $(
122 #[cfg($cfg)]
123 (DispatchTensorKind::$Backend(lhs), DispatchTensorKind::$Backend(rhs)) => {
124 let DispatchAutodiffContext::Enabled(strategy) = $autodiff else {
125 panic!("an autodiff float primitive must have an enabled autodiff context")
126 };
127 with_autodiff_backend!($Backend, strategy, |B| {
128 let output = B::q_matmul(
129 TensorPrimitive::Float(lhs.autodiff()),
130 TensorPrimitive::QFloat(rhs.quantized()),
131 );
132 wrap_q_matmul_autodiff!($Backend, output, strategy)
133 })
134 }
135 )*
136 #[allow(unreachable_patterns)]
137 _ => panic!("q_matmul inputs are on different backends"),
138 },
139 #[allow(unreachable_patterns)]
140 _ => panic!("q_matmul inputs are on different backends"),
141 }
142 }};
143}
144
145macro_rules! q_matmul_qf_arms {
146 ($lhs:expr, $rhs:expr, $autodiff:expr; $([$Backend:ident, $cfg:meta]),*) => {{
147 match ($lhs.kind, $rhs.kind) {
148 $(
149 #[cfg($cfg)]
150 (DispatchTensorKind::$Backend(lhs), DispatchTensorKind::$Backend(rhs)) => {
151 match $autodiff {
152 DispatchAutodiffContext::Disabled => {
153 type B = crate::backends::$Backend;
154 let output = B::q_matmul(
155 TensorPrimitive::QFloat(lhs.quantized()),
156 TensorPrimitive::Float(rhs.float()),
157 );
158 wrap_q_matmul_concrete!($Backend, output, DispatchAutodiffContext::Disabled)
159 }
160 #[cfg(feature = "autodiff")]
161 DispatchAutodiffContext::Enabled(strategy) => {
162 with_autodiff_backend!($Backend, strategy, |B| {
163 let rhs = <B as burn_backend::AutodiffBackend>::from_inner(rhs.float());
164 let output = B::q_matmul(
165 TensorPrimitive::QFloat(lhs.quantized()),
166 TensorPrimitive::Float(rhs),
167 );
168 wrap_q_matmul_autodiff!($Backend, output, strategy)
169 })
170 }
171 #[cfg(not(feature = "autodiff"))]
172 DispatchAutodiffContext::Enabled(_) => {
173 panic!("autodiff context requires the `autodiff` feature")
174 }
175 }
176 }
177 )*
178 #[cfg(feature = "autodiff")]
179 (lhs, DispatchTensorKind::Autodiff(rhs)) => match (lhs, *rhs) {
180 $(
181 #[cfg($cfg)]
182 (DispatchTensorKind::$Backend(lhs), DispatchTensorKind::$Backend(rhs)) => {
183 let DispatchAutodiffContext::Enabled(strategy) = $autodiff else {
184 panic!("an autodiff float primitive must have an enabled autodiff context")
185 };
186 with_autodiff_backend!($Backend, strategy, |B| {
187 let output = B::q_matmul(
188 TensorPrimitive::QFloat(lhs.quantized()),
189 TensorPrimitive::Float(rhs.autodiff()),
190 );
191 wrap_q_matmul_autodiff!($Backend, output, strategy)
192 })
193 }
194 )*
195 #[allow(unreachable_patterns)]
196 _ => panic!("q_matmul inputs are on different backends"),
197 },
198 #[allow(unreachable_patterns)]
199 _ => panic!("q_matmul inputs are on different backends"),
200 }
201 }};
202}
203
204#[backend_dispatch]
205impl QTensorOps<Self> for Dispatch {
206 fn q_from_data(data: TensorData, device: &DispatchDevice) -> QuantizedTensor<Self> {
207 B::q_from_data(data, device)
208 }
209
210 fn quantize(
211 tensor: FloatTensor<Self>,
212 scheme: &QuantScheme,
213 qparams: QuantizationParametersPrimitive<Self>,
214 ) -> QuantizedTensor<Self> {
215 B::quantize(tensor, scheme, qparams)
216 }
217
218 fn dequantize(tensor: QuantizedTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
219 B::dequantize(tensor, dtype)
220 }
221
222 #[backend_dispatch(skip)]
223 fn q_to_device(
224 tensor: QuantizedTensor<Self>,
225 device: &DispatchDevice,
226 ) -> QuantizedTensor<Self> {
227 to_device!(
228 Quantized,
229 quantized,
230 tensor,
231 device,
232 q_to_device,
233 |inner, device| {
234 let data =
235 burn_backend::read_sync(B1::q_into_data(inner)).expect("Should read data");
236 B2::q_from_data(data, device)
237 }
238 )
239 }
240
241 fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
242 B::q_reshape(tensor, shape)
243 }
244
245 async fn q_into_data(tensor: QuantizedTensor<Self>) -> Result<TensorData, ExecutionError> {
246 B::q_into_data(tensor).await
247 }
248
249 fn q_expand(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
250 B::q_expand(tensor, shape)
251 }
252
253 fn q_swap_dims(
254 tensor: QuantizedTensor<Self>,
255 dim1: usize,
256 dim2: usize,
257 ) -> QuantizedTensor<Self> {
258 B::q_swap_dims(tensor, dim1, dim2)
259 }
260
261 fn q_permute(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
262 B::q_permute(tensor, axes)
263 }
264
265 fn q_flip(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
266 B::q_flip(tensor, axes)
267 }
268
269 fn q_select(
270 tensor: QuantizedTensor<Self>,
271 dim: usize,
272 indices: IntTensor<Self>,
273 ) -> QuantizedTensor<Self> {
274 B::q_select(tensor, dim, indices)
275 }
276
277 fn q_slice(tensor: QuantizedTensor<Self>, slices: &[Slice]) -> QuantizedTensor<Self> {
278 B::q_slice(tensor, slices)
279 }
280
281 #[backend_dispatch(skip)]
282 fn q_matmul(lhs: TensorPrimitive<Self>, rhs: TensorPrimitive<Self>) -> TensorPrimitive<Self> {
283 match (lhs, rhs) {
284 (TensorPrimitive::QFloat(lhs), TensorPrimitive::QFloat(rhs)) => {
285 let autodiff = lhs.autodiff.merge(rhs.autodiff);
286 backend_list!(q_matmul_qq_arms, lhs, rhs, autodiff)
288 }
289 (TensorPrimitive::Float(lhs), TensorPrimitive::QFloat(rhs)) => {
290 let autodiff = lhs.autodiff.merge(rhs.autodiff);
291 #[cfg(feature = "autodiff")]
292 match (
293 matches!(&lhs.kind, DispatchTensorKind::Autodiff(_)),
294 lhs.autodiff,
295 ) {
296 (true, DispatchAutodiffContext::Enabled(_))
297 | (false, DispatchAutodiffContext::Disabled) => {}
298 (true, DispatchAutodiffContext::Disabled) => {
299 panic!("an autodiff float primitive must have an enabled autodiff context")
300 }
301 (false, DispatchAutodiffContext::Enabled(_)) => {
302 panic!("an enabled float tensor must use an autodiff primitive")
303 }
304 }
305 backend_list!(q_matmul_fq_arms, lhs, rhs, autodiff)
306 }
307 (TensorPrimitive::QFloat(lhs), TensorPrimitive::Float(rhs)) => {
308 let autodiff = lhs.autodiff.merge(rhs.autodiff);
309 #[cfg(feature = "autodiff")]
310 match (
311 matches!(&rhs.kind, DispatchTensorKind::Autodiff(_)),
312 rhs.autodiff,
313 ) {
314 (true, DispatchAutodiffContext::Enabled(_))
315 | (false, DispatchAutodiffContext::Disabled) => {}
316 (true, DispatchAutodiffContext::Disabled) => {
317 panic!("an autodiff float primitive must have an enabled autodiff context")
318 }
319 (false, DispatchAutodiffContext::Enabled(_)) => {
320 panic!("an enabled float tensor must use an autodiff primitive")
321 }
322 }
323 backend_list!(q_matmul_qf_arms, lhs, rhs, autodiff)
324 }
325 _ => unreachable!(),
326 }
327 }
328}