tract_linalg/frame/block_quant/
helpers.rs

1use byteorder::{ReadBytesExt, WriteBytesExt, LE};
2use std::io::{Cursor, Read, Write};
3use tract_data::internal::*;
4
5pub struct NibbleReader<R> {
6    second_half: Option<i8>,
7    reader: R,
8}
9
10impl<'s> NibbleReader<Cursor<&'s [u8]>> {
11    pub fn for_slice(slice: &'s [u8]) -> Self {
12        NibbleReader::new(Cursor::new(slice))
13    }
14}
15
16impl<R: Read> NibbleReader<R> {
17    pub fn new(reader: R) -> NibbleReader<R> {
18        NibbleReader { reader, second_half: None }
19    }
20
21    pub fn read_f16(&mut self) -> f16 {
22        assert!(self.second_half.is_none());
23        f16::from_bits(self.reader.read_u16::<LE>().unwrap())
24    }
25
26    pub fn read_i4(&mut self) -> i8 {
27        if let Some(second) = self.second_half.take() {
28            second
29        } else {
30            let byte = self.reader.read_u8().unwrap();
31            self.second_half = Some((byte >> 4) as i8);
32            (byte & 0x0F) as i8
33        }
34    }
35}
36
37pub struct NibbleWriter<W> {
38    first_half: Option<i8>,
39    writer: W,
40}
41
42impl<'s> NibbleWriter<Cursor<&'s mut [u8]>> {
43    pub fn for_slice(slice: &'s mut [u8]) -> Self {
44        NibbleWriter::new(Cursor::new(slice))
45    }
46}
47
48impl<W: Write> NibbleWriter<W> {
49    pub fn new(writer: W) -> NibbleWriter<W> {
50        NibbleWriter { writer, first_half: None }
51    }
52
53    pub fn write_f16(&mut self, f: f16) {
54        assert!(self.first_half.is_none());
55        self.writer.write_u16::<LE>(f.to_bits()).unwrap()
56    }
57
58    pub fn write_i4(&mut self, q: i8) {
59        if let Some(first) = self.first_half.take() {
60            self.writer.write_u8(first as u8 | ((q as u8) << 4)).unwrap()
61        } else {
62            self.first_half = Some(q);
63        }
64    }
65}
tract_linalg/frame/block_quant/helpers.rs

tract_linalg/frame/block_quant/
helpers.rs
