Integrate energy model into perf estimator

4 files changed, 92 insertions, 27 deletions

diff --git a/src/energy/trackers.rs b/src/energy/trackers.rs
index 31e2f3e..1665177 100644
--- a/src/energy/trackers.rs
+++ b/src/energy/trackers.rs
@@ -10,5 +10,5 @@ pub use perf::*;
 pub trait Estimator: Send + 'static {
     fn start_trace(&mut self, pid: u64);
     fn stop_trace(&mut self, pid: u64);
-    fn read_consumption(&mut self, pid: u64) -> u64;
+    fn read_consumption(&mut self, pid: u64) -> Option<f64>;
 }
diff --git a/src/energy/trackers/kernel.rs b/src/energy/trackers/kernel.rs
index b71fee6..2a19e4c 100644
--- a/src/energy/trackers/kernel.rs
+++ b/src/energy/trackers/kernel.rs
@@ -38,12 +38,12 @@ impl Estimator for KernelDriver {
         let _ = STOP_TRACE.ioctl(&mut self.file, &pid);
     }
 
-    fn read_consumption(&mut self, pid: u64) -> u64 {
+    fn read_consumption(&mut self, pid: u64) -> Option<f64> {
         let mut arg = pid;
         if READ_POWER.ioctl(&mut self.file, &mut arg).is_ok() {
-            arg
+            Some(arg as f64)
         } else {
-            0
+            None
         }
     }
 }
diff --git a/src/energy/trackers/mock.rs b/src/energy/trackers/mock.rs
index eb23421..e3ef377 100644
--- a/src/energy/trackers/mock.rs
+++ b/src/energy/trackers/mock.rs
@@ -8,7 +8,7 @@ impl Estimator for MockEstimator {
 
     fn stop_trace(&mut self, _pid: u64) {}
 
-    fn read_consumption(&mut self, _pid: u64) -> u64 {
-        14
+    fn read_consumption(&mut self, _pid: u64) -> Option<f64> {
+        Some(14.)
     }
 }
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
index 17bc693..d340e1f 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -1,43 +1,76 @@
 use std::collections::HashMap;
 
+use burn::tensor::Tensor;
 use perf_event::{
-    events::{Event, Hardware},
+    events::{Event, Hardware, Software},
     Builder, Counter, Group,
 };
 
-use crate::energy::Estimator;
+use crate::energy::{rapl, Estimator};
+type ArrayBackend = burn_ndarray::NdArray<f32>;
 
-#[derive(Default)]
 pub struct PerfEstimator {
     registry: HashMap<u64, Counters>,
+    model: crate::model::Net<ArrayBackend>,
+    device: <ArrayBackend as burn::prelude::Backend>::Device,
+}
+
+impl Default for PerfEstimator {
+    fn default() -> Self {
+        let model = crate::model::load_model();
+        Self {
+            registry: Default::default(),
+            model,
+            device: Default::default(),
+        }
+    }
 }
 
 struct Counters {
     group: Group,
     counters: Vec<Counter>,
+    old_time: u64,
+    old_total_energy: f64,
 }
 
-static EVENT_TYPES: &[(f32, Event)] = &[
-    (1.0, Event::Hardware(Hardware::CPU_CYCLES)),
-    (2.0, Event::Hardware(Hardware::INSTRUCTIONS)),
+static EVENT_TYPES: &[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) {
-        let Ok(mut group) = Group::new_with_pid_and_cpu(-1, 0) else {
-            eprintln!("Failed to create performance counter group for PID {}", pid);
-            return;
+        //let Ok(mut group) = Group::new_with_pid_and_cpu(-1, 0) else {
+        let mut group = match Group::new_with_pid_and_cpu(pid as i32, -1) {
+            Ok(counters) => counters,
+            Err(e) => {
+                eprintln!(
+                    "Failed to create performance counter group for PID {}: {}",
+                    pid, e
+                );
+                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)| {
+            .map(|kind| {
                 Builder::new()
                     .group(&mut group)
                     .kind(kind.clone())
-                    // .observe_pid(pid as i32)
-                    .observe_pid(-1)
-                    .one_cpu(0)
+                    .observe_pid(pid as i32)
+                    //.observe_pid(-1)
+                    //.one_cpu(0)
                     .build()
             })
             .collect();
@@ -57,8 +90,14 @@ impl Estimator for PerfEstimator {
             eprintln!("Failed to enable performance counters: {}", e);
             return;
         }
+        group.reset();
 
-        let counters = Counters { counters, group };
+        let counters = Counters {
+            counters,
+            group,
+            old_time: 0,
+            old_total_energy: 0.,
+        };
         self.registry.insert(pid, counters);
     }
 
@@ -66,21 +105,47 @@ impl Estimator for PerfEstimator {
         self.registry.remove(&pid);
     }
 
-    fn read_consumption(&mut self, pid: u64) -> u64 {
+    fn read_consumption(&mut self, pid: u64) -> Option<f64> {
         let Some(counters) = self.registry.get_mut(&pid) else {
-            return 0;
+            return None;
         };
 
         let counts = match counters.group.read() {
             Ok(counts) => counts,
-            Err(_) => return 0,
+            Err(e) => {
+                println!("failed to read group: {e}");
+                return None;
+            }
         };
+        //let task_clock = counts[&counters.counters[7]];
+        let task_clock = counts[&counters.counters[7]];
+
+        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 sum = 0;
-        for ((factor, _ty), count) in EVENT_TYPES.iter().zip(counts.iter()) {
-            sum += *factor as u64 * count.1;
+        // TODO:: Respect per core frequency for the running task
+        let mut values = vec![crate::benchmark::read_cpu_frequency(0).unwrap()];
+        for ty in counters.counters.iter().take(7) {
+            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));
 
-        sum
+        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)
     }
 }