refactor policies to per-process scheduling granularity

sched/include/xsched/sched/policy/cfs.h

@@ -26,8 +26,9 @@ class CompletelyFairSchedulerPolicy : public Policy {
     virtual void RecvHint(std::shared_ptr<const Hint> hint) override;
 
 private:
-    std::unordered_map<XQueueHandle, CFSNode> cfs_infos_;
-    std::chrono::microseconds time_slice_{1000}; // set timeslice to 1ms
+    std::unordered_map<PID, CFSNode> cfs_infos_;
+    std::unordered_map<XQueueHandle, std::pair<Priority, double>> pending_hints_;
+    std::chrono::microseconds time_slice_{1000};
 };
 
 } // namespace xsched::sched

sched/include/xsched/sched/policy/kedf.h

@@ -9,9 +9,9 @@
 namespace xsched::sched
 {
 
-struct DeadlineEntry
+struct ProcessDeadlineEntry
 {
-    XQueueHandle xqueue;
+    PID pid;
     std::chrono::system_clock::time_point deadline;
 };
 
@@ -25,6 +25,8 @@ class KEarliestDeadlineFirstPolicy : public Policy
     virtual void RecvHint(std::shared_ptr<const Hint> hint) override;
 
 private:
+    void SwitchProcess(PID pid, const Status &status);
+
     size_t k_ = 1;
     std::unordered_map<XQueueHandle, Deadline> deadlines_;
 };

sched/include/xsched/sched/policy/lax.h

@@ -25,7 +25,10 @@ class LaxityPolicy : public Policy
         Priority lax_prio;
         Priority crit_prio;
     };
-    std::unordered_map<XQueueHandle, LaxityInfo> laxity_infos_;
+    std::unordered_map<PID, LaxityInfo> laxity_infos_;
+    std::unordered_map<XQueueHandle, PID> handle_to_pid_;
+
+    void SwitchProcess(PID pid, const Status &status);
 };
 
 } // namespace xsched::sched

sched/include/xsched/sched/policy/mlfq.h

