From 92c34f2e9e61f5e8ff0ab42778b830aeee0d7c4e Mon Sep 17 00:00:00 2001
From: hanhuihui <hanhuihui5@h-partners.com>
Date: Thu, 14 May 2026 20:58:15 +0800
Subject: [PATCH] add task load judgment and optimized core selection logic

---
 scheds/rust/scx_soft_domain/src/bpf/intf.h    |  21 +-
 .../rust/scx_soft_domain/src/bpf/main.bpf.c   | 288 ++++++++++++++++--
 scheds/rust/scx_soft_domain/src/main.rs       | 247 ++++++++++++++-
 3 files changed, 520 insertions(+), 36 deletions(-)

diff --git a/scheds/rust/scx_soft_domain/src/bpf/intf.h b/scheds/rust/scx_soft_domain/src/bpf/intf.h
index 366a9c7..2c60fbb 100644
--- a/scheds/rust/scx_soft_domain/src/bpf/intf.h
+++ b/scheds/rust/scx_soft_domain/src/bpf/intf.h
@@ -1,9 +1,4 @@
 /* SPDX-License-Identifier: GPL-2.0 */
-/*
- * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
- * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
- * Copyright (c) 2022 David Vernet <dvernet@meta.com>
- */
 
 #ifndef __INTF_H
 #define __INTF_H
@@ -46,6 +41,12 @@ typedef unsigned long long u64;
 #define TASK_COMM_LEN 16
 #endif
 
