Merge pull request #52 from MollySophia/deserialize

Reduce deserialization memory footprint on iOS devices

Author: cryscan

crates/web-rwkv-derive/src/serde/de.rs

@@ -542,9 +542,9 @@ fn deserialize_seq(
         .filter(|field| !field.attrs.skip_deserializing())
         .count();
     let expecting = if deserialized_count == 1 {
-        format!("{} with 1 element", expecting)
+        format!("{expecting} with 1 element")
     } else {
-        format!("{} with {} elements", expecting, deserialized_count)
+        format!("{expecting} with {deserialized_count} elements")
     };
     let expecting = cattrs.expecting().unwrap_or(&expecting);
 
@@ -1126,7 +1126,7 @@ fn deserialize_adjacently_tagged_enum(
         .collect();
 
     let rust_name = params.type_name();
-    let expecting = format!("adjacently tagged enum {}", rust_name);
+    let expecting = format!("adjacently tagged enum {rust_name}");
     let expecting = cattrs.expecting().unwrap_or(&expecting);
     let type_name = cattrs.name().deserialize_name();
     let deny_unknown_fields = cattrs.deny_unknown_fields();
@@ -2377,7 +2377,7 @@ fn deserialize_map(
 }
 
 fn field_i(i: usize) -> Ident {
-    Ident::new(&format!("__field{}", i), Span::call_site())
+    Ident::new(&format!("__field{i}"), Span::call_site())
 }
 
 /// This function wraps the expression in `#[serde(deserialize_with = "...")]`

crates/web-rwkv-derive/src/serde/internals/attr.rs

@@ -556,7 +556,7 @@ impl Container {
                 } else {
                     let path = meta.path.to_token_stream().to_string().replace(' ', "");
                     return Err(
-                        meta.error(format_args!("unknown serde container attribute `{}`", path))
+                        meta.error(format_args!("unknown serde container attribute `{path}`"))
                     );
                 }
                 Ok(())
@@ -974,7 +974,7 @@ impl Variant {
                 } else {
                     let path = meta.path.to_token_stream().to_string().replace(' ', "");
                     return Err(
-                        meta.error(format_args!("unknown serde variant attribute `{}`", path))
+                        meta.error(format_args!("unknown serde variant attribute `{path}`"))
                     );
                 }
                 Ok(())
@@ -1129,8 +1129,7 @@ impl Field {
                 if let Some(lifetimes) = &borrow_attribute.lifetimes {
                     for lifetime in lifetimes {
                         if !borrowable.contains(lifetime) {
-                            let msg =
-                                format!("field `{}` does not have lifetime {}", ident, lifetime);
+                            let msg = format!("field `{ident}` does not have lifetime {lifetime}");
                             cx.error_spanned_by(field, msg);
                         }
                     }
@@ -1232,8 +1231,7 @@ impl Field {
                             for lifetime in &lifetimes {
                                 if !borrowable.contains(lifetime) {
                                     let msg = format!(
-                                        "field `{}` does not have lifetime {}",
-                                        ident, lifetime,
+                                        "field `{ident}` does not have lifetime {lifetime}",
                                     );
                                     cx.error_spanned_by(field, msg);
                                 }
@@ -1256,9 +1254,7 @@ impl Field {
                     flatten.set_true(&meta.path);
                 } else {
                     let path = meta.path.to_token_stream().to_string().replace(' ', "");
-                    return Err(
-                        meta.error(format_args!("unknown serde field attribute `{}`", path))
-                    );
+                    return Err(meta.error(format_args!("unknown serde field attribute `{path}`")));
                 }
                 Ok(())
             }) {
@@ -1451,8 +1447,7 @@ where
                 }
             } else {
                 return Err(meta.error(format_args!(
-                    "malformed {0} attribute, expected `{0}(serialize = ..., deserialize = ...)`",
-                    attr_name,
+                    "malformed {attr_name} attribute, expected `{attr_name}(serialize = ..., deserialize = ...)`",
                 )));
             }
             Ok(())
@@ -1517,16 +1512,15 @@ fn get_lit_str2(
         if !suffix.is_empty() {
             cx.error_spanned_by(
                 lit,
-                format!("unexpected suffix `{}` on string literal", suffix),
+                format!("unexpected suffix `{suffix}` on string literal"),
             );
         }
         Ok(Some(lit.clone()))
     } else {
         cx.error_spanned_by(
             expr,
             format!(
-                "expected serde {} attribute to be a string: `{} = \"...\"`",
-                attr_name, meta_item_name
+                "expected serde {attr_name} attribute to be a string: `{meta_item_name} = \"...\"`"
             ),
         );
         Ok(None)
@@ -1637,10 +1631,7 @@ fn parse_lit_into_lifetimes(
         while !input.is_empty() {
             let lifetime: Lifetime = input.parse()?;
             if !set.insert(lifetime.clone()) {
-                cx.error_spanned_by(
-                    &string,
-                    format!("duplicate borrowed lifetime `{}`", lifetime),
-                );
+                cx.error_spanned_by(&string, format!("duplicate borrowed lifetime `{lifetime}`"));
             }
             if input.is_empty() {
                 break;
@@ -1813,7 +1804,7 @@ fn borrowable_lifetimes(
     let mut lifetimes = BTreeSet::new();
     collect_lifetimes(&field.ty, &mut lifetimes);
     if lifetimes.is_empty() {
-        let msg = format!("field `{}` has no lifetimes to borrow", name);
+        let msg = format!("field `{name}` has no lifetimes to borrow");
         cx.error_spanned_by(field, msg);
         Err(())
     } else {

crates/web-rwkv-derive/src/serde/internals/check.rs

@@ -38,7 +38,7 @@ fn check_default_on_tuple(cx: &Ctxt, cont: &Container) {
                     if let Some(first) = first_default_index {
                         cx.error_spanned_by(
                             field.ty,
-                            format!("field must have #[serde(default)] because previous field {} has #[serde(default)]", first),
+                            format!("field must have #[serde(default)] because previous field {first} has #[serde(default)]"),
                         );
                     }
                     continue;
@@ -311,7 +311,7 @@ fn check_internal_tag_field_name_conflict(cx: &Ctxt, cont: &Container) {
     let diagnose_conflict = || {
         cx.error_spanned_by(
             cont.original,
-            format!("variant field name `{}` conflicts with internal tag", tag),
+            format!("variant field name `{tag}` conflicts with internal tag"),
         );
     };
 
@@ -358,10 +358,7 @@ fn check_adjacent_tag_conflict(cx: &Ctxt, cont: &Container) {
     if type_tag == content_tag {
         cx.error_spanned_by(
             cont.original,
-            format!(
-                "enum tags `{}` for type and content conflict with each other",
-                type_tag
-            ),
+            format!("enum tags `{type_tag}` for type and content conflict with each other"),
         );
     }
 }
@@ -447,7 +444,7 @@ fn check_transparent(cx: &Ctxt, cont: &mut Container, derive: Derive) {
 
 fn member_message(member: &Member) -> String {
     match member {
-        Member::Named(ident) => format!("`{}`", ident),
+        Member::Named(ident) => format!("`{ident}`"),
         Member::Unnamed(i) => format!("#{}", i.index),
     }
 }

src/context.rs

@@ -155,7 +155,7 @@ impl ContextBuilder {
                     let data = read_back_buffer(&device, &buffer);
                     let _ = sender.send(data);
                 }
-                log::info!("context dropped: {}", id);
+                log::info!("context dropped: {id}");
             });
         }
 

src/runtime/mod.rs

@@ -100,7 +100,7 @@ where
         let handle = tokio::spawn(Self::run(bundle.into(), receiver));
         tokio::spawn(async move {
             if let Err(err) = handle.await {
-                log::error!("{}", err);
+                log::error!("{err}");
             }
         });
         Self(sender)

src/tensor/ops.rs

@@ -168,25 +168,25 @@ impl std::fmt::Display for Activation {
 impl Macros {
     /// Define a `u32` macro `NF4_BLOCK_SIZE`.
     pub fn nf4(mut self, block_size: u32) -> Self {
-        self.insert("NF4_BLOCK_SIZE".into(), format!("{}u", block_size));
+        self.insert("NF4_BLOCK_SIZE".into(), format!("{block_size}u"));
         self
     }
 
     /// Define a `u32` macro `NF4_BLOCK_SIZE`.
     pub fn int8(mut self, block_size: u32) -> Self {
-        self.insert("INT8_BLOCK_SIZE".into(), format!("{}u", block_size));
+        self.insert("INT8_BLOCK_SIZE".into(), format!("{block_size}u"));
         self
     }
 
     /// Define a `f32` macro with a given name.
     pub fn f32(mut self, name: impl Into<String>, value: f32) -> Self {
-        self.insert(name.into(), format!("{}", value));
+        self.insert(name.into(), format!("{value}"));
         self
     }
 
     /// Define a `usize` macro with a given name.
     pub fn u32(mut self, name: impl Into<String>, value: u32) -> Self {
-        self.insert(name.into(), format!("{}u", value));
+        self.insert(name.into(), format!("{value}u"));
         self
     }
 
@@ -254,8 +254,8 @@ fn custom_tanh(x: vec4<f32>) -> vec4<f32> {
     /// Define a macro with custom display name and prefix.
     pub fn custom(mut self, value: impl std::fmt::Display, prefix: Option<&'_ str>) -> Self {
         match prefix {
-            None => self.insert(format!("{}", value), Default::default()),
-            Some(prefix) => self.insert(format!("{}_{}", prefix, value), Default::default()),
+            None => self.insert(format!("{value}"), Default::default()),
+            Some(prefix) => self.insert(format!("{prefix}_{value}"), Default::default()),
         };
         self
     }
@@ -273,7 +273,7 @@ fn custom_tanh(x: vec4<f32>) -> vec4<f32> {
     pub fn subgroup(self, min: u32, max: u32) -> Self {
         self.u32("MIN_SUBGROUP_SIZE", min)
             .u32("MAX_SUBGROUP_SIZE", max)
-            .define(format!("SUBGROUP_SIZE_{}_{}", min, max), true)
+            .define(format!("SUBGROUP_SIZE_{min}_{max}"), true)
     }
 }
 

src/tensor/serialization.rs

@@ -144,7 +144,9 @@ impl<'de, T: Scalar + Deserialize<'de>, K: Kind> DeserializeSeed<'de>
     {
         let context = &self.context;
         let tensor: TensorBlob<'de> = Deserialize::deserialize(deserializer)?;
-        TensorGpu::from_data_u8(context, tensor.shape, &tensor.data).map_err(D::Error::custom)
+        let tensor = TensorGpu::from_data_u8(context, tensor.shape, &tensor.data);
+        context.queue.submit(None);
+        tensor.map_err(D::Error::custom)
     }
 }