@@ -9,21 +9,20 @@ namespace xsched {
 namespace sched {
 
 struct MLFQNode {
-    int priority = 0;               // Current priority level (0 is highest)
-    bool is_running = false;        // Whether it's currently running
-    bool was_ready_last_tick = false; // Edge detection for idle/ready
+    int priority = 0;
+    bool is_running = false;
+    bool was_ready_last_tick = false;
 
     using TimePoint = std::chrono::time_point<std::chrono::system_clock>;
-    TimePoint i_a;                  // (1) Time became idle
-    TimePoint p_a;                  // (2) Time of last priority update
-    TimePoint q_a;                  // (3) Time of most recent request
+    TimePoint i_a;
+    TimePoint p_a;
+    TimePoint q_a;
 
-    // For pending time tracking and time slice tracking
     std::chrono::microseconds accumulated_pending_time{0};
-    TimePoint last_pending_start;   // When it entered pending state
+    TimePoint last_pending_start;
 
     std::chrono::microseconds time_slice_used{0};
-    TimePoint last_resume_time;     // When it started running
+    TimePoint last_resume_time;
 };
 
 class MultiLevelFeedbackQueuePolicy : public Policy {
@@ -35,16 +34,13 @@ class MultiLevelFeedbackQueuePolicy : public Policy {
     void RecvHint(std::shared_ptr<const Hint> hint) override;
 
 private:
-    std::map<XQueueHandle, MLFQNode> mlfq_infos_;
+    std::map<PID, MLFQNode> mlfq_infos_;
 
-    // Configuration
-    const int max_priority_ = 3;    // Levels 0, 1, 2, 3
-    const std::chrono::microseconds recovery_threshold_{100000}; // 100ms
-    const std::chrono::microseconds default_tick_{5000}; // 5ms scheduling tick
+    const int max_priority_ = 3;
+    const std::chrono::microseconds recovery_threshold_{100000};
+    const std::chrono::microseconds default_tick_{5000};
 
-    // Get time slice based on priority level
     std::chrono::microseconds get_time_slice(int prio) const {
-        // e.g., prio 0: 10ms, prio 1: 20ms, prio 2: 40ms, prio 3: 80ms
         return std::chrono::microseconds(10000 * (1 << prio));
     }
 };

sched/include/xsched/sched/policy/up.h

@@ -22,12 +22,14 @@ class UtilizationPartitionPolicy : public Policy
 
 private:
     std::chrono::microseconds GetBudget(Utilization util);
+    bool ProcessReady(PID pid, const Status &status);
     void SwitchToAny(const Status &status);
-    void SwitchTo(XQueueHandle handle, Utilization util, const Status &status);
+    void SwitchProcess(PID pid, Utilization util, const Status &status);
 
-    XQueueHandle cur_running_ = 0;
+    PID cur_running_ = 0;
     std::chrono::system_clock::time_point cur_end_;
-    std::map<XQueueHandle, Utilization> utils_;
+    std::map<PID, Utilization> utils_;
+    std::map<XQueueHandle, PID> handle_to_pid_;
     std::chrono::microseconds timeslice_ = 
         std::chrono::microseconds(TIMESLICE_DEFAULT);
 };

sched/src/policy/cfs.cpp

@@ -1,6 +1,7 @@
 #include <cmath>
 #include <limits>
 #include <map>
+#include <set>
 
 #include "xsched/utils/xassert.h"
 #include "xsched/sched/policy/cfs.h"
@@ -12,30 +13,35 @@ void CompletelyFairSchedulerPolicy::Sched(const Status &status)
     auto now = std::chrono::system_clock::now();
     bool has_ready_tasks = false;
 
-    // update vruntime for running queues
     for (auto &st : status.xqueue_status) {
+        PID pid = st.second->pid;
         XQueueHandle handle = st.second->handle;
-        auto it = cfs_infos_.find(handle);
+
+        auto pit = pending_hints_.find(handle);
+        if (pit != pending_hints_.end()) {
+            cfs_infos_[pid].priority = pit->second.first;
+            cfs_infos_[pid].weight = pit->second.second;
+            pending_hints_.erase(pit);
+        }
+
+        auto it = cfs_infos_.find(pid);
         if (it != cfs_infos_.end() && it->second.is_running) {
             auto delta_us = std::chrono::duration_cast<std::chrono::microseconds>(now - it->second.last_resume_time).count();
-            // Update virtual time: real time * (base weight / current task weight)
             it->second.vruntime += delta_us * (1024.0 / it->second.weight);
         }
     }
 
-    // Find the ready task with the minimum vruntime on each physical GPU
-    std::map<XDevice, XQueueHandle> min_vruntime_handles;
+    std::map<XDevice, PID> min_vruntime_pids;
     std::map<XDevice, double> min_vruntimes;
 
-    // find the current min_vruntime for each device among existing tasks
     for (auto &st : status.xqueue_status) {
         if (!st.second->ready) continue;
         has_ready_tasks = true;
 
         XDevice device = st.second->device;
-        XQueueHandle handle = st.second->handle;
+        PID pid = st.second->pid;
 
-        auto it = cfs_infos_.find(handle);
+        auto it = cfs_infos_.find(pid);
         if (it != cfs_infos_.end()) {
             double current_vruntime = it->second.vruntime;
             if (min_vruntimes.find(device) == min_vruntimes.end() || current_vruntime < min_vruntimes[device]) {
@@ -44,62 +50,66 @@ void CompletelyFairSchedulerPolicy::Sched(const Status &status)
         }
     }
 
-    // handle initialization for new tasks and select the final task to run
     for (auto &st : status.xqueue_status) {
         if (!st.second->ready) continue;
 
         XDevice device = st.second->device;
-        XQueueHandle handle = st.second->handle;
+        PID pid = st.second->pid;
 
-        // If this is the first time seeing this queue, initialize its CFS info
-        if (cfs_infos_.find(handle) == cfs_infos_.end()) {
-            cfs_infos_[handle] = CFSNode();
-            cfs_infos_[handle].last_resume_time = now;
-            // Inherit the minimum vruntime of the current device to prevent new tasks from starving old tasks
+        if (cfs_infos_.find(pid) == cfs_infos_.end()) {
+            cfs_infos_[pid] = CFSNode();
+            cfs_infos_[pid].last_resume_time = now;
             if (min_vruntimes.find(device) != min_vruntimes.end()) {
-                cfs_infos_[handle].vruntime = min_vruntimes[device];
+                cfs_infos_[pid].vruntime = min_vruntimes[device];
             } else {
-                cfs_infos_[handle].vruntime = 0.0;
+                cfs_infos_[pid].vruntime = 0.0;
             }
-            // Ensure min_vruntimes contains the vruntime of the new task (mainly for cases where the device has no old tasks)
             if (min_vruntimes.find(device) == min_vruntimes.end()) {
-                min_vruntimes[device] = cfs_infos_[handle].vruntime;
+                min_vruntimes[device] = cfs_infos_[pid].vruntime;
             }
         }
 
-        double current_vruntime = cfs_infos_[handle].vruntime;
+        double current_vruntime = cfs_infos_[pid].vruntime;
 
-        // Find the final minimum value and corresponding handle on this device
-        if (min_vruntime_handles.find(device) == min_vruntime_handles.end() || current_vruntime < cfs_infos_[min_vruntime_handles[device]].vruntime) {
-            min_vruntime_handles[device] = handle;
+        if (min_vruntime_pids.find(device) == min_vruntime_pids.end() ||
+            current_vruntime < cfs_infos_[min_vruntime_pids[device]].vruntime) {
+            min_vruntime_pids[device] = pid;
         }
     }
 
-    // Resume the task with the minimum vruntime, Suspend all others
     for (auto &st : status.xqueue_status) {
-        XDevice device = st.second->device;
-        XQueueHandle handle = st.second->handle;
-
+        PID pid = st.second->pid;
         if (!st.second->ready) {
-            cfs_infos_[handle].is_running = false;
+            cfs_infos_[pid].is_running = false;
             continue;
         }
+    }
 
-        if (min_vruntime_handles[device] == handle) {
-            if (!cfs_infos_[handle].is_running) {
-                this->Resume(handle);
-                cfs_infos_[handle].is_running = true;
-                cfs_infos_[handle].last_resume_time = now;
-            } else {
-                cfs_infos_[handle].last_resume_time = now;
-            }
+    std::set<PID> running_pids;
+    for (const auto &pair : min_vruntime_pids) {
+        PID best_pid = pair.second;
+        auto &node = cfs_infos_[best_pid];
+        if (!node.is_running) {
+            node.is_running = true;
+            node.last_resume_time = now;
         } else {
-            this->Suspend(handle);
-            cfs_infos_[handle].is_running = false;
+            node.last_resume_time = now;
+        }
+        running_pids.insert(best_pid);
+    }
+
+    for (auto &st : status.xqueue_status) {
+        PID pid = st.second->pid;
+        if (!st.second->ready) continue;
+
+        if (running_pids.count(pid)) {
+            this->Resume(st.first);
+        } else {
+            this->Suspend(st.first);
+            cfs_infos_[pid].is_running = false;
         }
     }
 
-    // force a new scheduling round after the time slice
     if (has_ready_tasks) {
         this->AddTimer(now + time_slice_);
     }
@@ -111,14 +121,10 @@ void CompletelyFairSchedulerPolicy::RecvHint(std::shared_ptr<const Hint> hint)
     auto h = std::dynamic_pointer_cast<const PriorityHint>(hint);
     if (h == nullptr) return;
 
-    XQueueHandle handle = h->Handle();
     Priority prio = h->Prio();
-
-    // Calculate weight: assuming base is 1024. For each priority increase, weight increases by 20%
     double weight = 1024.0 * std::pow(1.2, prio);
 
-    cfs_infos_[handle].priority = prio;
-    cfs_infos_[handle].weight = weight;
+    pending_hints_[h->Handle()] = {prio, weight};
 
-    XINFO("CFS: set priority %d (weight %.2f) for XQueue 0x" FMT_64X, prio, weight, handle);
+    XINFO("CFS: set priority %d (weight %.2f) for XQueue 0x" FMT_64X, prio, weight, h->Handle());
 }

sched/src/policy/kedf.cpp

@@ -9,44 +9,69 @@ using namespace xsched::sched;
 
 void KEarliestDeadlineFirstPolicy::Sched(const Status &status)
 {
-    std::vector<DeadlineEntry> ddls;
-    ddls.reserve(status.xqueue_status.size());
+    std::vector<ProcessDeadlineEntry> ddls;
+    ddls.reserve(status.process_status.size());
 
-    // calculate the deadline of each xqueue
-    for (auto &status : status.xqueue_status) {
-        XQueueHandle handle = status.second->handle;
+    for (auto &process : status.process_status) {
+        PID pid = process.first;
         auto ddl = (std::chrono::system_clock::time_point::max)();
-        if (!status.second->ready) {
-            ddls.emplace_back(DeadlineEntry{.xqueue=handle,.deadline=ddl});
-            continue;
+
+        bool any_ready = false;
+        for (auto handle : process.second->running_xqueues) {
+            auto xq_it = status.xqueue_status.find(handle);
+            if (xq_it == status.xqueue_status.end()) continue;
+            if (!xq_it->second->ready) continue;
+            any_ready = true;
+            auto d_it = deadlines_.find(handle);
+            if (d_it == deadlines_.end()) continue;
+            auto xq_ddl = xq_it->second->ready_time + std::chrono::microseconds(d_it->second);
+            if (xq_ddl < ddl) ddl = xq_ddl;
+        }
+        for (auto handle : process.second->suspended_xqueues) {
+            auto xq_it = status.xqueue_status.find(handle);
+            if (xq_it == status.xqueue_status.end()) continue;
+            if (!xq_it->second->ready) continue;
+            any_ready = true;
+            auto d_it = deadlines_.find(handle);
+            if (d_it == deadlines_.end()) continue;
+            auto xq_ddl = xq_it->second->ready_time + std::chrono::microseconds(d_it->second);
+            if (xq_ddl < ddl) ddl = xq_ddl;
         }
 
-        auto it = deadlines_.find(handle);
-        if (it == deadlines_.end()) {
-            ddls.emplace_back(DeadlineEntry{.xqueue=handle,.deadline=ddl});
+        if (!any_ready) {
+            ddls.emplace_back(ProcessDeadlineEntry{.pid=pid,.deadline=(std::chrono::system_clock::time_point::max)()});
             continue;
         }
 
-        ddl = status.second->ready_time + std::chrono::microseconds(it->second);
-        ddls.emplace_back(DeadlineEntry{.xqueue=handle,.deadline=ddl});
+        ddls.emplace_back(ProcessDeadlineEntry{.pid=pid,.deadline=ddl});
     }
 
-    // sort the xqueues by deadline, from the earliest to the latest
-    std::sort(ddls.begin(), ddls.end(), [](const DeadlineEntry &a, const DeadlineEntry &b) {
+    std::sort(ddls.begin(), ddls.end(), [](const ProcessDeadlineEntry &a, const ProcessDeadlineEntry &b) {
         return a.deadline < b.deadline;
     });
 
-    // resume the first k_ xqueues
     for (size_t i = 0; i < k_ && i < ddls.size(); ++i) {
-        this->Resume(ddls[i].xqueue);
+        SwitchProcess(ddls[i].pid, status);
     }
-
-    // suspend all other xqueues
     for (size_t i = k_; i < ddls.size(); ++i) {
-        this->Suspend(ddls[i].xqueue);
+        const auto it = status.process_status.find(ddls[i].pid);
+        if (it == status.process_status.end()) continue;
+        std::vector<XQueueHandle> running;
+        for (auto xq : it->second->running_xqueues) running.push_back(xq);
+        for (auto xq : running) this->Suspend(xq);
     }
 }
 
+void KEarliestDeadlineFirstPolicy::SwitchProcess(PID pid, const Status &status)
+{
+    const auto it = status.process_status.find(pid);
+    if (it == status.process_status.end()) return;
+
+    std::vector<XQueueHandle> suspended;
+    for (auto xq : it->second->suspended_xqueues) suspended.push_back(xq);
+    for (auto xq : suspended) this->Resume(xq);
+}
+
 void KEarliestDeadlineFirstPolicy::RecvHint(std::shared_ptr<const Hint> hint)
 {
     switch (hint->Type())