Inherit energy counters

2 files changed, 48 insertions, 33 deletions

diff --git a/src/energy.rs b/src/energy.rs
index b9cfced..66b87cf 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -190,6 +190,9 @@ impl EnergyService {
     fn handle_request(&mut self, request: Request) {
         match request {
             Request::NewTask(pid, task_info) => {
+                let Ok(process) = procfs::process::Process::new(pid) else {
+                    return;
+                };
                 if let Some(info) = self.process_info.write().unwrap().get_mut(&pid) {
                     let old_budget = task_info.read_budget();
                     let old_time = task_info.read_time_since_last_schedule_raw();
@@ -198,17 +201,20 @@ impl EnergyService {
                     info.task_info.set_last_scheduled_raw(old_time);
                     return;
                 }
-                if self
-                    .estimator
-                    .start_trace(
-                        pid as u64,
-                        task_info.read_cpu(),
-                        task_info.is_running_on_e_core(),
-                    )
-                    .is_err()
+                let main_thread = process.task_main_thread().unwrap().pid;
+                if !self.process_info.read().unwrap().contains_key(&main_thread)
+                    && self
+                        .estimator
+                        .start_trace(
+                            pid as u64,
+                            task_info.read_cpu(),
+                            task_info.is_running_on_e_core(),
+                        )
+                        .is_err()
                 {
                     return;
                 }
+
                 let parent = (|| {
                     let process = procfs::process::Process::new(pid)?;
                     process.stat().map(|stat| stat.ppid)
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
index 7223b12..96c73d9 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use burn::tensor::Tensor;
 use perf_event::{
-    events::{Event, Hardware},
+    events::{Event, Hardware, Software},
     Builder, Counter, Group,
 };
 
@@ -60,7 +60,7 @@ static EVENT_TYPES_E: &[Event] = &[
     Event::Hardware(Hardware::CACHE_REFERENCES),
     Event::Hardware(Hardware::CPU_CYCLES),
     Event::Hardware(Hardware::INSTRUCTIONS),
-    Event::Hardware(Hardware::REF_CPU_CYCLES),
+    // Event::Hardware(Hardware::REF_CPU_CYCLES),
 ];
 
 impl Estimator for PerfEstimator {
@@ -85,14 +85,15 @@ impl Estimator for PerfEstimator {
         .iter()
         .map(|kind| {
             Builder::new()
-                .group(&mut group)
+                // .group(&mut group)
                 .kind(kind.clone())
                 .observe_pid(pid as i32)
+                .inherit_thread(true)
                 .build()
         })
         .collect();
 
-        let counters = match counters {
+        let mut counters = match counters {
             Ok(counters) => counters,
             Err(e) => {
                 eprintln!(
@@ -103,16 +104,18 @@ impl Estimator for PerfEstimator {
             }
         };
 
-        if let Err(e) = group.enable() {
-            eprintln!("Failed to enable performance counters: {}", e);
-            return Err(());
-        }
-        if let Err(e) = group.reset() {
-            eprintln!("Failed to reset performance counters: {}", e);
-            return Err(());
+        for counter in counters.iter_mut() {
+            if let Err(e) = counter.enable() {
+                eprintln!("Failed to enable performance counters: {}", e);
+                return Err(());
+            }
+            if let Err(e) = counter.reset() {
+                eprintln!("Failed to reset performance counters: {}", e);
+                return Err(());
+            }
         }
 
-        let old_time = group.read().unwrap().time_running();
+        let old_time = counters[0].read_count_and_time().unwrap().time_running;
         let counters = Counters {
             counters,
             group,
@@ -151,17 +154,18 @@ impl Estimator for PerfEstimator {
             return None;
         };
 
-        let counts = match counters.group.read() {
-            Ok(counts) => counts,
-            Err(e) => {
-                println!("failed to read group: {e}");
-                return None;
-            }
-        };
-
-        let time_running_ns = counts.time_running();
-        if time_running_ns - counters.old_time == 0 || counts.iter().next().unwrap().1 == &0 {
-            //println!("The counters are zero although the task has been scheduled!!");
+        // let counts = match counters.group.read() {
+        //     Ok(counts) => counts,
+        //     Err(e) => {
+        //         println!("failed to read group: {e}");
+        //         return None;
+        //     }
+        // };
+        let time_running_ns = counters.counters[0]
+            .read_count_and_time()
+            .unwrap()
+            .time_running;
+        if time_running_ns - counters.old_time == 0 {
             return None;
         }
         let correction_factor = 10_000_000. / (time_running_ns - counters.old_time) as f64;
@@ -172,10 +176,12 @@ impl Estimator for PerfEstimator {
             (self.shared_cpu_current_frequencies.read().unwrap()[counters.cpu as usize] / 1000)
                 as f64,
         ];
-        for ty in counters.counters.iter() {
-            let count: u64 = counts[&ty];
+        for ty in counters.counters.iter_mut() {
+            // let count: u64 = counts[ty];
+            let count: u64 = ty.read().unwrap();
             values.push((count as f64) * correction_factor);
         }
+        values.push(values[4]);
 
         let result = if counters.running_on_e_core {
             &self.model_e
@@ -190,6 +196,9 @@ impl Estimator for PerfEstimator {
         }
         counters.old_total_energy += energy / correction_factor;
         counters.group.reset().unwrap();
+        for counter in counters.counters.iter_mut() {
+            counter.reset().unwrap();
+        }
         Some(energy / correction_factor)
     }
 }