fix(ggus): 改正张量信息读入，支持非对齐情况

refactor(xtask): 初步实现变换框架

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

Cargo.toml

@@ -1,3 +1,6 @@
 [workspace]
 members = ["ggus", "xtask"]
 resolver = "2"
+
+[workspace.dependencies]
+indexmap = "2.3"

ggus/Cargo.toml

@@ -12,4 +12,4 @@ keywords = ["gguf", "ggml", "llama-cpp"]
 categories = ["parsing", "parser-implementations", "encoding"]
 
 [dependencies]
-indexmap = "2.3"
+indexmap.workspace = true

ggus/src/tensor.rs

@@ -1,7 +1,10 @@
 use crate::{GGufReadError, GGufReader};
 use indexmap::IndexMap;
 use std::{
-    hash::Hash, marker::PhantomData, mem::size_of, slice::from_raw_parts, str::from_utf8_unchecked,
+    alloc::{alloc, dealloc, Layout},
+    hash::Hash,
+    ptr::{copy_nonoverlapping, NonNull},
+    slice::from_raw_parts,
 };
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
@@ -42,7 +45,7 @@ pub enum GGmlType {
 }
 
 impl GGmlType {
-    fn nbytes(self) -> usize {
+    pub fn nbytes(self) -> usize {
         match self {
             Self::F32 => size_of::<f32>(),
             Self::F16 => 2,
@@ -135,9 +138,15 @@ impl<'a> GGufTensors<'a> {
     }
 }
 
-// | name::ptr | name::len | offset | ggml_type;ndim | shape ..
-#[repr(transparent)]
-pub struct GGufTensorInfo<'a>(Box<[u64]>, PhantomData<&'a ()>);
+pub struct GGufTensorInfo<'a> {
+    name: &'a str,
+    ty: GGmlType,
+    ndim: u32,
+    offset: usize,
+    shape: NonNull<u64>,
+}
+
+unsafe impl Sync for GGufTensorInfo<'_> {}
 
 impl PartialEq for GGufTensorInfo<'_> {
     #[inline]
@@ -155,60 +164,65 @@ impl Hash for GGufTensorInfo<'_> {
     }
 }
 
+impl Drop for GGufTensorInfo<'_> {
+    #[inline]
+    fn drop(&mut self) {
+        unsafe {
+            dealloc(
+                self.shape.as_ptr().cast(),
+                Layout::array::<u64>(self.ndim as _).unwrap(),
+            )
+        }
+    }
+}
+
 impl<'a> GGufTensorInfo<'a> {
     fn new(name: &'a str) -> Self {
         unsafe {
             let ptr = name.as_ptr().add(name.len());
-            let ndim = ptr.cast::<u32>().read_unaligned() as usize;
+            let ndim = ptr.cast::<u32>().read_unaligned();
+            let layout = Layout::array::<u64>(ndim as _).unwrap();
 
             let ptr = ptr.add(size_of::<u32>());
-            let shape = ptr;
+            let shape_ = ptr;
 
-            let ptr = ptr.add(ndim * size_of::<u64>());
-            let ggml_type = ptr.cast::<GGmlType>().read_unaligned();
+            let ptr = ptr.add(layout.size());
+            let ty = ptr.cast::<GGmlType>().read_unaligned();
 
             let ptr = ptr.add(size_of::<u32>());
-            let offset = ptr.cast::<u64>().read_unaligned();
-
-            let mut body = vec![0u64; 4 + ndim].into_boxed_slice();
-            body[0] = name.as_ptr() as _;
-            body[1] = name.len() as _;
-            body[2] = offset;
-            let ptr = body[3..].as_mut_ptr().cast::<u32>();
-            ptr.cast::<GGmlType>().write(ggml_type);
-            ptr.add(1).write(ndim as _);
-            body[4..].copy_from_slice(from_raw_parts(shape.cast(), ndim));
-            Self(body, PhantomData)
-        }
-    }
+            let offset = ptr.cast::<u64>().read_unaligned() as _;
 
-    #[inline]
-    pub fn name(&self) -> &'a str {
-        unsafe {
-            let ptr = self.0.as_ptr().read();
-            let len = self.0.as_ptr().add(1).read();
-            from_utf8_unchecked(from_raw_parts(ptr as _, len as _))
+            let shape = alloc(layout).cast::<u64>();
+            copy_nonoverlapping(shape_, shape.cast(), layout.size());
+
+            Self {
+                name,
+                ty,
+                ndim,
+                offset,
+                shape: NonNull::new_unchecked(shape),
+            }
         }
     }
 
     #[inline]