+#define MAX_CPUS_PER_NODE 192
+#define MAX_CPUS_PER_LLC 192
+#define IRQ_THRESHOLD 307               /* 30% */
+#define LOAD_THRESHOLD 614              /* 60% */
+#define SCX_DSQ_ID_BASE 1000
+
 typedef enum {
     ALLOW_LLC_CPUS  = 0,
     ALLOW_NODE_CPUS = 1,
@@ -54,10 +55,19 @@ typedef enum {
 
 #define CPUMASK_BYTES ((MAX_CPUS + 7) / 8)
 
+struct task_ctx {
+	u64 last_sum_exec;
+	u64 last_running_ns;
+	u64 last_stopping_ns;
+	u64 last_runnable_ns;
+	u64 runtime_ewma_sum_exec;
+};
+
 struct cpu_ctx {
     u32 cpu;
     u32 node_id;
     u32 llc_id;
+    u32 task_pid;
 } __attribute__((aligned(8)));
 
 struct llc_ctx {
@@ -89,4 +99,3 @@ struct topo_ctx {
 } __attribute__((aligned(8)));
 
 #endif /* __INTF_H */
-
diff --git a/scheds/rust/scx_soft_domain/src/bpf/main.bpf.c b/scheds/rust/scx_soft_domain/src/bpf/main.bpf.c
index 420c9db..3e91a07 100644
--- a/scheds/rust/scx_soft_domain/src/bpf/main.bpf.c
+++ b/scheds/rust/scx_soft_domain/src/bpf/main.bpf.c
@@ -19,12 +19,58 @@
 #endif
 
 #include "intf.h"
-
 char _license[] SEC("license") = "GPL";
 
 const volatile u32 max_cpus = MAX_CPUS;
 const volatile s32 allowed_node = -1;
 const volatile char target_comm[TASK_COMM_LEN] = "";
+static const __u16 pelt_subperiod_q10[9] = {
+	1024, /*   0us */
+	1021, /* 128us */
+	1018, /* 256us */
+	1016, /* 384us */
+	1013, /* 512us */
+	1010, /* 640us */
+	1007, /* 768us */
+	1005, /* 896us */
+	1002  /* 1024us ~= 1 period */
+};
+
+static const __u16 pelt_period_q10[33] = {
+	1024, /*  0 */
+	1002, /*  1 */
+	 981, /*  2 */
+	 960, /*  3 */
+	 939, /*  4 */
+	 919, /*  5 */
+	 899, /*  6 */
+	 880, /*  7 */
+	 861, /*  8 */
+	 843, /*  9 */
+	 825, /* 10 */
+	 807, /* 11 */
+	 790, /* 12 */
+	 773, /* 13 */
+	 756, /* 14 */
+	 740, /* 15 */
+	 724, /* 16 */
+	 709, /* 17 */
+	 693, /* 18 */
+	 679, /* 19 */
+	 664, /* 20 */
+	 650, /* 21 */
+	 636, /* 22 */
+	 622, /* 23 */
+	 609, /* 24 */
+	 596, /* 25 */
+	 583, /* 26 */
+	 571, /* 27 */
+	 558, /* 28 */
+	 546, /* 29 */
+	 535, /* 30 */
+	 523, /* 31 */
+	 512  /* 32 => 0.5 */
+};
 
 UEI_DEFINE(uei);
 
@@ -56,11 +102,68 @@ struct {
 	__uint(max_entries, 1);
 } topo_ctxs SEC(".maps");
 
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, MAX_LLCS * MAX_CPUS_PER_LLC);
+	__type(key, u32);
+	__type(value, u32);
+} llc_sorted_cpu_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, MAX_NUMA_NODES * MAX_CPUS_PER_NODE);
+	__type(key, u32);
+	__type(value, u32);
+} node_sorted_cpu_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, MAX_LLCS);
+	__type(key, u32);
+	__type(value, u32);
+} llc_sched_idx SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, MAX_NUMA_NODES);
+	__type(key, u32);
+	__type(value, u32);
+} node_sched_idx SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, MAX_CPUS);
+	__type(key, u32);
+	__type(value, u32);
+} cpu_load_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, MAX_CPUS);
+	__type(key, u32);
+	__type(value, u32);
+} cpu_load_irq_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_LRU_HASH);
+	__type(key, u32);
+	__type(value, s32);
+	__uint(max_entries, 65535);
+} cpus_selected SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, u32);
+	__type(value, struct task_ctx);
+} task_storage SEC(".maps");
+
 static __maybe_unused struct cpu_ctx *lookup_cpu_ctx(int cpu)
 {
 	struct cpu_ctx *cpuc;
 
 	cpuc = bpf_map_lookup_elem(&cpu_ctxs, &cpu);
+
 	if (!cpuc) {
 		scx_bpf_error("no cpu_ctx for cpu %d", cpu);
 		return NULL;
@@ -95,6 +198,27 @@ static struct llc_ctx *lookup_llc_ctx(uint32_t llc)
 	return llcx;
 }
 
+static __always_inline s32 get_next_idle_cpu_in_llc(u32 llc)
+{
+	u32 key = llc;
+	u32 *idx_ptr = bpf_map_lookup_elem(&llc_sched_idx, &key);
+	if (!idx_ptr)
+		return -1;
+
+	u32 old_idx = __sync_fetch_and_add(idx_ptr, 1);
+	struct llc_ctx *llc_c = lookup_llc_ctx(llc);
+	if (!llc_c)
+		return -1;
+	u32 idx_in_llc = old_idx % llc_c->nr_cpus;
+
+	u32 map_key = llc * MAX_CPUS_PER_LLC + idx_in_llc;
+	u32 *cpu_ptr = bpf_map_lookup_elem(&llc_sorted_cpu_map, &map_key);
+	if (!cpu_ptr)
+		return -1;
+
+	return (s32)*cpu_ptr;
+}
+
 static __always_inline s32 round_robin_select_llc_in_node(struct topo_ctx *topo, u32 node_id)
 {
 	if (!topo || topo->nr_llcs <= 0 || topo->nr_nodes <= 0) {
@@ -205,12 +329,22 @@ static __always_inline bool scx_bpf_check_comm(struct task_struct *p)
 	return true;
 }
 
-struct {
-	__uint(type, BPF_MAP_TYPE_LRU_HASH);
-	__type(key, u32);
-	__type(value, s32);
-	__uint(max_entries, 65535);
-} cpus_selected SEC(".maps");
+static __always_inline void update_cpu_tasks_by_pid(u32 pid_key, s32 cpu_new){
+	s32 *cpu_old = bpf_map_lookup_elem(&cpus_selected, &pid_key);
+	if (cpu_old){
+		if (*cpu_old == cpu_new){
+			return;
+		}
+		struct cpu_ctx *cpuc_old = lookup_cpu_ctx(*cpu_old);
+		if (cpuc_old) {
+			cpuc_old->task_pid = 0;
+		}
+	}
+	struct cpu_ctx *cpuc_new = lookup_cpu_ctx(cpu_new);
+	if (cpuc_new) {
+		cpuc_new->task_pid = pid_key;
+	}
+}
 
 static __always_inline void update_cpu_by_pid(u32 pid, s32 cpu_id)
 {
@@ -219,6 +353,19 @@ static __always_inline void update_cpu_by_pid(u32 pid, s32 cpu_id)
 	}
 }
 
+static __always_inline bool can_fast_insert_local(struct task_struct *p, s32 cpu)
+{
+	u32 *cpu_load = bpf_map_lookup_elem(&cpu_load_map, &cpu);
+	struct task_ctx *taskc = bpf_task_storage_get(&task_storage, p, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
+	if (!taskc || !cpu_load) {
+		return false;
+	}
+	if (*cpu_load < LOAD_THRESHOLD || (taskc->runtime_ewma_sum_exec) * 100 / *cpu_load > 80){
+		return true;
+	}
+	return false;
+}
+
 static __always_inline s32 scx_prev_select_cpu(struct task_struct *p, s32 prev_cpu, u64 wake_flags)
 {
 	s32 cpu;
@@ -262,6 +409,7 @@ s32 BPF_STRUCT_OPS(soft_domain_select_cpu, struct task_struct *p, s32 prev_cpu,
 		cpu_sid = bpf_map_lookup_elem(&cpus_selected, &pid_parent);
 		if (cpu_sid) {
 			cpu = *cpu_sid;
+			update_cpu_tasks_by_pid(pid_key, cpu);
 			update_cpu_by_pid(pid_key, cpu);
 			goto insert_cpu;
 		}
@@ -278,6 +426,7 @@ s32 BPF_STRUCT_OPS(soft_domain_select_cpu, struct task_struct *p, s32 prev_cpu,
 			cpu_sid = bpf_map_lookup_elem(&cpus_selected, &pid_parent);
 			if (cpu_sid) {
 				cpu = *cpu_sid;
+				update_cpu_tasks_by_pid(pid_key, cpu);
 				update_cpu_by_pid(pid_key, cpu);
 				goto insert_cpu;
 			}
@@ -307,6 +456,7 @@ s32 BPF_STRUCT_OPS(soft_domain_select_cpu, struct task_struct *p, s32 prev_cpu,
 
 	cpu = round_robin_select_cpu(topo, allowed_node, cpu_allowed_flag);
 	if (cpu >= 0) {
+		update_cpu_tasks_by_pid(pid_key, cpu);
 		update_cpu_by_pid(pid_key, cpu);
 		goto insert_cpu;
 	}
@@ -334,29 +484,68 @@ void BPF_STRUCT_OPS(soft_domain_enqueue, struct task_struct *p, u64 enq_flags)
 		cpu = scx_bpf_pick_idle_cpu(p->cpus_ptr, 0);
 		if (cpu >= 0) {
 			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, SCX_SLICE_DFL, 0);
+			scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE);
 			return;
 		}
 
 		cpu = scx_bpf_pick_any_cpu(p->cpus_ptr, 0);
 		if (cpu >= 0) {
 			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, SCX_SLICE_DFL, 0);
+			scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE);
 			return;
 		}
 
 		return;
 	}
-
 	u32 pid = p->pid;
 	s32 *cpu_ptr = bpf_map_lookup_elem(&cpus_selected, &pid);
-	if (cpu_ptr) {
-		s32 fixed_cpu = *cpu_ptr;
-		if (scx_bpf_task_cpu(p) != fixed_cpu) {
-			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | fixed_cpu, SCX_SLICE_DFL, 0);
-			return;
-		}
+	if (!cpu_ptr) {
+		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
+		return;
 	}
 
-	scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
+	if (can_fast_insert_local(p, *cpu_ptr)) {
+		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | *cpu_ptr, SCX_SLICE_DFL, 0);
+		return;
+	}
+
+	struct task_ctx *taskc = bpf_task_storage_get(&task_storage, p, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
+	struct cpu_ctx *cpuc  = lookup_cpu_ctx(*cpu_ptr);
+	if (!taskc || !cpuc) {
+		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
+		return;
+	}
+
+	u64 dsq_id = SCX_DSQ_ID_BASE + cpuc->llc_id;
+	scx_bpf_dsq_insert(p, dsq_id, SCX_SLICE_DFL, 0);
+	u32 idle_cpu = get_next_idle_cpu_in_llc(cpuc->llc_id);
+	u32 *idle_cpu_load = bpf_map_lookup_elem(&cpu_load_map, &idle_cpu);
+	if (!idle_cpu_load || *idle_cpu_load > 307 || taskc->runtime_ewma_sum_exec + *idle_cpu_load > 900) {
+		return;
+	}
+
+	struct cpu_ctx *cpuc_idle_cpu = lookup_cpu_ctx(idle_cpu);
+	if (!cpuc_idle_cpu) {
+		return;
+	}
+	u32 pid_idle_cpu = cpuc_idle_cpu->task_pid;
+	if (!pid_idle_cpu) {
+		update_cpu_tasks_by_pid(pid, idle_cpu);
+		update_cpu_by_pid(pid, idle_cpu);
+		scx_bpf_kick_cpu(idle_cpu, SCX_KICK_IDLE);
+		return;
+	}
+	scx_bpf_kick_cpu(idle_cpu, SCX_KICK_IDLE);
+}
+
+void BPF_STRUCT_OPS(soft_domain_dispatch, s32 cpu, struct task_struct *p)
+{
+	struct cpu_ctx *cpuc = lookup_cpu_ctx(cpu);
+	if (cpuc) {
+		u64 dsq_id = SCX_DSQ_ID_BASE + cpuc->llc_id;
+		scx_bpf_dsq_move_to_local(dsq_id);
+		return;
+	}
 }
 
 static s32 create_topo_cpumask(struct bpf_cpumask **kptr)
