1use burn_backend::{
2 DeviceOps, ExecutionError, FloatDType, Shape, Slice, TensorData, TensorMetadata,
3 TensorPrimitive,
4 ops::QTensorOps,
5 quantization::{QuantPropagation, QuantScheme, QuantizationParametersPrimitive},
6 tensor::{FloatTensor, IntTensor, QuantizedTensor},
7};
8
9use crate::{Dispatch, DispatchDevice};
10
11impl QTensorOps<Self> for Dispatch {
12 fn q_from_data(data: TensorData, device: &DispatchDevice) -> QuantizedTensor<Self> {
13 creation_op!(Quantized, device, |device| B::q_from_data(data, device))
14 }
15
16 fn quantize(
17 tensor: FloatTensor<Self>,
18 scheme: &QuantScheme,
19 qparams: QuantizationParametersPrimitive<Self>,
20 ) -> QuantizedTensor<Self> {
21 binary_op!(
22 (tensor, float),
23 (qparams.scales, float),
24 |tensor, scales| {
25 B::quantize(tensor, scheme, QuantizationParametersPrimitive { scales })
26 } => Quantized
27 )
28 }
29
30 fn dequantize(tensor: QuantizedTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
31 unary_op!(tensor, quantized, |tensor| B::dequantize(tensor, dtype) => Float)
32 }
33
34 fn q_to_device(
35 tensor: QuantizedTensor<Self>,
36 device: &DispatchDevice,
37 ) -> QuantizedTensor<Self> {
38 to_device!(
39 Quantized,
40 quantized,
41 tensor,
42 device,
43 q_to_device,
44 |inner, device| {
45 let data =
46 burn_backend::read_sync(B1::q_into_data(inner)).expect("Should read data");
47 B2::q_from_data(data, device)
48 }
49 )
50 }
51
52 fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
53 unary_op!(tensor, quantized, |tensor| B::q_reshape(tensor, shape) => Quantized)
54 }
55
56 async fn q_into_data(tensor: QuantizedTensor<Self>) -> Result<TensorData, ExecutionError> {
57 unary_op!(tensor, quantized, |tensor| B::q_into_data(tensor).await)
58 }
59
60 fn q_expand(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
61 unary_op!(tensor, quantized, |tensor| B::q_expand(tensor, shape) => Quantized)
62 }
63
64 fn q_swap_dims(
65 tensor: QuantizedTensor<Self>,
66 dim1: usize,
67 dim2: usize,
68 ) -> QuantizedTensor<Self> {
69 unary_op!(tensor, quantized, |tensor| B::q_swap_dims(tensor, dim1, dim2) => Quantized)
70 }
71
72 fn q_permute(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
73 unary_op!(tensor, quantized, |tensor| B::q_permute(tensor, axes) => Quantized)
74 }
75
76 fn q_flip(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
77 unary_op!(tensor, quantized, |tensor| B::q_flip(tensor, axes) => Quantized)
78 }
79
80 fn q_select(
81 tensor: QuantizedTensor<Self>,
82 dim: usize,
83 indices: IntTensor<Self>,
84 ) -> QuantizedTensor<Self> {
85 binary_op!(
86 (tensor, quantized),
87 (indices, int),
88 |tensor, indices| B::q_select(tensor, dim, indices) => Quantized
89 )
90 }
91
92 fn q_slice(tensor: QuantizedTensor<Self>, slices: &[Slice]) -> QuantizedTensor<Self> {
93 unary_op!(tensor, quantized, |tensor| B::q_slice(tensor, slices) => Quantized)
94 }
95
96 fn q_matmul(lhs: TensorPrimitive<Self>, rhs: TensorPrimitive<Self>) -> TensorPrimitive<Self> {
97 match (lhs, rhs) {
99 (TensorPrimitive::QFloat(lhs), TensorPrimitive::QFloat(rhs)) => {
100 let propagation = lhs.device().defaults().quantization.propagation;
101 if matches!(propagation, QuantPropagation::Propagate) {
102 let out = binary_op!(
103 (lhs, quantized),
104 (rhs, quantized),
105 |lhs, rhs| {
106 if let TensorPrimitive::QFloat(out) = B::q_matmul(
107 TensorPrimitive::QFloat(lhs),
108 TensorPrimitive::QFloat(rhs),
109 ) {
110 out
111 } else {
112 unreachable!()
113 }
114 } => Quantized
115 );
116 TensorPrimitive::QFloat(out)
117 } else {
118 let out = binary_op!(
119 (lhs, quantized),
120 (rhs, quantized),
121 |lhs, rhs| {
122 if let TensorPrimitive::Float(out) = B::q_matmul(
123 TensorPrimitive::QFloat(lhs),
124 TensorPrimitive::QFloat(rhs),
125 ) {
126 out
127 } else {
128 unreachable!()
129 }
130 } => Float
131 );
132 TensorPrimitive::Float(out)
133 }
134 }
135 (TensorPrimitive::Float(lhs), TensorPrimitive::QFloat(rhs)) => {
136 let propagation = rhs.device().defaults().quantization.propagation;
137 if matches!(propagation, QuantPropagation::Propagate) {
138 let out = binary_op!(
139 (lhs, float),
140 (rhs, quantized),
141 |lhs, rhs| {
142 if let TensorPrimitive::QFloat(out) = B::q_matmul(
143 TensorPrimitive::Float(lhs),
144 TensorPrimitive::QFloat(rhs),
145 ) {
146 out
147 } else {
148 unreachable!()
149 }
150 } => Quantized
151 );
152 TensorPrimitive::QFloat(out)
153 } else {
154 let out = binary_op!(
155 (lhs, float),
156 (rhs, quantized),
157 |lhs, rhs| {
158 if let TensorPrimitive::Float(out) = B::q_matmul(
159 TensorPrimitive::Float(lhs),
160 TensorPrimitive::QFloat(rhs),
161 ) {
162 out
163 } else {
164 unreachable!()
165 }
166 } => Float
167 );
168 TensorPrimitive::Float(out)
169 }
170 }
171 (TensorPrimitive::QFloat(lhs), TensorPrimitive::Float(rhs)) => {
172 let propagation = lhs.device().defaults().quantization.propagation;
173 if matches!(propagation, QuantPropagation::Propagate) {
174 let out = binary_op!(
175 (lhs, quantized),
176 (rhs, float),
177 |lhs, rhs| {
178 if let TensorPrimitive::QFloat(out) = B::q_matmul(
179 TensorPrimitive::QFloat(lhs),
180 TensorPrimitive::Float(rhs),
181 ) {
182 out
183 } else {
184 unreachable!()
185 }
186 } => Quantized
187 );
188 TensorPrimitive::QFloat(out)
189 } else {
190 let out = binary_op!(
191 (lhs, quantized),
192 (rhs, float),
193 |lhs, rhs| {
194 if let TensorPrimitive::Float(out) = B::q_matmul(
195 TensorPrimitive::QFloat(lhs),
196 TensorPrimitive::Float(rhs),
197 ) {
198 out
199 } else {
200 unreachable!()
201 }
202 } => Float
203 );
204 TensorPrimitive::Float(out)
205 }
206 }
207 _ => unreachable!(),
208 }
209 }
210}