Implement energy model for e cores

4 files changed, 32 insertions, 24 deletions

diff --git a/src/energy/estimator.rs b/src/energy/estimator.rs
index fbab744..03034db 100644
--- a/src/energy/estimator.rs
+++ b/src/energy/estimator.rs
@@ -1,5 +1,5 @@
 pub trait Estimator: Send + 'static {
-    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool);
+    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool) -> Result<(), ()>;
     fn stop_trace(&mut self, pid: u64);
     fn update_information(&mut self, pid: u64, cpu: i32, is_ecore: bool);
     fn read_consumption(&mut self, pid: u64) -> Option<f64>;
diff --git a/src/energy/trackers/kernel.rs b/src/energy/trackers/kernel.rs
index f42bb16..a0178a4 100644
--- a/src/energy/trackers/kernel.rs
+++ b/src/energy/trackers/kernel.rs
@@ -30,8 +30,9 @@ const STOP_TRACE: Ioctl<Write, &u64> = unsafe { PERF_MON.write(0x81) };
 const READ_POWER: Ioctl<WriteRead, &u64> = unsafe { PERF_MON.write_read(0x82) };
 
 impl Estimator for KernelDriver {
-    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool) {
+    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool) -> Result<(), ()> {
         let _ = START_TRACE.ioctl(&mut self.file, &pid);
+        Ok(())
     }
 
     fn stop_trace(&mut self, pid: u64) {
diff --git a/src/energy/trackers/mock.rs b/src/energy/trackers/mock.rs
index d9ede50..cd08c34 100644
--- a/src/energy/trackers/mock.rs
+++ b/src/energy/trackers/mock.rs
@@ -3,7 +3,9 @@ use crate::energy::estimator::Estimator;
 pub struct MockEstimator;
 
 impl Estimator for MockEstimator {
-    fn start_trace(&mut self, _pid: u64, _cpu: i32, _running_on_e_core: bool) {}
+    fn start_trace(&mut self, _pid: u64, _cpu: i32, _running_on_e_core: bool) -> Result<(), ()> {
+        Ok(())
+    }
 
     fn stop_trace(&mut self, _pid: u64) {}
 
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
index e59057d..38cefe9 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -22,8 +22,10 @@ pub struct PerfEstimator {
 
 impl PerfEstimator {
     pub fn new(shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>) -> Self {
-        let model_p = crate::model::load_model_p();
-        let model_e = crate::model::load_model_e();
+        // let model_p = crate::model::load_model("perf_pcore.pt");
+        let model_p = crate::model::load_model("perf.pt");
+        let model_e = crate::model::load_model("perf.pt");
+        // let model_e = crate::model::load_model("perf_ecore.pt");
         Self {
             registry: Default::default(),
             model_p,
@@ -51,22 +53,18 @@ static EVENT_TYPES_P: &[Event] = &[
     Event::Hardware(Hardware::CPU_CYCLES),
     Event::Hardware(Hardware::INSTRUCTIONS),
     Event::Hardware(Hardware::REF_CPU_CYCLES),
-    Event::Software(Software::TASK_CLOCK),
 ];
-//TODO: use correct counter
 static EVENT_TYPES_E: &[Event] = &[
-    Event::Hardware(Hardware::BRANCH_INSTRUCTIONS),
     Event::Hardware(Hardware::BRANCH_MISSES),
     Event::Hardware(Hardware::CACHE_MISSES),
     Event::Hardware(Hardware::CACHE_REFERENCES),
     Event::Hardware(Hardware::CPU_CYCLES),
     Event::Hardware(Hardware::INSTRUCTIONS),
     Event::Hardware(Hardware::REF_CPU_CYCLES),
-    Event::Software(Software::TASK_CLOCK),
 ];
 
 impl Estimator for PerfEstimator {
-    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool) {
+    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool) -> Result<(), ()> {
         let mut group = match Group::new_with_pid_and_cpu(pid as i32, -1) {
             Ok(counters) => counters,
             Err(e) => {
@@ -74,11 +72,12 @@ impl Estimator for PerfEstimator {
                     "Failed to create performance counter group for PID {}: {}",
                     pid, e
                 );
-                return;
+                return Err(());
             }
         };
 
-        let counters: Result<Vec<_>, _> = if running_on_e_core {
+        let counters: Result<Vec<_>, _> = if running_on_e_core || true {
+            // println!("starting e core counter");
             EVENT_TYPES_E
         } else {
             EVENT_TYPES_P
@@ -100,28 +99,30 @@ impl Estimator for PerfEstimator {
                     "Failed to create performance counter group for PID {}: {}",
                     pid, e
                 );
-                return;
+                return Err(());
             }
         };
 
         if let Err(e) = group.enable() {
             eprintln!("Failed to enable performance counters: {}", e);
-            return;
+            return Err(());
         }
         if let Err(e) = group.reset() {
             eprintln!("Failed to reset performance counters: {}", e);
-            return;
+            return Err(());
         }
 
+        let old_time = group.read().unwrap().time_running();
         let counters = Counters {
             counters,
             group,
-            old_time: 0,
+            old_time,
             old_total_energy: 0.,
             cpu,
             running_on_e_core,
         };
         self.registry.insert(pid, counters);
+        Ok(())
     }
 
     fn stop_trace(&mut self, pid: u64) {
@@ -132,18 +133,21 @@ impl Estimator for PerfEstimator {
         let mut core_type_changed = false;
         if let Some(info) = self.registry.get_mut(&pid) {
             info.cpu = cpu;
+            info.running_on_e_core = is_ecore;
             core_type_changed = is_ecore != info.running_on_e_core;
         } else {
             eprintln!("Tried to update an unknown task")
         }
         if core_type_changed {
-            self.stop_trace(pid);
-            self.stop_trace(pid);
+            // println!("migrating task to {}", cpu);
+            // self.stop_trace(pid);
+            // self.start_trace(pid, cpu, is_ecore);
         }
     }
 
     fn read_consumption(&mut self, pid: u64) -> Option<f64> {
         let Some(counters) = self.registry.get_mut(&pid) else {
+            println!("did not find counters for {pid}");
             return None;
         };
 
@@ -154,21 +158,19 @@ impl Estimator for PerfEstimator {
                 return None;
             }
         };
-        let num_counter = counters.counters.len();
-        let task_clock = counts[&counters.counters[num_counter - 1]];
-
-        if task_clock == 0 {
+        let time_running_ns = counts.time_running();
+        if time_running_ns - counters.old_time == 0 || counts.iter().next().unwrap().1 == &0 {
             return None;
         }
-        let time_running_ns = counts.time_running();
         let correction_factor = 10_000_000. / (time_running_ns - counters.old_time) as f64;
         counters.old_time = time_running_ns;
 
         let mut values = vec![
+            //if counters.running_on_e_core { 1. } else { 0. },
             (self.shared_cpu_current_frequencies.read().unwrap()[counters.cpu as usize] / 1000)
                 as f64,
         ];
-        for ty in counters.counters.iter().take(num_counter - 1) {
+        for ty in counters.counters.iter() {
             let count: u64 = counts[&ty];
             values.push((count as f64) * correction_factor);
         }
@@ -181,6 +183,9 @@ impl Estimator for PerfEstimator {
         .forward(Tensor::from_floats(&values.as_slice()[0..], &self.device));
 
         let energy = result.into_scalar() as f64;
+        if counters.running_on_e_core {
+            // dbg!(energy);
+        }
         counters.old_total_energy += energy / correction_factor;
         counters.group.reset().unwrap();
         Some(energy / correction_factor)