Add migration message and e core handling in the estimator

4 files changed, 53 insertions, 40 deletions

diff --git a/src/energy/estimator.rs b/src/energy/estimator.rs
index 85cb0eb..7964b00 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);
+    fn start_trace(&mut self, pid: u64, running_on_e_core: bool);
     fn stop_trace(&mut self, pid: u64);
     fn read_consumption(&mut self, pid: u64, cpu: i32) -> Option<f64>;
 }
diff --git a/src/energy/trackers/kernel.rs b/src/energy/trackers/kernel.rs
index 89a0d89..f7e0811 100644
--- a/src/energy/trackers/kernel.rs
+++ b/src/energy/trackers/kernel.rs
@@ -30,7 +30,7 @@ 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) {
+    fn start_trace(&mut self, pid: u64, running_on_e_core: bool) {
         let _ = START_TRACE.ioctl(&mut self.file, &pid);
     }
 
diff --git a/src/energy/trackers/mock.rs b/src/energy/trackers/mock.rs
index 597cab0..b2a504e 100644
--- a/src/energy/trackers/mock.rs
+++ b/src/energy/trackers/mock.rs
@@ -3,7 +3,7 @@ use crate::energy::estimator::Estimator;
 pub struct MockEstimator;
 
 impl Estimator for MockEstimator {
-    fn start_trace(&mut self, _pid: u64) {}
+    fn start_trace(&mut self, _pid: u64, _running_on_e_core: bool) {}
 
     fn stop_trace(&mut self, _pid: u64) {}
 
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
index 71e87d8..8da2bf8 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -9,22 +9,25 @@ use perf_event::{
 use crate::freq::FrequencyKHZ;
 use std::sync::{Arc, RwLock};
 
-use crate::energy::{rapl, Estimator};
+use crate::energy::Estimator;
 use crate::model::ArrayBackend;
 
 pub struct PerfEstimator {
     registry: HashMap<u64, Counters>,
-    model: crate::model::Net<ArrayBackend>,
+    model_p: crate::model::Net<ArrayBackend>,
+    model_e: crate::model::Net<ArrayBackend>,
     device: <ArrayBackend as burn::prelude::Backend>::Device,
     shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
 }
 
 impl PerfEstimator {
     pub fn new(shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>) -> Self {
-        let model = crate::model::load_model();
+        let model_p = crate::model::load_model_p();
+        let model_e = crate::model::load_model_e();
         Self {
             registry: Default::default(),
-            model,
+            model_p,
+            model_e,
             device: Default::default(),
             shared_cpu_current_frequencies,
         }
@@ -36,9 +39,21 @@ struct Counters {
     counters: Vec<Counter>,
     old_time: u64,
     old_total_energy: f64,
+    running_on_e_core: bool,
 }
 
-static EVENT_TYPES: &[Event] = &[
+static EVENT_TYPES_P: &[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),
+];
+//TODO: use correct counter
+static EVENT_TYPES_E: &[Event] = &[
     Event::Hardware(Hardware::BRANCH_INSTRUCTIONS),
     Event::Hardware(Hardware::BRANCH_MISSES),
     Event::Hardware(Hardware::CACHE_MISSES),
@@ -50,8 +65,7 @@ static EVENT_TYPES: &[Event] = &[
 ];
 
 impl Estimator for PerfEstimator {
-    fn start_trace(&mut self, pid: u64) {
-        //let Ok(mut group) = Group::new_with_pid_and_cpu(-1, 0) else {
+    fn start_trace(&mut self, pid: u64, running_on_e_core: bool) {
         let mut group = match Group::new_with_pid_and_cpu(pid as i32, -1) {
             Ok(counters) => counters,
             Err(e) => {
@@ -62,23 +76,21 @@ impl Estimator for PerfEstimator {
                 return;
             }
         };
-        //let Ok(mut group) = Group::new_with_pid_and_cpu(pid as i32, -1) else {
-        //    eprintln!("Failed to create performance counter group for PID {}", pid);
-        //    return;
-        //};
 
-        let counters: Result<Vec<_>, _> = EVENT_TYPES
-            .iter()
-            .map(|kind| {
-                Builder::new()
-                    .group(&mut group)
-                    .kind(kind.clone())
-                    .observe_pid(pid as i32)
-                    //.observe_pid(-1)
-                    //.one_cpu(0)
-                    .build()
-            })
-            .collect();
+        let counters: Result<Vec<_>, _> = if running_on_e_core {
+            EVENT_TYPES_E
+        } else {
+            EVENT_TYPES_P
+        }
+        .iter()
+        .map(|kind| {
+            Builder::new()
+                .group(&mut group)
+                .kind(kind.clone())
+                .observe_pid(pid as i32)
+                .build()
+        })
+        .collect();
 
         let counters = match counters {
             Ok(counters) => counters,
@@ -95,13 +107,17 @@ impl Estimator for PerfEstimator {
             eprintln!("Failed to enable performance counters: {}", e);
             return;
         }
-        group.reset();
+        if let Err(e) = group.reset() {
+            eprintln!("Failed to reset performance counters: {}", e);
+            return;
+        }
 
         let counters = Counters {
             counters,
             group,
             old_time: 0,
             old_total_energy: 0.,
+            running_on_e_core,
         };
         self.registry.insert(pid, counters);
     }
@@ -122,36 +138,33 @@ impl Estimator for PerfEstimator {
                 return None;
             }
         };
-        //let task_clock = counts[&counters.counters[7]];
-        let task_clock = counts[&counters.counters[7]];
+        let num_counter = counters.counters.len();
+        let task_clock = counts[&counters.counters[num_counter - 1]];
 
         if task_clock == 0 {
             return None;
         }
-        let current_time = std::time::Instant::now();
         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;
-        //dbg!(time_running_ns);
-        //dbg!(counters.old_time.elapsed().as_millis());
 
         let mut values =
             vec![(self.shared_cpu_current_frequencies.read().unwrap()[cpu as usize] / 1000) as f64];
-        //let mut values = vec![crate::benchmark::read_cpu_frequency(0).unwrap()];
-        for ty in counters.counters.iter().take(7) {
+        for ty in counters.counters.iter().take(num_counter - 1) {
             let count: u64 = counts[&ty];
             values.push((count as f64) * correction_factor);
         }
-        //dbg!(&values);
-        let result = self
-            .model
-            .forward(Tensor::from_floats(&values.as_slice()[0..8], &self.device));
+
+        let result = if counters.running_on_e_core {
+            &self.model_e
+        } else {
+            &self.model_p
+        }
+        .forward(Tensor::from_floats(&values.as_slice()[0..], &self.device));
 
         let energy = result.into_scalar() as f64;
-        //dbg!(energy, correction_factor);
         counters.old_total_energy += energy / correction_factor;
         counters.group.reset().unwrap();
-        //dbg!(counters.old_total_energy);
         Some(energy / correction_factor)
     }
 }