test: 支持从测例文件加载测例

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

operators/Cargo.toml

@@ -44,5 +44,9 @@ search-cuda-tools.workspace = true
 search-corex-tools.workspace = true
 
 [dev-dependencies]
+ggus = "0.4"
+memmap2 = "0.9"
+patricia_tree = "0.9"
+
 gemm = "0.18"
 rand = "0.9"

operators/src/lib.rs

@@ -2,6 +2,8 @@
 
 mod common;
 mod handle;
+#[cfg(test)]
+mod test;
 
 pub mod add;
 pub mod add_rows;

operators/src/test/gguf.rs

@@ -0,0 +1,104 @@
+use super::MetaValue;
+use digit_layout::DigitLayout;
+use ggus::{
+    GGuf, GGufError, GGufMetaDataValueType, GGufMetaError, GGufMetaMap, GGufMetaValueArray,
+    GGufReader, GENERAL_ALIGNMENT,
+};
+use patricia_tree::StringPatriciaMap;
+
+/// GGuf 文件内容，元信息和张量。
+pub(super) struct Content<'a> {
+    pub meta_kvs: StringPatriciaMap<MetaValue<'a>>,
+    pub tensors: StringPatriciaMap<GGufTensor<'a>>,
+}
+
+impl GGufMetaMap for Content<'_> {
+    fn get(&self, key: &str) -> Option<(GGufMetaDataValueType, &[u8])> {
+        self.meta_kvs.get(key).map(|v| (v.ty, &*v.value))
+    }
+}
+
+#[derive(Clone, Debug)]
+pub(super) struct GGufTensor<'a> {
+    pub ty: DigitLayout,
+    pub shape: Vec<usize>,
+    pub data: &'a [u8],
+}
+
+impl<'a> Content<'a> {
+    /// 从分片的 GGuf 文件解析内容。
+    pub fn new(files: &[&'a [u8]]) -> Result<Self, GGufError> {
+        std::thread::scope(|s| {
+            let mut ans = Self {
+                meta_kvs: Default::default(),
+                tensors: Default::default(),
+            };
+            // 在多个线程中并行解析多个文件，并逐个合并到单独的结构体中
+            for thread in files
+                .into_iter()
+                .map(|data| s.spawn(|| GGuf::new(data)))
+                .collect::<Vec<_>>()
+                .into_iter()
+            {
+                thread
+                    .join()
+                    .unwrap()
+                    .and_then(|gguf| ans.merge_file(gguf))?;
+            }
+
+            Ok(ans)
+        })
+    }
+
+    fn merge_file(&mut self, others: GGuf<'a>) -> Result<(), GGufError> {
+        // 合并元信息
+        for (k, kv) in others.meta_kvs {
+            if k == GENERAL_ALIGNMENT || k.starts_with("split.") {
+                continue;
+            }
+            let value = MetaValue {
+                ty: kv.ty(),
+                value: kv.value_bytes(),
+            };
+            if self.meta_kvs.insert(k.to_string(), value).is_some() {
+                return Err(GGufError::DuplicateMetaKey(k.into()));
+            }
+        }
+        // 合并张量，并将形状转换到 usize 类型
+        for (name, tensor) in others.tensors {
+            let tensor = tensor.to_info();
+            let tensor = GGufTensor {
+                ty: tensor.ty().to_digit_layout(),
+                shape: tensor.shape().iter().map(|&d| d as _).collect(),
+                data: &others.data[tensor.offset()..][..tensor.nbytes()],
+            };
+            if self.tensors.insert(name.to_string(), tensor).is_some() {
+                return Err(GGufError::DuplicateTensorName(name.into()));
+            }
+        }
+        Ok(())
+    }
+}
+
+impl MetaValue<'_> {
+    /// 从元信息读取 isize 数组，用于解析 strides
+    pub fn to_vec_isize(&self) -> Result<Vec<isize>, GGufMetaError> {
+        use GGufMetaDataValueType as Ty;
+
+        let mut reader = GGufReader::new(&self.value);
+        let (ty, len) = match self.ty {
+            Ty::Array => reader.read_arr_header().map_err(GGufMetaError::Read)?,
+            ty => return Err(GGufMetaError::TypeMismatch(ty)),
+        };
+
+        match ty {
+            Ty::I32 => Ok(GGufMetaValueArray::<i32>::new(reader, len)
+                .map(|x| x.unwrap() as _)
+                .collect()),
+            Ty::I64 => Ok(GGufMetaValueArray::<i64>::new(reader, len)
+                .map(|x| x.unwrap() as _)
+                .collect()),
+            _ => Err(GGufMetaError::ArrTypeMismatch(ty)),
+        }
+    }
+}

