Fix warnings and general cleanup

6 files changed, 33 insertions, 90 deletions

diff --git a/src/energy/budget.rs b/src/energy/budget.rs
index e29554d..b049814 100644
--- a/src/energy/budget.rs
+++ b/src/energy/budget.rs
@@ -5,7 +5,6 @@ use std::sync::Arc;
 
 type Pid = i32;
 
-const MAX_BUDGET_FACTOR: f64 = 5.0;
 const MAX_BUDGET: u64 = 30000;
 
 pub trait BudgetPolicy: Send + 'static {
diff --git a/src/energy/estimator.rs b/src/energy/estimator.rs
index 03034db..54d8c47 100644
--- a/src/energy/estimator.rs
+++ b/src/energy/estimator.rs
@@ -1,6 +1,6 @@
 pub trait Estimator: Send + 'static {
-    fn start_trace(&mut self, pid: u64, cpu: i32, running_on_e_core: bool) -> Result<(), ()>;
+    fn start_trace(&mut self, pid: u64, cpu: i32) -> Result<(), ()>;
     fn stop_trace(&mut self, pid: u64);
-    fn update_information(&mut self, pid: u64, cpu: i32, is_ecore: bool);
+    fn update_information(&mut self, pid: u64, cpu: i32);
     fn read_consumption(&mut self, pid: u64) -> Option<f64>;
 }
diff --git a/src/energy/rapl.rs b/src/energy/rapl.rs
index b88a959..a570d5c 100644
--- a/src/energy/rapl.rs
+++ b/src/energy/rapl.rs
@@ -1,3 +1,4 @@
+#![allow(unused)]
 use std::fs;
 use std::io;
 use std::path::Path;
diff --git a/src/energy/trackers/kernel.rs b/src/energy/trackers/kernel.rs
index a0178a4..8ce8469 100644
--- a/src/energy/trackers/kernel.rs
+++ b/src/energy/trackers/kernel.rs
@@ -1,4 +1,3 @@
-// energy/trackers/kernel.rs
 use iocuddle::*;
 use std::fs::File;
 use std::path::Path;
