feat(ggus): GGuf 文件格式移动到 ggus 库

Signed-off-by: YdrMaster <ydrml@hotmail.com>

xtask/src/gguf_file.rs → ggus/src/file.rs

@@ -1,17 +1,16 @@
-use ggus::{GGufFileHeader, GGufMetaKVPairs, GGufReadError, GGufTensors};
+use crate::{pad, GGufFileHeader, GGufMetaKVPairs, GGufReadError, GGufTensors};
 use std::{error::Error, fmt};
 
 #[derive(Clone)]
-pub(crate) struct GGufFile<'a> {
-    #[allow(unused)]
+pub struct GGuf<'a> {
     pub header: GGufFileHeader,
     pub meta_kvs: GGufMetaKVPairs<'a>,
     pub tensors: GGufTensors<'a>,
     pub data: &'a [u8],
 }
 
 #[derive(Debug)]
-pub(crate) enum GGufError {
+pub enum GGufError {
     MagicMismatch,
     EndianNotSupport,
     VersionNotSupport,
@@ -31,8 +30,8 @@ impl fmt::Display for GGufError {
 
 impl Error for GGufError {}
 
-impl<'a> GGufFile<'a> {
-    pub fn new(data: &'a [u8]) -> Result<Self, GGufError> {
+impl<'a> GGuf<'a> {
+    pub fn scan(data: &'a [u8]) -> Result<Self, GGufError> {
         let header = unsafe { data.as_ptr().cast::<GGufFileHeader>().read() };
         if !header.is_magic_correct() {
             return Err(GGufError::MagicMismatch);
@@ -66,8 +65,3 @@ impl<'a> GGufFile<'a> {
         })
     }
 }
-
-#[inline(always)]
-pub(crate) const fn pad(pos: usize, align: usize) -> usize {
-    (align - pos % align) % align
-}

ggus/src/header.rs

@@ -1,5 +1,4 @@
-use crate::sizeof;
-use std::str::Utf8Error;
+use std::str::Utf8Error;
 
 #[derive(Clone, Default, Debug)]
 #[repr(C)]
@@ -45,6 +44,6 @@ impl GGufFileHeader {
 
     #[inline]
     pub const fn nbytes(&self) -> usize {
-        sizeof!(Self)
+        size_of::<Self>()
     }
 }

ggus/src/lib.rs

@@ -1,27 +1,26 @@
 #![doc = include_str!("../README.md")]
 #![deny(warnings)]
 
+mod file;
 mod header;
 mod metadata;
 mod name;
-mod reader;
+mod read;
 mod tensor;
-mod writer;
+mod write;
 
+pub use file::{GGuf, GGufError};
 pub use header::GGufFileHeader;
 pub use metadata::{
     utok, GGufArray, GGufFileType, GGufMetaDataValueType, GGufMetaKV, GGufMetaKVPairs,
     GGufTokenType, DEFAULT_ALIGNMENT, GENERAL_ALIGNMENT,
 };
 pub use name::GGufFileName;
-pub use reader::{GGufReadError, GGufReader};
+pub use read::{GGufReadError, GGufReader};
 pub use tensor::{GGmlType, GGufTensorInfo, GGufTensors};
-pub use writer::{GGufMetaWriter, GGufSimulator, GGufTensorWriter};
+pub use write::{GGufMetaWriter, GGufSimulator, GGufTensorWriter};
 
-macro_rules! sizeof {
-    ($ty:ty) => {
-        std::mem::size_of::<$ty>()
-    };
+#[inline(always)]
+const fn pad(pos: usize, align: usize) -> usize {
+    (align - pos % align) % align
 }
-
-use sizeof;

ggus/src/metadata/general.rs

@@ -1,4 +1,4 @@
-use super::{GGufFileType, GGufMetaDataValueType as Ty, GGufMetaKVPairs};
+use crate::{GGufFileType, GGufMetaDataValueType as Ty, GGufMetaKVPairs};
 
 pub const GENERAL_ALIGNMENT: &str = "general.alignment";
 pub const DEFAULT_ALIGNMENT: usize = 32;

ggus/src/metadata/llm.rs

@@ -1,4 +1,4 @@
-use super::{GGufMetaDataValueType, GGufMetaKVPairs};
+use crate::{GGufMetaDataValueType, GGufMetaKVPairs};
 
 pub struct LlmMeta<'a>(GGufMetaKVPairs<'a>, &'a str);
 

ggus/src/metadata/mod.rs

@@ -4,10 +4,7 @@ mod general;
 mod llm;
 mod tokenizer;
 
-use crate::{
-    reader::{GGufReadError, GGufReader},
-    sizeof,
-};
+use crate::{GGufReadError, GGufReader};
 use indexmap::IndexMap;
 use std::{hash::Hash, slice::from_raw_parts};
 
@@ -234,7 +231,7 @@ impl<'a> GGufMetaKV<'a> {
                 self.key
                     .as_ptr()
                     .add(self.key.len())
-                    .add(sizeof!(GGufMetaDataValueType)),
+                    .add(size_of::<GGufMetaDataValueType>()),
                 self.len,
             )
         }

ggus/src/metadata/tokenizer.rs

@@ -1,5 +1,4 @@
-use super::{GGufMetaDataValueType, GGufMetaKVPairs};
-use crate::reader::GGufReader;
+use crate::{GGufMetaDataValueType, GGufMetaKVPairs, GGufReader};
 use std::marker::PhantomData;
 
 #[repr(transparent)]
@@ -78,7 +77,7 @@ impl<'a> GGufArray<'a, str> {
     }
 }
 
-impl<'a, T: Copy> Iterator for GGufArray<'a, T> {
+impl<'a, T: Copy + 'static> Iterator for GGufArray<'a, T> {
     type Item = T;
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {

ggus/src/name.rs

@@ -1,5 +1,4 @@
-use core::fmt;
-use std::str::FromStr;
+use std::{fmt, str::FromStr};
 
 #[derive(Clone, Debug)]
 pub struct GGufFileName {

ggus/src/reader.rs → ggus/src/read.rs

@@ -1,4 +1,4 @@
-use crate::{sizeof, GGufMetaDataValueType};
+use crate::GGufMetaDataValueType;
 use std::str::Utf8Error;
 
 pub struct GGufReader<'a> {
@@ -25,8 +25,8 @@ impl<'a> GGufReader<'a> {
         self.cursor
     }
 
-    pub(crate) fn skip<T: Copy>(&mut self, len: usize) -> Result<(), GGufReadError> {
-        let len = len * sizeof!(T);
+    pub(crate) fn skip<T: Copy + 'static>(&mut self, len: usize) -> Result<(), GGufReadError> {
+        let len = len * size_of::<T>();
         let data = &self.data[self.cursor..];
         if data.len() >= len {
             self.cursor += len;
@@ -36,7 +36,7 @@ impl<'a> GGufReader<'a> {
         }
     }
 
-    pub fn read<T: Copy>(&mut self) -> Result<T, GGufReadError> {
+    pub fn read<T: Copy + 'static>(&mut self) -> Result<T, GGufReadError> {
         let ptr = self.data[self.cursor..].as_ptr().cast::<T>();
         self.skip::<T>(1)?;
         Ok(unsafe { ptr.read_unaligned() })

ggus/src/tensor.rs

@@ -1,9 +1,8 @@
-use crate::{
-    reader::{GGufReadError, GGufReader},
-    sizeof,
-};
+use crate::{GGufReadError, GGufReader};
 use indexmap::IndexMap;
-use std::{hash::Hash, marker::PhantomData, slice::from_raw_parts, str::from_utf8_unchecked};
+use std::{
+    hash::Hash, marker::PhantomData, mem::size_of, slice::from_raw_parts, str::from_utf8_unchecked,
+};
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
 #[repr(u32)]
@@ -45,7 +44,7 @@ pub enum GGmlType {
 impl GGmlType {
     fn nbytes(self) -> usize {
         match self {
-            Self::F32 => sizeof!(f32),
+            Self::F32 => size_of::<f32>(),
             Self::F16 => 2,
             Self::Q4_0 => todo!(),
             Self::Q4_1 => todo!(),
@@ -67,11 +66,11 @@ impl GGmlType {
             Self::IQ3S => todo!(),
             Self::IQ2S => todo!(),
             Self::IQ4XS => todo!(),
-            Self::I8 => sizeof!(i8),
-            Self::I16 => sizeof!(i16),
-            Self::I32 => sizeof!(i32),
-            Self::I64 => sizeof!(i64),
-            Self::F64 => sizeof!(f64),
+            Self::I8 => size_of::<i8>(),
+            Self::I16 => size_of::<i16>(),
+            Self::I32 => size_of::<i32>(),
+            Self::I64 => size_of::<i64>(),
+            Self::F64 => size_of::<f64>(),
             Self::IQ1M => todo!(),
             _ => unimplemented!(),
         }
@@ -162,13 +161,13 @@ impl<'a> GGufTensorInfo<'a> {
             let ptr = name.as_ptr().add(name.len());
             let ndim = ptr.cast::<u32>().read_unaligned() as usize;
 
-            let ptr = ptr.add(sizeof!(u32));
+            let ptr = ptr.add(size_of::<u32>());
             let shape = ptr;
 
-            let ptr = ptr.add(ndim * sizeof!(u64));
+            let ptr = ptr.add(ndim * size_of::<u64>());
             let ggml_type = ptr.cast::<GGmlType>().read_unaligned();
 
-            let ptr = ptr.add(sizeof!(u32));
+            let ptr = ptr.add(size_of::<u32>());
             let offset = ptr.cast::<u64>().read_unaligned();
 
             let mut body = vec![0u64; 4 + ndim].into_boxed_slice();

ggus/src/write/internal.rs

@@ -0,0 +1,118 @@
+use crate::{pad, GGmlType, GGufFileHeader, GGufMetaDataValueType, GENERAL_ALIGNMENT};
+use internal::Internal;
+use std::{
+    io::{Result, Write},
+    slice::from_raw_parts,
+};
+
+pub(super) struct GGufWriter<T: Write>(Internal<T>);
+
+impl<T: Write> GGufWriter<T> {
+    #[inline]
+    fn write<U: Copy + 'static>(&mut self, val: &[U]) -> Result<()> {
+        self.0
+            .write_bytes(unsafe { from_raw_parts(val.as_ptr().cast(), size_of_val(val)) })
+    }
+
+    #[inline]
+    fn write_str(&mut self, val: impl AsRef<str>) -> Result<()> {
+        let val = val.as_ref().as_bytes();
+        self.write(&[val.len() as u64])?;
+        self.write(val)
+    }
+
+    #[inline]
+    pub fn new(writer: T, header: GGufFileHeader) -> Result<Self> {
+        let mut ans = Self(Internal::new(writer));
+        ans.write(unsafe {
+            from_raw_parts(
+                &header as *const _ as *const u8,
+                size_of::<GGufFileHeader>(),
+            )
+        })?;
+        Ok(ans)
+    }
+
+    #[inline]
+    pub const fn written_bytes(&self) -> usize {
+        self.0.written_bytes()
+    }
+
+    #[inline]
+    pub fn write_alignment(&mut self, align: usize) -> Result<usize> {
+        self.write_meta_kv(
+            GENERAL_ALIGNMENT,
+            GGufMetaDataValueType::U32,
+            (align as u32).to_le_bytes(),
+        )
+        .map(Option::unwrap)
+    }
+
+    pub fn write_meta_kv(
+        &mut self,
+        key: impl AsRef<str>,
+        ty: GGufMetaDataValueType,
+        val: impl AsRef<[u8]>,
+    ) -> Result<Option<usize>> {
+        let key = key.as_ref();
+        let val = val.as_ref();
+
+        self.write_str(key)?;
+        self.write(&[ty])?;
+        self.write(val)?;
+
+        Ok(if key == GENERAL_ALIGNMENT {
+            let &[a, b, c, d] = val else {
+                panic!("general.alignment must be an u32")
+            };
+            Some(u32::from_le_bytes([a, b, c, d]) as _)
+        } else {
+            None
+        })
+    }
+
+    pub fn write_tensor_info(
+        &mut self,
+        name: &str,
+        shape: &[u64],
+        data_type: GGmlType,
+        offset: u64,
+    ) -> Result<()> {
+        self.write_str(name)?;
+        self.write(&[shape.len() as u32])?;
+        self.write(shape)?;
+        self.write(&[data_type])?;
+        self.write(&[offset])
+    }
+
+    pub fn write_data(&mut self, data: &[u8], align: usize) -> Result<()> {
+        for _ in 0..pad(self.written_bytes(), align) {
+            self.write(&[0u8])?;
+        }
+        self.write(data)
+    }
+}
+
+mod internal {
+    use std::io::{BufWriter, Result, Write};
+
+    pub(super) struct Internal<T: Write>(BufWriter<T>, usize);
+
+    impl<T: Write> Internal<T> {
+        #[inline]
+        pub fn new(writer: T) -> Self {
+            Self(BufWriter::new(writer), 0)
+        }
+
+        #[inline]
+        pub const fn written_bytes(&self) -> usize {
+            self.1
+        }
+
+        #[inline]
+        pub fn write_bytes(&mut self, val: &[u8]) -> Result<()> {
+            self.1 += val.len();
+            self.0.write_all(val.as_ref())
+        }
+    }
+}

ggus/src/write/mod.rs

@@ -0,0 +1,6 @@
+mod internal;
+mod simulator;
+mod writer;
+
+pub use simulator::GGufSimulator;
+pub use writer::{GGufMetaWriter, GGufTensorWriter};