Inherit energy counters

3 files changed, 67 insertions, 33 deletions

diff --git a/perf-event/src/lib.rs b/perf-event/src/lib.rs
index d1478fe..e09de18 100644
--- a/perf-event/src/lib.rs
+++ b/perf-event/src/lib.rs
@@ -533,6 +533,25 @@ impl<'a> Builder<'a> {
         self
     }
 
+    /// Set whether this counter is inherited by new threads.
+    ///
+    /// When this flag is set, this counter observes activity in new threads
+    /// created by any thread already being observed.
+    ///
+    /// By default, the flag is unset: counters are not inherited, and observe
+    /// only the threads specified when they are created.
+    ///
+    /// This flag cannot be set if the counter belongs to a `Group`. Doing so
+    /// will result in an error when the counter is built. This is a kernel
+    /// limitation.
+    pub fn inherit_thread(mut self, inherit: bool) -> Builder<'a> {
+        let flag = if inherit { 1 } else { 0 };
+        self.attrs.set_inherit(flag);
+        self.attrs.set_inherit_stat(flag);
+        self.attrs.set_inherit_thread(flag);
+        self
+    }
+
     /// Count events of the given kind. This accepts an [`Event`] value,
     /// or any type that can be converted to one, so you can pass [`Hardware`],
     /// [`Software`] and [`Cache`] values directly.
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)
     }
 }