@@ -30,7 +29,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, cpu: i32, running_on_e_core: bool) -> Result<(), ()> {
+    fn start_trace(&mut self, pid: u64, _cpu: i32) -> Result<(), ()> {
         let _ = START_TRACE.ioctl(&mut self.file, &pid);
         Ok(())
     }
@@ -39,9 +38,7 @@ impl Estimator for KernelDriver {
         let _ = STOP_TRACE.ioctl(&mut self.file, &pid);
     }
 
-    fn update_information(&mut self, pid: u64, cpu: i32, is_ecore: bool) {
-        todo!()
-    }
+    fn update_information(&mut self, _pid: u64, _cpu: i32) {}
 
     fn read_consumption(&mut self, pid: u64) -> Option<f64> {
         let mut arg = pid;
diff --git a/src/energy/trackers/mock.rs b/src/energy/trackers/mock.rs
index cd08c34..231eb8f 100644
--- a/src/energy/trackers/mock.rs
+++ b/src/energy/trackers/mock.rs
@@ -3,13 +3,13 @@ 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) -> Result<(), ()> {
+    fn start_trace(&mut self, _pid: u64, _cpu: i32) -> Result<(), ()> {
         Ok(())
     }
 
     fn stop_trace(&mut self, _pid: u64) {}
 
-    fn update_information(&mut self, _pid: u64, _cpu: i32, _is_ecore: bool) {}
+    fn update_information(&mut self, _pid: u64, _cpu: i32) {}
 
     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 96c73d9..8862e91 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -2,34 +2,30 @@ use std::collections::HashMap;
 
 use burn::tensor::Tensor;
 use perf_event::{
-    events::{Event, Hardware, Software},
-    Builder, Counter, Group,
+    events::{Event, Hardware},
+    Builder, Counter,
 };
 
 use crate::freq::FrequencyKHZ;
 use std::sync::{Arc, RwLock};
 
 use crate::energy::Estimator;
-use crate::model::ArrayBackend;
+use crate::model::BurnBackend;
 
 pub struct PerfEstimator {
     registry: HashMap<u64, Counters>,
-    model_p: crate::model::Net<ArrayBackend>,
-    model_e: crate::model::Net<ArrayBackend>,
-    device: <ArrayBackend as burn::prelude::Backend>::Device,
+    model: crate::model::Net<BurnBackend>,
+    device: <BurnBackend 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_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 = crate::model::load_model("perf.pt");
         // let model_e = crate::model::load_model("perf_ecore.pt");
         Self {
             registry: Default::default(),
-            model_p,
-            model_e,
+            model,
             device: Default::default(),
             shared_cpu_current_frequencies,
         }
@@ -37,67 +33,41 @@ impl PerfEstimator {
 }
 
 struct Counters {
-    group: Group,
     counters: Vec<Counter>,
     old_time: u64,
     old_total_energy: f64,
     cpu: i32,
-    running_on_e_core: bool,
 }
 
 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),
+    #[cfg(not(feature = "amd"))]
     Event::Hardware(Hardware::REF_CPU_CYCLES),
 ];
-static EVENT_TYPES_E: &[Event] = &[
-    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),
-];
 
 impl Estimator for PerfEstimator {
-    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) => {
-                eprintln!(
-                    "Failed to create performance counter group for PID {}: {}",
-                    pid, e
-                );
-                return Err(());
-            }
-        };
-
-        let counters: Result<Vec<_>, _> = if running_on_e_core || true {
-            // println!("starting e core counter");
-            EVENT_TYPES_E
-        } else {
-            EVENT_TYPES_P
-        }
-        .iter()
-        .map(|kind| {
-            Builder::new()
-                // .group(&mut group)
-                .kind(kind.clone())
-                .observe_pid(pid as i32)
-                .inherit_thread(true)
-                .build()
-        })
-        .collect();
+    fn start_trace(&mut self, pid: u64, cpu: i32) -> Result<(), ()> {
+        let counters: Result<Vec<_>, _> = EVENT_TYPES_P
+            .iter()
+            .map(|kind| {
+                Builder::new()
+                    // .group(&mut group)
+                    .kind(kind.clone())
+                    .observe_pid(pid as i32)
+                    .inherit_thread(true)
+                    .build()
+            })
+            .collect();
 
         let mut counters = match counters {
             Ok(counters) => counters,
             Err(e) => {
                 eprintln!(
-                    "Failed to create performance counter group for PID {}: {}",
+                    "Failed to create performance counter for PID {}: {}",
                     pid, e
                 );
                 return Err(());
@@ -118,11 +88,9 @@ impl Estimator for PerfEstimator {
         let old_time = counters[0].read_count_and_time().unwrap().time_running;
         let counters = Counters {
             counters,
-            group,
             old_time,
             old_total_energy: 0.,
             cpu,
-            running_on_e_core,
         };
         self.registry.insert(pid, counters);
         Ok(())
@@ -132,20 +100,12 @@ impl Estimator for PerfEstimator {
         self.registry.remove(&pid);
     }
 
-    fn update_information(&mut self, pid: u64, cpu: i32, is_ecore: bool) {
-        let mut core_type_changed = false;
+    fn update_information(&mut self, pid: u64, cpu: i32) {
         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 {
-            // 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> {
@@ -154,13 +114,6 @@ 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 = counters.counters[0]
             .read_count_and_time()
             .unwrap()
@@ -172,7 +125,6 @@ impl Estimator for PerfEstimator {
         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,
         ];
@@ -181,21 +133,15 @@ impl Estimator for PerfEstimator {
             let count: u64 = ty.read().unwrap();
             values.push((count as f64) * correction_factor);
         }
+        #[cfg(feature = "amd")]
         values.push(values[4]);
 
-        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 result = self
+            .model
+            .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();
         for counter in counters.counters.iter_mut() {
             counter.reset().unwrap();
         }