@@ -475,10 +664,10 @@ int32_t BPF_STRUCT_OPS_SLEEPABLE(soft_domain_init)
 	struct topo_ctx *topo;
 
 	topo = bpf_map_lookup_elem(&topo_ctxs, &key);
-	if (!topo) {
-		bpf_printk("topo_ctxs lookup failed\n");
+    if (!topo) {
+        bpf_printk("topo_ctxs lookup failed\n");
 		return -ENOENT;
-	}
+    }
 
 	bpf_for(i, 0, topo->nr_nodes) {
 		ret = init_node_bpf_mask(i);
@@ -494,18 +683,69 @@ int32_t BPF_STRUCT_OPS_SLEEPABLE(soft_domain_init)
 		}
 	}
 
+	bpf_for(i, 0, topo->nr_llcs) {
+		u64 dsq_id = SCX_DSQ_ID_BASE + i;
+		ret = scx_bpf_create_dsq(dsq_id, -1);
+		if (ret) {
+			return ret;
+		}
+	}
+
 	return 0;
 }
 
+void BPF_STRUCT_OPS(soft_domain_runnable, struct task_struct *p)
+{
+	struct task_ctx *taskc;
+
+	taskc = bpf_task_storage_get(&task_storage, p, 0, BPF_LOCAL_STORAGE_GET_F_CREATE);
+	if (!taskc) 
+		return;
+	taskc->last_runnable_ns = bpf_ktime_get_ns();
+	taskc->last_sum_exec = BPF_CORE_READ(p, se.sum_exec_runtime);
+}
+
+static __always_inline u64 update_pelt_load(u64 runtime_ewma_sum_exec, u64 delta_running, u64 delta_runnable)
+{
+	u64 elapsed_us = delta_runnable / 1000;
+	u64 decay_periods = elapsed_us / 1024;
+	u64 rem_us  = elapsed_us % 1024;
+	u64 sub_idx = rem_us / 128;
+	if (decay_periods > 32)
+		decay_periods = 32;
+	if (sub_idx > 8)
+		sub_idx = 8;
+	if (!decay_periods){
+		if (!sub_idx)
+			sub_idx = 1;
+		return ((runtime_ewma_sum_exec * pelt_subperiod_q10[sub_idx]) >> 10) + delta_running * (1024 - pelt_subperiod_q10[sub_idx]) / delta_runnable;
+	}
+	return ((runtime_ewma_sum_exec * pelt_period_q10[decay_periods]) >> 10) + delta_running * (1024 - pelt_period_q10[decay_periods]) / delta_runnable;
+}
+
+void BPF_STRUCT_OPS(soft_domain_quiescent, struct task_struct *p)
+{
+	u64 now, delta_runnable, delta_running_exec;
+	struct task_ctx *taskc = bpf_task_storage_get(&task_storage, p, 0, 0);
+	if (!taskc || !taskc->last_runnable_ns)
+		return;
+	now = bpf_ktime_get_ns();
+	delta_running_exec = BPF_CORE_READ(p, se.sum_exec_runtime) - taskc->last_sum_exec;
+	delta_runnable = now - taskc->last_runnable_ns;
+	taskc->runtime_ewma_sum_exec = update_pelt_load(taskc->runtime_ewma_sum_exec, delta_running_exec, delta_runnable);
+}
+
 void BPF_STRUCT_OPS(soft_domain_exit, struct scx_exit_info *ei)
 {
 	UEI_RECORD(uei, ei);
 }
 
 SCX_OPS_DEFINE(soft_domain_ops,
-	.select_cpu		= (void *)soft_domain_select_cpu,
-	.enqueue 		= (void *)soft_domain_enqueue,
-	.init			= (void *)soft_domain_init,
-	.exit			= (void *)soft_domain_exit,
-	.name			= "soft_domain");
-
+		.select_cpu		= (void *)soft_domain_select_cpu,
+		.enqueue 		= (void *)soft_domain_enqueue,
+		.dispatch 		= (void *)soft_domain_dispatch,
+		.runnable		= (void *)soft_domain_runnable,
+		.quiescent		= (void *)soft_domain_quiescent,
+		.init			= (void *)soft_domain_init,
+		.exit			= (void *)soft_domain_exit,
+		.name			= "soft_domain");
\ No newline at end of file
diff --git a/scheds/rust/scx_soft_domain/src/main.rs b/scheds/rust/scx_soft_domain/src/main.rs
index 8141f8b..27dbbc8 100644
--- a/scheds/rust/scx_soft_domain/src/main.rs
+++ b/scheds/rust/scx_soft_domain/src/main.rs
@@ -1,5 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0
 use std::collections::HashSet;
