From eca801deb10dea25bc781b4cc7a7313bf0121cc3 Mon Sep 17 00:00:00 2001
From: Vadim Smirnov <ukint-vs@proton.me>
Date: Thu, 27 Jul 2023 20:58:04 +0300
Subject: [PATCH] feat(weights-diff): Add memory weights

---
 utils/weight-diff/src/main.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/utils/weight-diff/src/main.rs b/utils/weight-diff/src/main.rs
index 22dcdda5abf..f99233311d9 100644
--- 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();