tract_linalg/frame/block_quant/
mod.rs

1use downcast_rs::{Downcast, impl_downcast};
2use dyn_clone::{DynClone, clone_box};
3use dyn_eq::DynEq;
4use dyn_hash::DynHash;
5use num_traits::Zero;
6use tract_data::internal::*;
7use tract_data::itertools::Itertools;
8
9use std::alloc::Layout;
10use std::borrow::Cow;
11use std::fmt::{Debug, Display};
12use std::hash::Hash;
13use std::sync::Arc;
14
15mod helpers;
16mod q20t;
17mod q4_0;
18mod q8_1;
19mod storage;
20mod value;
21
22pub use helpers::{CrumbReader, CrumbWriter, NibbleReader, NibbleWriter};
23pub use q4_0::Q4_0;
24pub use q8_1::Q8_1;
25pub use q20t::{BaseQ2_0_T, Q2_0_T};
26pub use storage::{BlockQuantStorage, block_quant_slice};
27pub use value::{BlockQuantFact, PackedBlockQuantFact};
28
29use crate::mmm::{EagerPackedInput, MMMInputFormat};
30use crate::pack::PackedFormat;
31
32use crate::WeightType;
33
34use super::mmm::MMMInputValue;
35
36pub trait BlockQuant:
37    Debug + Display + Send + Sync + DynClone + DynHash + dyn_eq::DynEq + Downcast
38{
39    fn block_len(&self) -> usize;
40
41    fn block_bytes(&self) -> usize;
42
43    fn dequant_block_f32(&self, quant: &[u8], block: &mut [f32]);
44    fn dequant_block_f16(&self, quant: &[u8], block: &mut [f16]);
45    fn quant_block_f16(&self, block: &[f16], quant: &mut [u8]);
46    fn quant_block_f32(&self, block: &[f32], quant: &mut [u8]);
47
48    fn quant_f16(&self, input: &[f16]) -> TractResult<Blob> {
49        unsafe {
50            let blocks = input.len() / self.block_len();
51            let mut quant = Blob::for_layout(
52                Layout::from_size_align(blocks * self.block_bytes(), 128).unwrap(),
53            );
54            for b in 0..blocks {
55                let block = &input[b * self.block_len()..][..self.block_len()];
56                let qblock = &mut quant[b * self.block_bytes()..][..self.block_bytes()];
57                self.quant_block_f16(block, qblock);
58            }
59            Ok(quant)
60        }
61    }
62
63    fn quant_f32(&self, input: &[f32]) -> TractResult<Blob> {
64        unsafe {
65            let blocks = input.len() / self.block_len();
66            let mut quant = Blob::for_layout(
67                Layout::from_size_align(blocks * self.block_bytes(), 128).unwrap(),
68            );
69            for b in 0..blocks {
70                let block = &input[b * self.block_len()..][..self.block_len()];
71                let qblock = &mut quant[b * self.block_bytes()..][..self.block_bytes()];
72                self.quant_block_f32(block, qblock);
73            }
74            Ok(quant)
75        }
76    }
77
78    fn dequant_f32(&self, input: &[u8]) -> TractResult<Tensor> {
79        unsafe {
80            let blocks = input.len() / self.block_bytes();
81            let mut tensor = Tensor::uninitialized::<f32>(&[blocks * self.block_len()])?;
82            let mut tensor_plain = tensor.try_as_plain_mut()?;
83            let slice = tensor_plain.as_slice_mut::<f32>()?;
84            for b in 0..blocks {
85                let block = &mut slice[b * self.block_len()..][..self.block_len()];
86                let qblock = &input[b * self.block_bytes()..][..self.block_bytes()];
87                self.dequant_block_f32(qblock, block);
88            }
89            Ok(tensor)
90        }
91    }
92
93    fn dequant_f16(&self, input: &[u8]) -> TractResult<Tensor> {
94        unsafe {
95            let blocks = input.len() / self.block_bytes();
96            let mut tensor = Tensor::uninitialized::<f16>(&[blocks * self.block_len()])?;
97            let mut tensor_plain = tensor.try_as_plain_mut()?;
98            let slice = tensor_plain.as_slice_mut::<f16>()?;
99            for b in 0..blocks {
100                let block = &mut slice[b * self.block_len()..][..self.block_len()];
101                let qblock = &input[b * self.block_bytes()..][..self.block_bytes()];
102                self.dequant_block_f16(qblock, block);
103            }
104            Ok(tensor)
105        }
106    }
107
108    fn extract_at_offset_f16(&self, input: &[u8], offset: usize) -> f16 {
109        let len = self.block_len();
110        let block_id = offset / len;
111        let mut block = vec![f16::zero(); self.block_len()];
112        self.dequant_block_f16(
113            &input[block_id * self.block_bytes()..][..self.block_bytes()],
114            &mut block,
115        );
116        block[offset % len]
117    }
118
119    fn extract_at_offset_f32(&self, input: &[u8], offset: usize) -> f32 {
120        let len = self.block_len();
121        let block_id = offset / len;
122        let mut block = vec![f32::zero(); self.block_len()];
123        self.dequant_block_f32(
124            &input[block_id * self.block_bytes()..][..self.block_bytes()],
125            &mut block,
126        );
127        block[offset % len]
128    }
129
130    fn simulate_precision_loss(
131        &self,
132        mut tensor: Tensor,
133        block_axis: usize,
134    ) -> TractResult<Tensor> {
135        ensure!(block_axis == tensor.rank() - 1);
136        ensure!(tensor.shape()[block_axis] % self.block_len() == 0);
137        let mut scratch = vec![0u8; self.block_bytes()];
138        if tensor.datum_type() == f32::datum_type() {
139            let mut tensor_plain = tensor.try_as_plain_mut()?;
140            for block in tensor_plain.as_slice_mut::<f32>()?.chunks_mut(self.block_len()) {
141                self.quant_block_f32(block, &mut scratch);
142                self.dequant_block_f32(&scratch, block);
143            }
144            Ok(tensor)
145        } else if tensor.datum_type() == f16::datum_type() {
146            let mut tensor_plain = tensor.try_as_plain_mut()?;
147            for block in tensor_plain.as_slice_mut::<f16>()?.chunks_mut(self.block_len()) {
148                self.quant_block_f16(block, &mut scratch);
149                self.dequant_block_f16(&scratch, block);
150            }
151            Ok(tensor)
152        } else {
153            todo!()
154        }
155    }
156
157    fn pack(
158        &self,
159        input: &[u8],
160        k: usize,
161        r: usize,
162        zip: usize,
163        scales_at_end: bool,
164    ) -> TractResult<EagerPackedInput>;
165
166    unsafe fn extract_packed_panel(
167        &self,
168        value: &EagerPackedInput,
169        target: &PackedFormat,
170        panel: usize,
171        scratch: *mut u8,
172    ) -> TractResult<()>;
173
174    fn extract_at_mn_f16(
175        &self,
176        value: &EagerPackedInput,
177        mn: usize,
178        target: &mut [f16],
179    ) -> TractResult<()>;
180
181    fn extract_at_mn_f32(
182        &self,
183        value: &EagerPackedInput,
184        mn: usize,
185        target: &mut [f32],
186    ) -> TractResult<()>;
187}
188
189dyn_clone::clone_trait_object!(BlockQuant);
190dyn_hash::hash_trait_object!(BlockQuant);
191dyn_eq::eq_trait_object!(BlockQuant);
192impl_downcast!(BlockQuant);
193
194#[allow(clippy::derived_hash_with_manual_eq)]
195#[derive(Clone, Hash)]
196pub struct PackedBlockQuantFormat {
197    pub bq: Box<dyn BlockQuant>,
198    pub r: usize,
199    pub zip: usize,
200    pub scales_at_end: bool,
201}
202
203impl PartialEq for PackedBlockQuantFormat {
204    fn eq(&self, other: &Self) -> bool {
205        *self.bq == *other.bq
206            && self.r == other.r
207            && self.zip == other.zip
208            && self.scales_at_end == other.scales_at_end
209    }
210}
211
212impl Eq for PackedBlockQuantFormat {}
213
214impl Display for PackedBlockQuantFormat {
215    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
216        write!(f, "Packed{}[{}]", &*self.bq, self.r)?;
217        if self.zip != 0 {
218            write!(f, "Z{}", self.zip)?;
219        }
220        if self.scales_at_end {
221            write!(f, "Se")?;
222        }
223        Ok(())
224    }
225}
226
227impl Debug for PackedBlockQuantFormat {
228    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
229        <Self as Display>::fmt(self, f)
230    }
231}
232
233impl PackedBlockQuantFormat {
234    pub fn new(bq: &dyn BlockQuant, r: usize, zip: usize, scales_at_end: bool) -> Self {
235        PackedBlockQuantFormat { bq: clone_box(bq), r, zip, scales_at_end }
236    }
237
238    pub fn simulate_precision_loss(
239        &self,
240        tensor: Tensor,
241        block_axis: usize,
242    ) -> TractResult<Tensor> {
243        self.bq.simulate_precision_loss(tensor, block_axis)
244    }
245
246    pub fn pack(&self, input: &[u8], k: usize) -> TractResult<EagerPackedInput> {
247        self.bq.pack(input, k, self.r, self.zip, self.scales_at_end)
248    }
249}
250
251impl MMMInputFormat for PackedBlockQuantFormat {
252    fn prepare_tensor(&self, t: &Tensor, _k_axis: usize, _mn_axis: usize) -> TractResult<Tensor> {
253        let bqs = t.try_storage_as::<BlockQuantStorage>()?;
254        let num_groups: usize =
255            if t.rank() > 2 { t.shape()[..t.rank() - 2].iter().product() } else { 1 };
256        let m_per_group = t.shape()[t.rank().saturating_sub(2)];
257        let k = *t.shape().last().unwrap();
258        let values = (0..num_groups)
259            .map(|g| {
260                let slice = block_quant_slice(bqs.value(), &*self.bq, m_per_group, k, g);
261                let packed = self.pack(slice, k)?;
262                Ok(Box::new(packed) as Box<dyn MMMInputValue>)
263            })
264            .collect::<TractResult<Vec<_>>>()?;
265        let leading_shape = &t.shape()[..t.rank().saturating_sub(2)];
266        Ok(crate::mmm::PackedMatrixStorage::new_batched(leading_shape, values)
267            .into_tensor(t.datum_type()))
268    }
269
270    fn prepare_one(
271        &self,
272        t: &Tensor,
273        k_axis: usize,
274        mn_axis: usize,
275    ) -> TractResult<Box<dyn MMMInputValue>> {
276        // this code path is essentially there for test scenarios
277        let t = if t.is_plain() && t.datum_type().is_number() {
278            let k = t.shape()[k_axis];
279            let m = t.shape()[mn_axis];
280            assert!(k % self.bq.block_len() == 0);
281            let t: Cow<Tensor> = if k_axis == 1 && mn_axis == 0 {
282                Cow::Borrowed(t)
283            } else {
284                Cow::Owned(t.clone().move_axis(1, 0)?)
285            };
286            let quant = if t.datum_type() == f32::datum_type() {
287                self.bq.quant_f32(t.try_as_plain()?.as_slice()?)?
288            } else if t.datum_type() == f16::datum_type() {
289                self.bq.quant_f16(t.try_as_plain()?.as_slice()?)?
290            } else {
291                todo!()
292            };
293            Cow::Owned(
294                BlockQuantStorage::new(self.bq.clone(), m, k, Arc::new(quant))?
295                    .into_tensor_with_shape(t.datum_type(), &[1, m, k]),
296            )
297        } else {
298            Cow::Borrowed(t)
299        };
300        ensure!(mn_axis == 0);
301        ensure!(k_axis == 1);
302        let bqs = t.try_storage_as::<BlockQuantStorage>()?;
303        let k = *t.shape().last().unwrap();
304        let packed = self.pack(bqs.value(), k)?;
305        Ok(Box::new(packed))
306    }
307
308    fn precursor(&self) -> WeightType {
309        WeightType::BlockQuant(self.bq.clone())
310    }
311
312    fn k_alignment(&self) -> usize {
313        self.bq.block_len()
314    }
315
316    fn r(&self) -> usize {
317        self.r
318    }
319
320    fn mem_size(&self, k: TDim, mn: TDim) -> TDim {
321        k * mn * self.bq.block_bytes() / self.bq.block_len()
322    }
323
324    fn extract_at_mn_f16(
325        &self,
326        data: &EagerPackedInput,
327        mn: usize,
328        slice: &mut [f16],
329    ) -> TractResult<()> {
330        self.bq.extract_at_mn_f16(data, mn, slice)
331    }
332
333    fn extract_at_mn_f32(
334        &self,
335        data: &EagerPackedInput,
336        mn: usize,
337        slice: &mut [f32],
338    ) -> TractResult<()> {
339        self.bq.extract_at_mn_f32(data, mn, slice)
340    }
341}
tract_linalg/frame/block_quant/mod.rs

tract_linalg/frame/block_quant/
mod.rs