operators/src/test/mod.rs

@@ -0,0 +1,129 @@
+mod gguf;
+
+use digit_layout::DigitLayout;
+use gguf::{Content, GGufTensor};
+use ggus::{GGufMetaDataValueType, GGufMetaMap, GGufMetaMapExt};
+use memmap2::Mmap;
+use ndarray_layout::{ArrayLayout, Endian::BigEndian};
+use std::{collections::HashMap, env::var_os, fs::File};
+
+/// 测例数据。
+pub struct TestCase<'a> {
+    /// 测例在文件中的序号。
+    pub index: usize,
+    /// 测例元信息，即传递给算子非张量参数。
+    pub attributes: HashMap<String, MetaValue<'a>>,
+    /// 测例张量，包括算子的输入和正确答案。
+    pub tensors: HashMap<String, Tensor<'a>>,
+}
+
+/// 元信息键值对。
+pub struct MetaValue<'a> {
+    pub ty: GGufMetaDataValueType,
+    pub value: &'a [u8],
+}
+
+/// 测例张量。
+pub struct Tensor<'a> {
+    pub ty: DigitLayout,
+    pub layout: ArrayLayout<4>,
+    pub data: &'a [u8],
+}
+
+impl GGufMetaMap for TestCase<'_> {
+    fn get(&self, key: &str) -> Option<(GGufMetaDataValueType, &[u8])> {
+        self.attributes.get(key).map(|v| (v.ty, &*v.value))
+    }
+}
+
+impl<'a> Content<'a> {
+    pub fn into_cases(self) -> HashMap<String, Vec<TestCase<'a>>> {
+        assert_eq!(self.general_architecture().unwrap(), "infiniop-test");
+        let mut ans = HashMap::new();
+
+        let ntest = self.get_usize("test_count").unwrap();
+        let Self {
+            mut meta_kvs,
+            mut tensors,
+        } = self;
+        for i in 0..ntest {
+            let prefix = format!("test.{i}.");
+
+            let mut meta_kvs = meta_kvs
+                .split_by_prefix(&prefix)
+                .into_iter()
+                .map(|(k, v)| (k[prefix.len()..].to_string(), v))
+                .collect::<HashMap<_, _>>();
+            let tensors = tensors
+                .split_by_prefix(&prefix)
+                .into_iter()
+                .map(|(k, v)| {
+                    let GGufTensor {
+                        ty,
+                        mut shape,
+                        data,
+                    } = v;
+                    shape.reverse();
+
+                    let k = k[prefix.len()..].to_string();
+                    let element_size = ty.nbytes();
+                    let layout = if let Some(strides) = meta_kvs.remove(&format!("{k}.strides")) {
+                        let mut strides = strides.to_vec_isize().unwrap();
+                        for x in &mut strides {
+                            *x *= element_size as isize
+                        }
+                        strides.reverse();
+
+                        ArrayLayout::<4>::new(&shape, &strides, 0)
+                    } else {
+                        ArrayLayout::<4>::new_contiguous(&shape, BigEndian, element_size)
+                    };
+                    (k, Tensor { ty, layout, data })
+                })
+                .collect();
+
+            let case = TestCase {
+                index: i,
+                attributes: meta_kvs,
+                tensors,
+            };
+
+            let op_name = case.get_str("op_name").unwrap().to_string();
+            ans.entry(op_name).or_insert_with(Vec::new).push(case);
+        }
+
+        ans
+    }
+}
+
+#[test]
+fn test() {
+    let Some(name) = var_os("TEST_CASES") else {
+        eprintln!("TEST_CASES not set");
+        return;
+    };
+    let Ok(file) = File::open(&name) else {
+        eprintln!("Failed to open {}", name.to_string_lossy());
+        return;
+    };
+    let mmap = unsafe { Mmap::map(&file).unwrap() };
+    let cases = Content::new(&[&*mmap]).unwrap().into_cases();
+
+    for (op_name, cases) in cases {
+        for case in cases {
+            println!("Test case {}: {op_name}", case.index);
+            for (k, v) in &case.attributes {
+                println!("  {k}: {:?}", v.ty)
+            }
+            for (k, v) in &case.tensors {
+                println!(
+                    "  {k}: {} {:?} / {:?} {}",
+                    v.ty,
+                    v.layout.shape(),
+                    v.layout.strides(),
+                    v.data.len()
+                )
+            }
+        }
+    }
+}