feat(xtask): convert 支持合并权重矩阵

Signed-off-by: YdrMaster <ydrml@hotmail.com>
InfiniTensor · Aug 15, 2024 · 80a6330 · 80a6330
1 parent 44ee36a
commit 80a6330
Show file tree

Hide file tree

Showing 6 changed files with 419 additions and 51 deletions.
diff --git a/xtask/src/convert.rs b/xtask/src/convert.rs
@@ -46,6 +46,8 @@ impl ConvertArgs {
                     Operator::cast(content)
                 } else if let Some(content) = op.strip_prefix("transpose:") {
                     Operator::transpose_linear(content)
+                } else if let Some(content) = op.strip_prefix("merge:") {
+                    Operator::merge_linear(content)
                 } else {
                     panic!("Unsupported operation: {}", op)
                 }

diff --git a/xtask/src/utils/mod.rs b/xtask/src/utils/mod.rs
@@ -87,6 +87,12 @@ enum DataPromise<'a> {
     Lazy(Arc<dyn LazyData + Send + Sync + 'a>),
 }
 
+impl<'a> DataPromise<'a> {
+    fn lazy(f: impl FnOnce() -> MmapMut + Send + Sync + 'a) -> Self {
+        Self::Lazy(Arc::new(LazyLock::new(f)))
+    }
+}
+
 impl ggus::DataFuture for DataPromise<'_> {
     #[inline]
     fn get(&self) -> &[u8] {

diff --git a/xtask/src/utils/operator/cast.rs b/xtask/src/utils/operator/cast.rs
@@ -6,7 +6,6 @@ use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterato
 use std::{
     alloc::Layout,
     slice::{from_raw_parts, from_raw_parts_mut},
-    sync::{Arc, LazyLock},
 };
 
 impl Operator {
@@ -31,15 +30,15 @@ impl Content<'_> {
                 let data = tensor.data.clone();
 
                 tensor.ty = ty;
-                tensor.data = DataPromise::Lazy(Arc::new(LazyLock::new(move || {
+                tensor.data = DataPromise::lazy(move || {
                     use GGmlType as Ty;
                     let data = data.get();
                     match (from, to) {
                         (Ty::F32, Ty::F16) => cast(data, |&x| f16::from_f32(x)),
                         (Ty::F16, Ty::F32) => cast(data, |&x| f16::to_f32(x)),
                         (_, _) => todo!("unsupported cast: {from:?} -> {to:?}"),
                     }
-                })));
+                });
             }
         }