5 files changed, 41 insertions, 12 deletions

diff --git a/include/litmus/gpu_affinity.h b/include/litmus/gpu_affinity.h
index f610f58b1f3b..197d4438d4f9 100644
--- a/include/litmus/gpu_affinity.h
+++ b/include/litmus/gpu_affinity.h
@@ -16,22 +16,33 @@ static inline void reset_gpu_tracker(struct task_struct* t)
 
 static inline void start_gpu_tracker(struct task_struct* t)
 {
-        t->rt_param.gpu_time_stamp = litmus_clock();
+        lt_t now = litmus_clock();
+
+        TRACE_TASK(t, "+++ starting gpu tracker @ %llu\n", now);
+//      WARN_ON(t->rt_param.gpu_time_stamp != 0); // already running!
+
+        if (likely(!t->rt_param.gpu_time_stamp))
+                t->rt_param.gpu_time_stamp = now;
 }
 
 static inline void stop_gpu_tracker(struct task_struct* t)
 {
         lt_t now = litmus_clock();
-        t->rt_param.accum_gpu_time += (now - t->rt_param.gpu_time_stamp);
-        t->rt_param.gpu_time_stamp = 0;
+
+        TRACE_TASK(t, "--- stopping gpu tracker @ %llu\n", now);
+//      WARN_ON(t->rt_param.gpu_time_stamp == 0); // already stopped!
+
+        if (likely(t->rt_param.gpu_time_stamp)) {
+                t->rt_param.accum_gpu_time += (now - t->rt_param.gpu_time_stamp);
+                t->rt_param.gpu_time_stamp = 0;
+        }
 }
 
 static inline lt_t get_gpu_time(struct task_struct* t)
 {
         lt_t accum = t->rt_param.accum_gpu_time;
-        if (t->rt_param.gpu_time_stamp != 0) {
+        if (t->rt_param.gpu_time_stamp != 0)
                 accum += (litmus_clock() - t->rt_param.gpu_time_stamp);
-        }
         return accum;
 }
 
diff --git a/litmus/gpu_affinity.c b/litmus/gpu_affinity.c
index e59814b998d3..b3415112eeb1 100644
--- a/litmus/gpu_affinity.c
+++ b/litmus/gpu_affinity.c
@@ -111,13 +111,13 @@ void update_gpu_estimate(struct task_struct *t, lt_t observed)
                 lower = est->avg - MIN(range, est->avg); // no underflow.
 
                 if (unlikely(observed < lower)) {
-                        TRACE_TASK(t, "Observation is too small: %llu\n", observed);
+                        TRACE_TASK(t, "Observation is too small. %llu < %llu (avg: %llu)\n", observed, lower, est->avg);
                         return;
                 }
 
                 upper = est->avg + range;
                 if (unlikely(observed > upper)) {
-                        TRACE_TASK(t, "Observation is too large: %llu\n", observed);
+                        TRACE_TASK(t, "Observation is too large: %llu > %llu (avg: %llu)\n", observed, upper, est->avg);
                         return;
                 }
         }
diff --git a/litmus/litmus.c b/litmus/litmus.c
index ceaa31405884..05221405ddc4 100644
--- a/litmus/litmus.c
+++ b/litmus/litmus.c
@@ -136,8 +136,7 @@ asmlinkage long sys_set_rt_task_param(pid_t pid, struct rt_task __user * param)
         {
                 printk(KERN_INFO "litmus: real-time task %d rejected "
                        "because task density > 1.0: exec_cost = %llu, deadline = %llu\n",
-                           tp.exec_cost, min(tp.relative_deadline, tp.period),
-                           pid);
+                           pid, tp.exec_cost, min(tp.relative_deadline, tp.period));
                 goto out_unlock;
         }
         if (tp.cls != RT_CLASS_HARD &&
@@ -583,13 +582,13 @@ long litmus_admit_task(struct task_struct* tsk)
                         min(get_rt_relative_deadline(tsk), get_rt_period(tsk)) ) {
                 printk("%s/%d "
                         "litmus admit: invalid task parameters "
-                        "(e = %lu, p = %lu, d = %lu)\n",
+                        "(e = %llu, p = %llu, d = %llu)\n",
                         tsk->comm, tsk->pid,
                         get_exec_cost(tsk), get_rt_period(tsk),
                         get_rt_relative_deadline(tsk));
                 TRACE_TASK(tsk,
                         "litmus admit: invalid task parameters "
-                        "(e = %lu, p = %lu, d = %lu)\n",
+                        "(e = %llu, p = %llu, d = %llu)\n",
                         get_exec_cost(tsk), get_rt_period(tsk),
                         get_rt_relative_deadline(tsk));
                 retval = -EINVAL;
diff --git a/litmus/locking.c b/litmus/locking.c
index 513b5d3844c8..c3f085c079f8 100644
--- a/litmus/locking.c
+++ b/litmus/locking.c
@@ -887,7 +887,11 @@ void suspend_for_lock(void)
 #if defined(CONFIG_LITMUS_AFFINITY_LOCKING) && defined(CONFIG_LITMUS_NVIDIA)
         // disable tracking
         if(tsk_rt(t)->held_gpus)
+        {
+                // tracking is actually stopped in schedule(), where it
+                // is also stopped upon preemption
                 tsk_rt(t)->suspend_gpu_tracker_on_block = 1;
+        }
 #endif
 
         schedule();
@@ -895,7 +899,10 @@ void suspend_for_lock(void)
 #if defined(CONFIG_LITMUS_AFFINITY_LOCKING) && defined(CONFIG_LITMUS_NVIDIA)
         // re-enable tracking
         if(tsk_rt(t)->held_gpus)
+        {
                 tsk_rt(t)->suspend_gpu_tracker_on_block = 0;
+//              start_gpu_tracker(t); // we have to turn tracking back on here
+        }
 #endif
 
 #if defined(CONFIG_REALTIME_AUX_TASKS) || defined(CONFIG_LITMUS_NVIDIA)
diff --git a/litmus/sched_litmus.c b/litmus/sched_litmus.c
index 3a594cd51a44..63f893c232af 100644
--- a/litmus/sched_litmus.c
+++ b/litmus/sched_litmus.c
@@ -5,6 +5,8 @@
 #include <litmus/sched_plugin.h>
 #include <litmus/preempt.h>
 
+#include <litmus/gpu_affinity.h>
+
 static void update_time_litmus(struct rq *rq, struct task_struct *p)
 {
         u64 delta = rq->clock - p->se.exec_start;
@@ -155,8 +157,18 @@ litmus_schedule(struct rq *rq, struct task_struct *prev)
                 next->rt_param.stack_in_use = rq->cpu;
                 next->se.exec_start = rq->clock;
 
-                if (is_realtime(next))
+                if (is_realtime(next)) {
                         budget_state_machine(next,on_scheduled);
+
+                        // turn gpu tracking back on if needed
+                        if(tsk_rt(next)->held_gpus) {
+                                // tracker might already be running if we're waking
+                                // up from using the GPU
+                                if(0 == tsk_rt(next)->gpu_time_stamp) {
+                                        start_gpu_tracker(next);
+                                }
+                        }
+                }
         }
 
         return next;