Fix warnings and general cleanup

16 files changed, 78 insertions, 236 deletions

diff --git a/.gitignore b/.gitignore
index 7c7ffae..94a5d27 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,6 @@
 intf.h
 main.bpf.c
 src/bpf.rs
+*.png
+*.csv
+*.pt
diff --git a/Cargo.toml b/Cargo.toml
index 55f3411..743fb9e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,6 +6,9 @@ name = "power_sched"
 version = "0.1.0"
 edition = "2021"
 
+[features]
+amd = []
+
 [dependencies]
 anyhow = "1.0.96"
 ctrlc = "3.4.5"
@@ -24,8 +27,6 @@ chrono = "0.4.40"
 burn = {version = "0.16.0", features = ["openblas-system","candle"] }
 bincode = "=2.0.0-rc.3"
 bincode_derive = "=2.0.0-rc.3"
-#burn-ndarray = {version = "0.16.0", default-features = false, features = ["burn-autodiff", "blas-openblas-system", "std"] }
-#burn-ndarray = {version = "0.16.0", default-features = false, features = ["burn-autodiff"] }
 burn-ndarray = {version = "0.16.0"}
 burn-import = "0.16.0"
 
diff --git a/perf-event/src/hooks.rs b/perf-event/src/hooks.rs
index d79331f..7cb9d40 100644
--- a/perf-event/src/hooks.rs
+++ b/perf-event/src/hooks.rs
@@ -134,6 +134,7 @@ macro_rules! expand_trait_method {
         #[doc = stringify!($ioctl)]
         /// .
         #[allow(non_snake_case)]
+        /// # Safety
         unsafe fn $name(&mut self, _fd: c_int, _arg: $arg_type) -> c_int {
             panic!(
                 "unimplemented `perf_event::hooks::Hooks` method: {}",
diff --git a/perf-event/src/lib.rs b/perf-event/src/lib.rs
index e09de18..babab41 100644
--- a/perf-event/src/lib.rs
+++ b/perf-event/src/lib.rs
@@ -429,7 +429,7 @@ pub struct CountAndTime {
     pub time_running: u64,
 }
 
-impl<'a> EventPid<'a> {
+impl EventPid<'_> {
     // Return the `pid` arg and the `flags` bits representing `self`.
     fn as_args(&self) -> (pid_t, u32) {
         match self {
@@ -1037,7 +1037,7 @@ impl std::ops::Index<&Counter> for Counts {
 
 impl std::fmt::Debug for Counts {
     fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
-        fmt.debug_map().entries(self.into_iter()).finish()
+        fmt.debug_map().entries(self).finish()
     }
 }
 
diff --git a/src/benchmark.rs b/src/benchmark.rs
index f7e30c5..c6795df 100644
--- a/src/benchmark.rs
+++ b/src/benchmark.rs
@@ -7,21 +7,18 @@ use perf_event::{
     events::{Event, Hardware},
     Builder, Counter, Group,
 };
-use rand::seq::IteratorRandom;
 use scx_utils::Topology;
 use scx_utils::UserExitInfo;
 use std::mem::MaybeUninit;
 use std::process;
 use std::thread;
 use std::time::{Duration, Instant};
-use std::{collections::HashMap, i32, ops::Range};
+use std::{collections::HashMap, ops::Range};
 use std::{fs::File, sync::atomic::AtomicI32};
 
 const SLICE_US: u64 = 5000;
 const LOG_INTERVAL_MS: u64 = 10;
 const RESHUFFLE_ROUNDS: usize = 1000; // Number of rounds before changing cpu
-const MAX_COUNTERS_AT_ONCE_P_CORE: usize = 7;
-const MAX_COUNTERS_AT_ONCE_E_CORE: usize = 6;
 type Pid = i32;
 static CPU: AtomicI32 = AtomicI32::new(0);
 
@@ -29,7 +26,6 @@ pub struct BenchmarkScheduler<'a> {
     bpf: BpfScheduler<'a>,
     own_pid: Pid,
     log_path: String,
-    mode: Mode,
     p_cores: Range<i32>,
     e_cores: Range<i32>,
 }
@@ -161,11 +157,7 @@ impl MeasurementDiff {
 }
 
 impl<'a> BenchmarkScheduler<'a> {
-    pub fn init(
-        open_object: &'a mut MaybeUninit<OpenObject>,
-        log_path: &str,
-        mode: Mode,
-    ) -> Result<Self> {
+    pub fn init(open_object: &'a mut MaybeUninit<OpenObject>, log_path: &str) -> Result<Self> {
         let bpf = BpfScheduler::init(
             open_object,
             0,     // exit_dump_len (default)
@@ -194,7 +186,6 @@ impl<'a> BenchmarkScheduler<'a> {
             bpf,
             own_pid: process::id() as i32,
             log_path: log_path.to_string(),
-            mode,
             p_cores,
             e_cores,
         })
@@ -228,11 +219,10 @@ impl<'a> BenchmarkScheduler<'a> {
 
     fn start_measurement_thread(&self) -> thread::JoinHandle<()> {
         let log_path = self.log_path.clone();
-        let mode = self.mode.clone();
         let e_cores = self.e_cores.clone();
         let p_cores = self.p_cores.clone();
         thread::spawn(move || {
-            if let Err(e) = run_measurement_loop(log_path, mode, p_cores.start, e_cores.start) {
+            if let Err(e) = run_measurement_loop(log_path, p_cores.start, e_cores.start) {
                 eprintln!("Measurement thread error: {:?}", e);
             }
         })
@@ -253,14 +243,13 @@ impl<'a> BenchmarkScheduler<'a> {
 }
 
 // Main measurement loop
-fn run_measurement_loop(log_path: String, mode: Mode, p_core: i32, e_core: i32) -> Result<()> {
+fn run_measurement_loop(log_path: String, p_core: i32, e_core: i32) -> Result<()> {
     // Define available hardware counters
     let available_events = define_available_events();
 
     // Initialize CSV writer with header
     let mut csv_writer = initialize_csv_writer(&log_path, &available_events)?;
 
-    let mut rng = rand::rng();
     let mut round_counter = 0;
 
     let mut cpu_to_monitor = p_core;
@@ -284,14 +273,7 @@ fn run_measurement_loop(log_path: String, mode: Mode, p_core: i32, e_core: i32)
             }
         };
 
-        // Select random subset of counters
-        let selected_events = available_events
-            .iter()
-            .choose_multiple(&mut rng, MAX_COUNTERS_AT_ONCE_P_CORE);
-
-        //let selected_events = available_events[0..MAX_COUNTERS_AT_ONCE_E_CORE].iter();
-
-        // println!("Selected {} events for monitoring", selected_events.len());
+        let selected_events = available_events.clone();
 
         // Build counters
         let mut counters = Vec::new();
diff --git a/src/energy.rs b/src/energy.rs
index 68a2a0e..0a49e07 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -5,7 +5,6 @@ mod trackers;
 
 use crate::energy::estimator::Estimator;
 use std::collections::{BTreeSet, HashMap};
-use std::ops::RangeInclusive;
 use std::sync::atomic::{AtomicBool, AtomicI32, AtomicI64, AtomicU64};
 use std::sync::{mpsc, Arc, RwLock};
 use std::thread;
@@ -39,10 +38,6 @@ impl TaskInfo {
     pub fn read_budget(&self) -> u64 {
         self.budget.load(std::sync::atomic::Ordering::Relaxed)
     }
-    pub fn is_running_on_e_core(&self) -> bool {
-        self.running_on_e_core
-            .load(std::sync::atomic::Ordering::Relaxed)
-    }
     pub fn read_time_since_last_schedule(&self) -> Option<Duration> {
         let old_time = self
             .last_scheduled
@@ -111,9 +106,6 @@ pub struct EnergyService {
     process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
     request_receiver: mpsc::Receiver<Request>,
     update_interval: Duration,
-    shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-    shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
     rapl_offset: f64,
     last_energy_diff: f64,
     last_time_between_measurements: Duration,
@@ -132,9 +124,6 @@ impl EnergyService {
         process_info: Arc<RwLock<HashMap<Pid, ProcessInfo>>>,
         request_receiver: mpsc::Receiver<Request>,
         update_interval: Duration,
-        shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-        shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-        shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
     ) -> Self {
         Self {
             estimator,
@@ -143,9 +132,6 @@ impl EnergyService {
             process_info,
             request_receiver,
             update_interval,
-            shared_cpu_frequency_ranges,
-            shared_policy_frequency_ranges,
-            shared_cpu_current_frequencies,
             rapl_offset: rapl::read_package_energy().unwrap(),
             last_energy_diff: 0f64,
             last_time_between_measurements: Duration::new(0, 0),
@@ -211,11 +197,7 @@ impl EnergyService {
                     .contains_key(&main_thread.pid)
                     && self
                         .estimator
-                        .start_trace(
-                            pid as u64,
-                            task_info.read_cpu(),
-                            task_info.is_running_on_e_core(),
-                        )
+                        .start_trace(pid as u64, task_info.read_cpu())
                         .is_err()
                 {
                     return;
@@ -276,11 +258,8 @@ impl EnergyService {
                     info.energy_delta = energy * self.bias;
                     info.energy += energy * self.bias;
                     info.tree_energy += energy * self.bias;
-                    self.estimator.update_information(
-                        *pid as u64,
-                        info.task_info.read_cpu(),
-                        info.task_info.is_running_on_e_core(),
-                    );
+                    self.estimator
+                        .update_information(*pid as u64, info.task_info.read_cpu());
                     let mut parent = info.parent;
                     while let Some(info) = process_info.get_mut(&parent) {
                         info.tree_energy += energy * self.bias;
@@ -297,25 +276,19 @@ impl EnergyService {
         let rapl_diff = rapl - self.old_rapl;
         self.last_energy_diff = rapl_diff;
         self.old_rapl = rapl;
-        let power_comsumption_watt = rapl_diff / elapsed.as_secs_f64();
         let idle_consumption = elapsed.as_secs_f64() * IDLE_CONSUMPTION_W;
+
         if let Some(init) = self.process_info.write().unwrap().get_mut(&1) {
             let est_diff = init.tree_energy - old_energy + idle_consumption;
-            // let offset_bias = (rapl / (init.tree_energy + idle_consumption)).clamp(0.1, 2.);
             let current_bias = if init.tree_energy - old_energy > idle_consumption * 0.5 {
                 (rapl_diff / est_diff).clamp(0.1, 2.)
             } else {
                 1.
             };
-            // let current_bias = (offset_bias + diff_bias) * 0.5;
             let alpha: f64 = 10. * elapsed.as_secs_f64().recip();
             self.bias = (self.bias * (alpha.recip() * current_bias + ((alpha - 1.) / alpha)))
                 .clamp(0.1, 5.);
             self.system_energy += est_diff;
-            // println!(
-            //     "Energy estimation: {:.1} rapl: {:.1}, est diff: {:.1} rapl diff: {:.1}, bias: {:.1}, power consumption: {:.1}",
-            //     self.system_energy, rapl, est_diff, rapl_diff, self.bias, power_comsumption_watt,
-            // );
         }
     }
 
@@ -341,23 +314,6 @@ impl EnergyService {
         }
     }
 
-    // Accessor methods for BudgetPolicy
-    pub fn active_processes(&self) -> &BTreeSet<Pid> {
-        &self.active_processes
-    }
-
-    fn package_energy(&mut self) -> f64 {
-        rapl::read_package_energy().unwrap()
-    }
-
-    pub fn process_energy(&self, pid: Pid) -> Option<f64> {
-        self.process_info
-            .read()
-            .unwrap()
-            .get(&pid)
-            .map(|info| info.energy)
-    }
-
     pub fn all_process_energy_deltas(&self) -> HashMap<Pid, f64> {
         self.process_info
             .read()
@@ -370,8 +326,6 @@ impl EnergyService {
 
 pub fn start_energy_service(
     use_mocking: bool,
-    shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-    shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
     shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
 ) -> std::io::Result<mpsc::SyncSender<Request>> {
     // Potentially convert back to bounded channel
@@ -401,9 +355,6 @@ pub fn start_energy_service(
         process_info.clone(),
         request_receiver,
         UPDATE_INTERVAL,
-        shared_cpu_frequency_ranges,
-        shared_policy_frequency_ranges,
-        shared_cpu_current_frequencies,
     );
 
     service.run();
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();
         }
diff --git a/src/freq.rs b/src/freq.rs
index aa575f9..17def9a 100644
--- a/src/freq.rs
+++ b/src/freq.rs
@@ -8,6 +8,7 @@ use std::time::Duration;
 
 pub type FrequencyKHZ = u32;
 
+#[allow(dead_code)]
 pub enum Governor {
     Conservative,
     Ondemand,
@@ -30,10 +31,10 @@ impl Governor {
     }
 }
 
+#[allow(dead_code)]
 pub enum Request {
-    GetPossibleCPUFrequencyRange,
-    GetPolicyCPUFrequency,
-    GetCurrentFrequencies,
+    UpdatePossibleCPUFrequencyRange,
+    UpdatePolicyCPUFrequency,
     SetGovernorForCore(u32, Governor),
     SetFrequencyRangeAllCores(RangeInclusive<FrequencyKHZ>),
     SetFrequencyRangeForCore(u32, RangeInclusive<FrequencyKHZ>),
@@ -245,18 +246,14 @@ impl SysFSFrequencyService {
     fn handle_requests(&mut self) -> io::Result<()> {
         while let Ok(request) = self.request_receiver.try_recv() {
             match request {
-                Request::GetPossibleCPUFrequencyRange => {
+                Request::UpdatePossibleCPUFrequencyRange => {
                     *self.cpu_frequency_ranges.write().unwrap() =
                         self.get_possible_cpu_frequency_range()?;
                 }
-                Request::GetPolicyCPUFrequency => {
+                Request::UpdatePolicyCPUFrequency => {
                     *self.policy_frequency_ranges.write().unwrap() =
                         self.get_policy_cpu_frequency_ranges()?;
                 }
-                Request::GetCurrentFrequencies => {
-                    *self.cpu_current_frequencies.write().unwrap() =
-                        self.get_current_frequencies()?;
-                }
                 Request::SetGovernorForCore(cpu, governor) => {
                     self.switch_governor(cpu, governor)?;
                 }
diff --git a/src/main.rs b/src/main.rs
index d5a3dcd..00c0fcc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,5 @@
 mod bpf_skel;
-use benchmark::{BenchmarkScheduler, Mode};
+use benchmark::BenchmarkScheduler;
 pub use bpf_skel::*;
 pub mod bpf_intf;
 
@@ -12,12 +12,14 @@ mod scheduler;
 mod socket;
 
 #[rustfmt::skip]
+#[allow(clippy::unwrap_or_default)]
+#[allow(clippy::useless_asref)]
+#[allow(clippy::single_match)]
+#[allow(clippy::needless_lifetimes)]
 mod bpf;
 
 use anyhow::Result;
-use burn::tensor::Tensor;
 use clap::{Arg, ArgAction, Command};
-use model::load_model;
 use scheduler::Scheduler;
 use std::mem::MaybeUninit;
 
@@ -37,36 +39,20 @@ fn main() -> Result<()> {
             Arg::new("benchmark")
                 .short('b')
                 .long("benchmark")
-                .help("Enable benchmarking mode. Use \"e\" for e cores and \"p\" for p cores")
+                .help("Enable benchmarking mode.")
                 .required(false)
-                .value_name("mode"),
+                .action(ArgAction::SetTrue),
         )
         .get_matches();
 
-    let device = Default::default();
-    let model = load_model("perf.pt");
-    let tensor = Tensor::from_floats(
-        [
-            800., 90678., 54734., 153646., 20354478., 40948418.,
-            89103105.,
-            //5200., 148947., 322426., 498965., 62340773., 144451046., 41976480.,
-        ],
-        &device,
-    );
-    let result = model.forward(tensor);
-    let energy: f32 = result.into_scalar();
-    println!("energy: {energy}");
-    // panic!();
-
     let use_perf = matches.get_flag("perf");
-    let benchmark = matches.get_one::<String>("benchmark");
+    let benchmark = matches.get_flag("benchmark");
 
     // Initialize and load the scheduler.
     let mut open_object = MaybeUninit::uninit();
     let log_path = "/tmp/logs.csv";
-    if let Some(mode) = benchmark {
-        let mode = mode.trim().chars().next().unwrap();
-        let mut sched = BenchmarkScheduler::init(&mut open_object, log_path, Mode::from(mode))?;
+    if benchmark {
+        let mut sched = BenchmarkScheduler::init(&mut open_object, log_path)?;
         sched.run()?;
         return Ok(());
     }
diff --git a/src/model.rs b/src/model.rs
index 6dbde1f..eb62c35 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -7,9 +7,7 @@ use nn::LinearConfig;
 use burn::record::{FullPrecisionSettings, Recorder};
 use burn_import::pytorch::PyTorchFileRecorder;
 
-//type ArrayBackend = burn_ndarray::NdArray<f32>;
-//pub type ArrayBackend = burn_wgpu::Wgpu;
-pub type ArrayBackend = burn::backend::candle::Candle;
+pub type BurnBackend = burn::backend::candle::Candle;
 
 #[derive(Module, Debug)]
 pub struct Net<B: Backend> {
@@ -44,17 +42,17 @@ impl<B: Backend> Net<B> {
         let x = self.relu1.forward(x);
         let x = self.lin2.forward(x);
         let x = self.relu2.forward(x);
-        let x = self.lin3.forward(x);
-        x
+
+        self.lin3.forward(x)
     }
 }
 
 /// Load the p core model from the file in your source code (not in build.rs or script).
-pub fn load_model(path: &str) -> Net<ArrayBackend> {
+pub fn load_model(path: &str) -> Net<BurnBackend> {
     let device = Default::default();
-    let record: NetRecord<ArrayBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
+    let record: NetRecord<BurnBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
         .load(path.into(), &device)
         .expect("Failed to decode state");
 
-    Net::<ArrayBackend>::init(&device).load_record(record)
+    Net::<BurnBackend>::init(&device).load_record(record)
 }
diff --git a/src/scheduler.rs b/src/scheduler.rs
index fa86ef2..0832049 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -10,10 +10,10 @@ use scx_utils::{Topology, UserExitInfo};
 use std::collections::{HashMap, VecDeque};
 use std::mem::MaybeUninit;
 use std::ops::{Range, RangeInclusive};
+use std::process;
 use std::sync::mpsc::TrySendError;
 use std::sync::{mpsc, Arc, RwLock};
 use std::time::Duration;
-use std::{process, usize};
 
 use crate::Pid;
 
@@ -24,23 +24,17 @@ pub struct Scheduler<'a> {
     task_queue: VecDeque<QueuedTask>,
     no_budget_task_queue: VecDeque<QueuedTask>,
     managed_tasks: HashMap<Pid, Arc<TaskInfo>>,
-    maximum_budget: u64,
     tasks_scheduled: u64,
     //TODO: also consider Pids of children
     own_pid: Pid,
     p_cores: Range<i32>,
     e_cores: Option<Range<i32>>,
     garbage_core: i32,
-    topology: Topology,
     to_remove: Vec<Pid>,
     e_core_selector: Box<dyn CoreSelector>,
-    p_core_selector: Box<dyn CoreSelector>,
     energy_sender: mpsc::SyncSender<EnergyRequest>,
     empty_task_infos: mpsc::Receiver<Arc<TaskInfo>>,
     frequency_sender: mpsc::SyncSender<FrequencyRequest>,
-    shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-    shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
-    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
 }
 
 impl<'a> Scheduler<'a> {
@@ -55,12 +49,8 @@ impl<'a> Scheduler<'a> {
             Arc::new(RwLock::new(Vec::new()));
 
         // Start energy tracking service
-        let energy_sender = energy::start_energy_service(
-            use_mocking,
-            shared_cpu_frequency_ranges.clone(),
-            shared_policy_frequency_ranges.clone(),
-            shared_cpu_current_frequencies.clone(),
-        )?;
+        let energy_sender =
+            energy::start_energy_service(use_mocking, shared_cpu_current_frequencies.clone())?;
 
         let (task_sender, empty_task_infos) = mpsc::sync_channel(200);
         std::thread::spawn(move || loop {
@@ -91,7 +81,6 @@ impl<'a> Scheduler<'a> {
             None
         };
 
-        let p_core_selector = Box::new(RoundRobinSelector::new(&p_cores));
         let e_core_selector = if let Some(e_cores) = &e_cores {
             // reserve the last e core as garbage core
             Box::new(RoundRobinSelector::new(
@@ -113,10 +102,10 @@ impl<'a> Scheduler<'a> {
         )?;
 
         frequency_sender
-            .try_send(FrequencyRequest::GetPolicyCPUFrequency)
+            .try_send(FrequencyRequest::UpdatePolicyCPUFrequency)
             .unwrap();
         frequency_sender
-            .try_send(FrequencyRequest::GetPossibleCPUFrequencyRange)
+            .try_send(FrequencyRequest::UpdatePossibleCPUFrequencyRange)
             .unwrap();
 
         std::thread::sleep(Duration::from_secs(1));
@@ -133,22 +122,16 @@ impl<'a> Scheduler<'a> {
             task_queue: VecDeque::new(),
             no_budget_task_queue: VecDeque::new(),
             managed_tasks: HashMap::new(),
-            maximum_budget: u64::MAX,
             own_pid: process::id() as i32,
             p_cores,
             empty_task_infos,
             tasks_scheduled: 0,
             e_cores,
             garbage_core: 0,
-            topology,
             e_core_selector,
-            p_core_selector,
             energy_sender,
             to_remove,
             frequency_sender,
-            shared_cpu_frequency_ranges,
-            shared_policy_frequency_ranges,
-            shared_cpu_current_frequencies,
         })
     }
 
@@ -211,9 +194,6 @@ impl<'a> Scheduler<'a> {
             }
         }
     }
-    fn dispatch_next_task(&mut self) {
-        self.batch_dispatch_next_tasks(1);
-    }
 
     fn batch_dispatch_next_tasks(&mut self, tasks: i32) {
         for _ in 0..tasks {