+use std::fs::File;
+use std::io::{BufRead, BufReader};
 use std::mem::MaybeUninit;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -40,6 +42,12 @@ struct Scheduler<'a> {
     skel: BpfSkel<'a>,
     _link: libbpf_rs::Link,
     shutdown: Arc<AtomicBool>,
+    nr_cpus: u32,
+    nr_llc: usize,
+    nr_node: usize,
+    cpu_llcs: Vec<u32>,
+    cpu_nodes: Vec<u32>,
+    prev_stats: Vec<u64>,
 }
 
 impl<'a> Scheduler<'a> {
@@ -58,8 +66,12 @@ impl<'a> Scheduler<'a> {
         let mut open_skel = scx_ops_open!(skel_builder, open_object, soft_domain_ops, open_opts)?;
 
         let nr_cpus = Self::init_globals(&mut open_skel, opts)?;
+        let cpu_llcs = Self::detect_cpu_llc_mapping(nr_cpus as usize)?;
+        let cpu_nodes = Self::detect_cpu_node_mapping(nr_cpus as usize)?;
+        let nr_llc = *cpu_llcs.iter().max().unwrap_or(&0) as usize + 1;
+        let nr_node = *cpu_nodes.iter().max().unwrap_or(&0) as usize + 1;
         let mut skel = scx_ops_load!(open_skel, soft_domain_ops, uei)?;
-        Self::init_topology(&mut skel, nr_cpus)?;
+        Self::init_topology(&mut skel, nr_cpus, nr_llc, nr_node)?;
         let link = scx_ops_attach!(skel, soft_domain_ops)?;
 
         info!("{} scheduler initialized", SCHEDULER_NAME);
@@ -67,9 +79,60 @@ impl<'a> Scheduler<'a> {
             skel,
             _link: link,
             shutdown: Arc::new(AtomicBool::new(false)),
+            nr_cpus,
+            nr_llc,
+            nr_node,
+            cpu_llcs,
+            cpu_nodes,
+            prev_stats: vec![0u64; MAX_CPUS as usize * 3],
         })
     }
 
