feat(weights-diff): Add memory weights (#2995)

gear-tech · Jul 28, 2023 · 51d77aa · 51d77aa
1 parent f2c7aa6
commit 51d77aa
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/scripts/weight-diff.sh b/scripts/weight-diff.sh
@@ -53,3 +53,4 @@ git checkout "$current_branch"
 
 cargo run --package gear-weight-diff --release -- diff "$dump_path1" "$dump_path2" "$runtime" "instruction" $flag
 cargo run --package gear-weight-diff --release -- diff "$dump_path1" "$dump_path2" "$runtime" "host-fn" $flag
+cargo run --package gear-weight-diff --release -- diff "$dump_path1" "$dump_path2" "$runtime" "memory" $flag
diff --git a/utils/weight-diff/src/main.rs b/utils/weight-diff/src/main.rs
@@ -76,6 +76,7 @@ enum Runtime {
 enum WeightsKind {
     Instruction,
     HostFn,
+    Memory,
 }
 
 #[derive(Debug, Serialize)]
@@ -97,6 +98,7 @@ struct DeserializableDump {
 struct DeserializableSchedule {
     instruction_weights: IndexMap<String, serde_json::Value>,
     host_fn_weights: IndexMap<String, serde_json::Value>,
+    memory_weights: IndexMap<String, serde_json::Value>,
 }
 
 impl DeserializableSchedule {
@@ -127,6 +129,18 @@ impl DeserializableSchedule {
 
         map
     }
+
+    fn memory_weights(&self) -> IndexMap<String, u64> {
+        let mut map = IndexMap::new();
+
+        for (k, v) in self.memory_weights.clone() {
+            if let Ok(v) = serde_json::from_value::<Weight>(v) {
+                map.insert(k, v.ref_time());
+            }
+        }
+
+        map
+    }
 }
 
 fn format_weight(weight: u64) -> String {
@@ -214,6 +228,7 @@ fn main() {
                     schedule2.instruction_weights(),
                 ),
                 WeightsKind::HostFn => (schedule1.host_fn_weights(), schedule2.host_fn_weights()),
+                WeightsKind::Memory => (schedule1.memory_weights(), schedule2.memory_weights()),
             };
 
             let mut result_map = IndexMap::new();