-    fn ndim(&self) -> usize {
-        unsafe { self.0.as_ptr().add(3).cast::<u32>().add(1).read() as _ }
+    pub fn name(&self) -> &'a str {
+        self.name
     }
 
     #[inline]
     pub fn shape(&self) -> &[u64] {
-        unsafe { from_raw_parts(self.0.as_ptr().add(4), self.ndim()) }
+        unsafe { from_raw_parts(self.shape.as_ptr(), self.ndim as _) }
     }
 
     #[inline]
     pub fn ggml_type(&self) -> GGmlType {
-        unsafe { self.0.as_ptr().add(3).cast::<GGmlType>().read() }
+        self.ty
     }
 
     #[inline]
     pub fn offset(&self) -> usize {
-        unsafe { self.0.as_ptr().add(2).read() as _ }
+        self.offset
     }
 
     #[inline]

ggus/src/write/internal.rs

@@ -1,5 +1,5 @@
 use crate::{pad, GGmlType, GGufFileHeader, GGufMetaDataValueType, GENERAL_ALIGNMENT};
-use internal::Internal;
+use internal_::Internal;
 use std::{
     io::{Result, Write},
     slice::from_raw_parts,
@@ -39,11 +39,11 @@ impl<T: Write> GGufWriter<T> {
     }
 
     #[inline]
-    pub fn write_alignment(&mut self, align: usize) -> Result<usize> {
+    pub fn write_alignment(&mut self, alignment: usize) -> Result<usize> {
         self.write_meta_kv(
             GENERAL_ALIGNMENT,
             GGufMetaDataValueType::U32,
-            (align as u32).to_le_bytes(),
+            (alignment as u32).to_le_bytes(),
         )
         .map(Option::unwrap)
     }
@@ -75,25 +75,25 @@ impl<T: Write> GGufWriter<T> {
         &mut self,
         name: &str,
         shape: &[u64],
-        data_type: GGmlType,
+        ty: GGmlType,
         offset: u64,
     ) -> Result<()> {
         self.write_str(name)?;
         self.write(&[shape.len() as u32])?;
         self.write(shape)?;
-        self.write(&[data_type])?;
+        self.write(&[ty])?;
         self.write(&[offset])
     }
 
-    pub fn write_data(&mut self, data: &[u8], align: usize) -> Result<()> {
-        for _ in 0..pad(self.written_bytes(), align) {
+    pub fn write_data(&mut self, data: &[u8], alignment: usize) -> Result<()> {
+        for _ in 0..pad(self.written_bytes(), alignment) {
             self.write(&[0u8])?;
         }
         self.write(data)
     }
 }
 
-mod internal {
+mod internal_ {
     use std::io::{BufWriter, Result, Write};
 
     pub(super) struct Internal<T: Write>(BufWriter<T>, usize);

ggus/src/write/simulator.rs

@@ -1,5 +1,5 @@
 use super::internal::GGufWriter;
-use crate::{pad, GGufFileHeader, GGufMetaDataValueType, GGufTensorInfo, DEFAULT_ALIGNMENT};
+use crate::{pad, GGmlType, GGufFileHeader, GGufMetaDataValueType, DEFAULT_ALIGNMENT};
 use std::io::{Result, Write};
 
 pub struct GGufSimulator {
@@ -28,41 +28,31 @@ impl GGufSimulator {
     }
 
     #[inline]
-    pub fn with_alignment(align: usize) -> Self {
+    pub fn with_alignment(alignment: usize) -> Self {
         let mut ans = Self::new();
-        ans.write_alignment(align);
+        ans.write_alignment(alignment);
         ans
     }
 
     #[inline]
-    pub fn write_alignment(&mut self, align: usize) {
-        self.alignment = self.writer.write_alignment(align).unwrap();
+    pub fn write_alignment(&mut self, alignment: usize) {
+        self.alignment = self.writer.write_alignment(alignment).unwrap();
     }
 
     #[inline]
-    pub fn write_meta_kv(
-        &mut self,
-        key: impl AsRef<str>,
-        ty: GGufMetaDataValueType,
-        val: impl AsRef<[u8]>,
-    ) {
-        if let Some(align) = self.writer.write_meta_kv(key, ty, val).unwrap() {
-            self.alignment = align;
+    pub fn write_meta_kv(&mut self, key: &str, ty: GGufMetaDataValueType, val: &[u8]) {
+        if let Some(alignment) = self.writer.write_meta_kv(key, ty, val).unwrap() {
+            self.alignment = alignment;
         }
     }
 
-    pub fn write_tensor(&mut self, info: &GGufTensorInfo) {
+    pub fn write_tensor(&mut self, name: &str, ty: GGmlType, shape: &[u64]) {
         self.offset += pad(self.offset, self.alignment);
         self.writer
-            .write_tensor_info(
-                info.name(),
-                info.shape(),
-                info.ggml_type(),
-                self.offset as _,
-            )
+            .write_tensor_info(name, shape, ty, self.offset as _)
             .unwrap();
 
-        let len = info.nbytes();
+        let len = shape.iter().product::<u64>() as usize * ty.nbytes();
         self.offset += len;
         self.data.push(len);
     }