+    fn detect_cpu_llc_mapping(nr_cpus: usize) -> Result<Vec<u32>> {
+        let mut map = vec![0u32; nr_cpus];
+        for cpu in 0..nr_cpus {
+            let path = format!("/sys/devices/system/cpu/cpu{}/topology/cluster_id", cpu);
+            if !std::path::Path::new(&path).exists() {
+                continue;
+            }
+            let content = std::fs::read_to_string(&path)?;
+            let cluster_id = content.trim().parse::<u32>().map_err(|e| {
+                std::io::Error::new(std::io::ErrorKind::InvalidData, e)
+            })?;
+            map[cpu] = cluster_id;
+        }
+        Ok(map)
+    }
+
+    fn detect_cpu_node_mapping(nr_cpus: usize) -> Result<Vec<u32>> {
+        let mut map = vec![0u32; nr_cpus];
+        for node in 0..MAX_NUMA_NODES {
+            let path = format!("/sys/devices/system/node/node{}/cpulist", node);
+            if !std::path::Path::new(&path).exists() {
+                break;
+            }
+            let content = std::fs::read_to_string(&path)?;
+            for range in content.trim().split(',') {
+                let parts: Vec<&str> = range.split('-').collect();
+                if parts.len() == 1 {
+                    let cpu = parts[0].parse::<usize>()?;
+                    if cpu < nr_cpus {
+                        map[cpu] = node;
+                    }
+                } else if parts.len() == 2 {
+                    let start = parts[0].parse::<usize>()?;
+                    let end = parts[1].parse::<usize>()?;
+                    for cpu in start..=end {
+                        if cpu < nr_cpus {
+                            map[cpu] = node;
+                        }
+                    }
+                }
+            }
+        }
+        Ok(map)
+    }
+
     fn init_globals(skel: &mut OpenBpfSkel, opts: &Opts) -> Result<u32> {
         let rodata = skel.maps.rodata_data.as_mut().unwrap();
 
@@ -92,7 +155,7 @@ impl<'a> Scheduler<'a> {
         Ok(rodata.max_cpus)
     }
 
-    fn init_topology(skel: &mut BpfSkel, nr_cpus: u32) -> Result<()> {
+    fn init_topology(skel: &mut BpfSkel, nr_cpus: u32, nr_llc: usize, nr_node: usize) -> Result<()> {
         let nr_cpus = nr_cpus as usize;
         eprintln!("[INIT] Initializing topology for {} CPUs", nr_cpus);
 
@@ -121,12 +184,12 @@ impl<'a> Scheduler<'a> {
             .collect();
         let nr_llcs = llc_list.len();
 
-        // cpu_ctxs
         for cpu in 0..nr_cpus {
             let ctx = cpu_ctx {
                 cpu: cpu as u32,
                 node_id: cpu_node[cpu],
                 llc_id: cpu_llc[cpu],
+                task_pid:0,
             };
             let key = (cpu as u32).to_ne_bytes();
             let value = unsafe {
@@ -139,7 +202,6 @@ impl<'a> Scheduler<'a> {
                 .with_context(|| format!("Failed to update cpu_ctx for CPU {}", cpu))?;
         }
 
-        // node_ctxs
         let nr_nodes = node_cpus.iter().filter(|v| !v.is_empty()).count();
         for node in 0..nr_nodes {
             let mut mask = [0u8; CPUMASK_BYTES as usize];
@@ -170,7 +232,6 @@ impl<'a> Scheduler<'a> {
                 .with_context(|| format!("Failed to update node_ctx for node {}", node))?;
         }
 
-        // llc_ctxs
         let mut llc_cpus_list = vec![Vec::new(); nr_llcs];
         for cpu in 0..nr_cpus {
             let llc_id = cpu_llc[cpu];
@@ -204,7 +265,6 @@ impl<'a> Scheduler<'a> {
                 .with_context(|| format!("Failed to update llc_ctx for llc index {}", idx))?;
         }
 
-        // topo_ctxs
         let topo = topo_ctx {
             nr_cpus: nr_cpus as u32,
             nr_nodes: nr_nodes as u32,
@@ -223,13 +283,153 @@ impl<'a> Scheduler<'a> {
         skel.maps.topo_ctxs.update(&key, value, MapFlags::ANY)
             .context("Failed to update topo_ctxs")?;
 
+        for llc in 0..nr_llc {
+            let key = (llc as u32).to_ne_bytes();
+            let zero = 0u32.to_ne_bytes();
+            skel.maps.llc_sched_idx.update(&key, &zero, MapFlags::ANY)
+                .context("Failed to init llc_sched_idx")?;
+        }
+
+        for llc in 0..nr_llc {
+            for idx in 0..MAX_CPUS {
+                let map_key = (llc as u32) * MAX_CPUS_PER_LLC as u32 + idx as u32;
+                let val = 0u32.to_ne_bytes();
+                let _ = skel.maps.llc_sorted_cpu_map.update(&map_key.to_ne_bytes(), &val, MapFlags::ANY);
+            }
+        }
+
+        for node in 0..nr_node {
+            let key = (node as u32).to_ne_bytes();
+            let zero = 0u32.to_ne_bytes();
+            skel.maps.node_sched_idx.update(&key, &zero, MapFlags::ANY)
+                .context("Failed to init node_sched_idx")?;
+        }
+
+        for node in 0..nr_node {
+            for idx in 0..MAX_CPUS {
+                let map_key = (node as u32) * MAX_CPUS_PER_NODE as u32 + idx as u32;
+                let val = 0u32.to_ne_bytes();
+                let _ = skel.maps.node_sorted_cpu_map.update(&map_key.to_ne_bytes(), &val, MapFlags::ANY);
+            }
+        }
+
         Ok(())
     }
 
+    fn update_cpu_loads(&mut self) {
+        let nr_cpus = self.nr_cpus as usize;
+
+        let curr_stats = match read_cpu_stats() {
+            Ok(stats) => stats,
+            Err(e) => {
+                debug!("Failed to read CPU stats: {}", e);
+                return;
+            }
+        };
+
+        if self.prev_stats.iter().all(|&v| v == 0) {
+            self.prev_stats = curr_stats;
+            return;
+        }
+
+        let mut loads = vec![0u32; MAX_CPUS as usize];
+        let mut irq_heavy = vec![0u32; MAX_CPUS as usize];
+        for cpu in 0..nr_cpus {
+            let prev_total = self.prev_stats[cpu * 3];
+            let prev_idle = self.prev_stats[cpu * 3 + 1];
+            let prev_irq = self.prev_stats[cpu * 3 + 2];
+            let curr_total = curr_stats[cpu * 3];
+            let curr_idle = curr_stats[cpu * 3 + 1];
+            let curr_irq = curr_stats[cpu * 3 + 2];
+
+            let total_diff = curr_total.saturating_sub(prev_total);
+            let idle_diff = curr_idle.saturating_sub(prev_idle);
+            let irq_diff = curr_irq.saturating_sub(prev_irq);
+
+            if total_diff > 0 {
+                let busy = total_diff.saturating_sub(idle_diff);
+                loads[cpu] = ((busy as u64 * 1024) / total_diff) as u32;
+                loads[cpu] = loads[cpu].min(1024);
+
+                irq_heavy[cpu] = ((irq_diff as u64 * 1024) / total_diff) as u32;
+                irq_heavy[cpu] = irq_heavy[cpu].min(1024);
+            }
+        }
+
+        for cpu in 0..nr_cpus {
+            let key = (cpu as u32).to_ne_bytes();
+            let val = loads[cpu].to_ne_bytes();
+            let val_irq = irq_heavy[cpu].to_ne_bytes();
+            if let Err(e) = self.skel.maps.cpu_load_map.update(&key, &val, MapFlags::ANY) {
+                debug!("Failed to update cpu_load_map[{}]: {}", cpu, e);
+            }
+            let _ = self.skel.maps.cpu_load_irq_map.update(&key, &val_irq, MapFlags::ANY);
+        }
+
+        let mut cpu_load_pairs: Vec<(u32, u32)> = (0..nr_cpus as u32)
+            .map(|cpu| (cpu, loads[cpu as usize]))
+            .collect();
+        cpu_load_pairs.sort_by_key(|(_, load)| *load);
+
+
+        let mut llc_cpus: Vec<Vec<u32>> = vec![Vec::new(); self.nr_llc];
+        for cpu in 0..self.nr_cpus {
+            let llc = self.cpu_llcs[cpu as usize] as usize;
+            llc_cpus[llc].push(cpu as u32);
+        }
+        for llc in 0..self.nr_llc {
+            llc_cpus[llc].sort_by_key(|&cpu| loads[cpu as usize]);
+
+            for (idx, &cpu) in llc_cpus[llc].iter().enumerate() {
+                let map_key = (llc as u32) * MAX_CPUS_PER_LLC as u32 + (idx as u32);
+                let val = (cpu as u32).to_ne_bytes();
+                let _ = self.skel.maps.llc_sorted_cpu_map.update(&map_key.to_ne_bytes(), &val, MapFlags::ANY);
+            }
+
+            let llc_key = (llc as u32).to_ne_bytes();
+            let zero = 0u32.to_ne_bytes();
+            let _ = self.skel.maps.llc_sched_idx.update(&llc_key, &zero, MapFlags::ANY);
+        }
+
+        let mut node_cpus: Vec<Vec<u32>> = vec![Vec::new(); self.nr_node];
+        for cpu in 0..self.nr_cpus {
+            let node = self.cpu_nodes[cpu as usize] as usize;
+            node_cpus[node].push(cpu as u32);
+        }
+
+        for node in 0..self.nr_node {
+            node_cpus[node].sort_by_key(|&cpu| loads[cpu as usize]);
+
+            for (idx, &cpu) in node_cpus[node].iter().enumerate() {
+                let map_key = (node as u32) * MAX_CPUS_PER_NODE as u32 + (idx as u32);
+                let val = (cpu as u32).to_ne_bytes();
+                let _ = self.skel.maps.node_sorted_cpu_map.update(&map_key.to_ne_bytes(), &val, MapFlags::ANY);
+            }
+
+            let node_key = (node as u32).to_ne_bytes();
+            let zero = 0u32.to_ne_bytes();
+            let _ = self.skel.maps.node_sched_idx.update(&node_key, &zero, MapFlags::ANY);
+        }
+
+        debug!(
+            "Updated CPU loads: node0 min={}, max={}; node1 min={}, max={}",
+            loads[node_cpus[0][0] as usize],
+            loads[*node_cpus[0].last().unwrap() as usize],
+            loads[node_cpus[1][0] as usize],
+            loads[*node_cpus[1].last().unwrap() as usize]
+        );
+
+        self.prev_stats = curr_stats;
+    }
+
+
     fn run(&mut self, _opts: &Opts) -> Result<UserExitInfo> {
         let shutdown = self.shutdown.clone();
+        self.update_cpu_loads();
+
         while !shutdown.load(Ordering::Relaxed) && !uei_exited!(&self.skel, uei) {
             std::thread::sleep(Duration::from_millis(100));
+            self.update_cpu_loads();
             if log::log_enabled!(log::Level::Debug) {
                 debug!("Scheduler is running");
             }
@@ -238,6 +438,40 @@ impl<'a> Scheduler<'a> {
     }
 }
 
+fn read_cpu_stats() -> Result<Vec<u64>> {
+    let file = File::open("/proc/stat")?;
+    let reader = BufReader::new(file);
+    let mut stats = Vec::new();
+
+    for line in reader.lines() {
+        let line = line?;
+        if line.starts_with("cpu ") {
+            continue;
+        }
+        if !line.starts_with("cpu") {
+            continue;
+        }
+        let parts: Vec<&str> = line.split_whitespace().collect();
+        if parts.len() < 5 {
+            continue;
+        }
+        let user: u64 = parts[1].parse().unwrap_or(0);
+        let nice: u64 = parts[2].parse().unwrap_or(0);
+        let system: u64 = parts[3].parse().unwrap_or(0);
+        let idle: u64 = parts[4].parse().unwrap_or(0);
+        let iowait: u64 = parts.get(5).and_then(|s| s.parse().ok()).unwrap_or(0);
+        let irq: u64 = parts.get(6).and_then(|s| s.parse().ok()).unwrap_or(0);
+        let softirq: u64 = parts.get(7).and_then(|s| s.parse().ok()).unwrap_or(0);
+        let steal: u64 = parts.get(8).and_then(|s| s.parse().ok()).unwrap_or(0);
+
+        let total = user + nice + system + idle + iowait + irq + softirq + steal;
+        stats.push(total);
+        stats.push(idle);
+        stats.push(irq + softirq);
+    }
+    Ok(stats)
+}
+
 fn read_sysfs_node_id(cpu: usize) -> Option<u32> {
     let cpu_path = format!("/sys/devices/system/cpu/cpu{}", cpu);
     if let Ok(entries) = std::fs::read_dir(&cpu_path) {
@@ -341,3 +575,4 @@ fn main() -> Result<()> {
     }
     Ok(())
 }
+
-- 
2.43.0