6 files changed, 493 insertions, 35 deletions

diff --git a/kernel/lockdep.c b/kernel/lockdep.c
index f2852a510232..ebff2cf715c5 100644
--- a/kernel/lockdep.c
+++ b/kernel/lockdep.c
@@ -530,7 +530,7 @@ static void print_lock(struct held_lock *hlock)
         print_ip_sym(hlock->acquire_ip);
 }
 
-static void lockdep_print_held_locks(struct task_struct *curr)
+void lockdep_print_held_locks(struct task_struct *curr)
 {
         int i, depth = curr->lockdep_depth;
 
@@ -546,6 +546,7 @@ static void lockdep_print_held_locks(struct task_struct *curr)
                 print_lock(curr->held_locks + i);
         }
 }
+EXPORT_SYMBOL(lockdep_print_held_locks);
 
 static void print_kernel_version(void)
 {
diff --git a/kernel/mutex.c b/kernel/mutex.c
index 200407c1502f..435685ecd068 100644
--- a/kernel/mutex.c
+++ b/kernel/mutex.c
@@ -496,3 +496,144 @@ int atomic_dec_and_mutex_lock(atomic_t *cnt, struct mutex *lock)
         return 1;
 }
 EXPORT_SYMBOL(atomic_dec_and_mutex_lock);
+
+
+
+
+
+
+
+
+//__mutex_lock_common(lock, TASK_UNINTERRUPTIBLE, 0, _RET_IP_);
+
+void mutex_lock_sfx(struct mutex *lock,
+                                   side_effect_t pre, unsigned long pre_arg,
+                                   side_effect_t post, unsigned long post_arg)
+{
+        long state = TASK_UNINTERRUPTIBLE;
+        unsigned int subclass = 0;
+        unsigned long ip = _RET_IP_;
+        
+
+        struct task_struct *task = current;
+        struct mutex_waiter waiter;
+        unsigned long flags;
+        
+        preempt_disable();
+        mutex_acquire(&lock->dep_map, subclass, 0, ip);
+
+        spin_lock_mutex(&lock->wait_lock, flags);
+        
+        if(pre)
+        {
+                if(unlikely(pre(pre_arg)))
+                {
+                        // this will fuck with lockdep's CONFIG_PROVE_LOCKING...
+                        spin_unlock_mutex(&lock->wait_lock, flags);
+                        preempt_enable();
+                        return;
+                }
+        }
+
+        debug_mutex_lock_common(lock, &waiter);
+        debug_mutex_add_waiter(lock, &waiter, task_thread_info(task));
+        
+        /* add waiting tasks to the end of the waitqueue (FIFO): */
+        list_add_tail(&waiter.list, &lock->wait_list);
+        waiter.task = task;
+        
+        if (atomic_xchg(&lock->count, -1) == 1)
+                goto done;
+        
+        lock_contended(&lock->dep_map, ip);
+        
+        for (;;) {
+                /*
+                 * Lets try to take the lock again - this is needed even if
+                 * we get here for the first time (shortly after failing to
+                 * acquire the lock), to make sure that we get a wakeup once
+                 * it's unlocked. Later on, if we sleep, this is the
+                 * operation that gives us the lock. We xchg it to -1, so
+                 * that when we release the lock, we properly wake up the
+                 * other waiters:
+                 */
+                if (atomic_xchg(&lock->count, -1) == 1)
+                        break;
+                
+                __set_task_state(task, state);
+                
+                /* didnt get the lock, go to sleep: */
+                spin_unlock_mutex(&lock->wait_lock, flags);
+                preempt_enable_no_resched();
+                schedule();
+                preempt_disable();
+                spin_lock_mutex(&lock->wait_lock, flags);
+        }
+        
+done:
+        lock_acquired(&lock->dep_map, ip);
+        /* got the lock - rejoice! */
+        mutex_remove_waiter(lock, &waiter, current_thread_info());
+        mutex_set_owner(lock);
+        
+        /* set it to 0 if there are no waiters left: */
+        if (likely(list_empty(&lock->wait_list)))
+                atomic_set(&lock->count, 0);
+        
+        if(post)
+                post(post_arg); 
+        
+        spin_unlock_mutex(&lock->wait_lock, flags);
+        
+        debug_mutex_free_waiter(&waiter);
+        preempt_enable();
+        
+        //return 0;     
+}
+EXPORT_SYMBOL(mutex_lock_sfx);
+
+
+
+//__mutex_unlock_common_slowpath(lock_count, 1);
+
+void mutex_unlock_sfx(struct mutex *lock,
+                                        side_effect_t pre, unsigned long pre_arg,
+                                        side_effect_t post, unsigned long post_arg)
+{
+        //struct mutex *lock = container_of(lock_count, struct mutex, count);
+        unsigned long flags;
+        
+        spin_lock_mutex(&lock->wait_lock, flags);
+        
+        if(pre)
+                pre(pre_arg);
+        
+        //mutex_release(&lock->dep_map, nested, _RET_IP_);
+        mutex_release(&lock->dep_map, 1, _RET_IP_);
+        debug_mutex_unlock(lock);
+        
+        /*
+         * some architectures leave the lock unlocked in the fastpath failure
+         * case, others need to leave it locked. In the later case we have to
+         * unlock it here
+         */
+        if (__mutex_slowpath_needs_to_unlock())
+                atomic_set(&lock->count, 1);
+        
+        if (!list_empty(&lock->wait_list)) {
+                /* get the first entry from the wait-list: */
+                struct mutex_waiter *waiter =
+                list_entry(lock->wait_list.next,
+                                   struct mutex_waiter, list);
+                
+                debug_mutex_wake_waiter(lock, waiter);
+                
+                wake_up_process(waiter->task);
+        }
+        
+        if(post)
+                post(post_arg);
+        
+        spin_unlock_mutex(&lock->wait_lock, flags);     
+}
+EXPORT_SYMBOL(mutex_unlock_sfx);
diff --git a/kernel/sched.c b/kernel/sched.c
index c5d775079027..3162605ffc91 100644
--- a/kernel/sched.c
+++ b/kernel/sched.c
@@ -82,6 +82,10 @@
 #include <litmus/sched_trace.h>
 #include <litmus/trace.h>
 
+#ifdef CONFIG_LITMUS_SOFTIRQD
+#include <litmus/litmus_softirq.h>
+#endif
+
 static void litmus_tick(struct rq*, struct task_struct*);
 
 #define CREATE_TRACE_POINTS
@@ -3789,6 +3793,7 @@ pick_next_task(struct rq *rq)
         }
 }
 
+
 /*
  * schedule() is the main scheduler function.
  */
@@ -3807,6 +3812,10 @@ need_resched:
         rcu_note_context_switch(cpu);
         prev = rq->curr;
 
+#ifdef CONFIG_LITMUS_SOFTIRQD
+        release_klitirqd_lock(prev);
+#endif
+        
         release_kernel_lock(prev);
 need_resched_nonpreemptible:
         TS_SCHED_START;
@@ -3882,15 +3891,20 @@ need_resched_nonpreemptible:
 
         if (sched_state_validate_switch() || unlikely(reacquire_kernel_lock(prev)))
                 goto need_resched_nonpreemptible;
-
+        
         preempt_enable_no_resched();
+
         if (need_resched())
                 goto need_resched;
 
+        reacquire_klitirqd_lock(prev);
+        
         srp_ceiling_block();
 }
 EXPORT_SYMBOL(schedule);
 
+
+
 #ifdef CONFIG_MUTEX_SPIN_ON_OWNER
 /*
  * Look out! "owner" is an entirely speculative pointer
@@ -4051,6 +4065,7 @@ static void __wake_up_common(wait_queue_head_t *q, unsigned int mode,
         }
 }
 
+
 /**
  * __wake_up - wake up threads blocked on a waitqueue.
  * @q: the waitqueue
@@ -4236,6 +4251,12 @@ void __sched wait_for_completion(struct completion *x)
 }
 EXPORT_SYMBOL(wait_for_completion);
 
+void __sched __wait_for_completion_locked(struct completion *x)
+{
+        do_wait_for_common(x, MAX_SCHEDULE_TIMEOUT, TASK_UNINTERRUPTIBLE);
+}
+EXPORT_SYMBOL(__wait_for_completion_locked);
+
 /**
  * wait_for_completion_timeout: - waits for completion of a task (w/timeout)
  * @x:  holds the state of this particular completion
diff --git a/kernel/semaphore.c b/kernel/semaphore.c
index 94a62c0d4ade..c947a046a6d7 100644
--- a/kernel/semaphore.c
+++ b/kernel/semaphore.c
@@ -33,11 +33,11 @@
 #include <linux/spinlock.h>
 #include <linux/ftrace.h>
 
-static noinline void __down(struct semaphore *sem);
+noinline void __down(struct semaphore *sem);
 static noinline int __down_interruptible(struct semaphore *sem);
 static noinline int __down_killable(struct semaphore *sem);
 static noinline int __down_timeout(struct semaphore *sem, long jiffies);
-static noinline void __up(struct semaphore *sem);
+noinline void __up(struct semaphore *sem);
 
 /**
  * down - acquire the semaphore
@@ -190,11 +190,13 @@ EXPORT_SYMBOL(up);
 
 /* Functions for the contended case */
 
+/*
 struct semaphore_waiter {
         struct list_head list;
         struct task_struct *task;
         int up;
 };
+ */
 
 /*
  * Because this function is inlined, the 'state' parameter will be
@@ -233,10 +235,12 @@ static inline int __sched __down_common(struct semaphore *sem, long state,
         return -EINTR;
 }
 
-static noinline void __sched __down(struct semaphore *sem)
+noinline void __sched __down(struct semaphore *sem)
 {
         __down_common(sem, TASK_UNINTERRUPTIBLE, MAX_SCHEDULE_TIMEOUT);
 }
+EXPORT_SYMBOL(__down);
+
 
 static noinline int __sched __down_interruptible(struct semaphore *sem)
 {
@@ -253,7 +257,7 @@ static noinline int __sched __down_timeout(struct semaphore *sem, long jiffies)
         return __down_common(sem, TASK_UNINTERRUPTIBLE, jiffies);
 }
 
-static noinline void __sched __up(struct semaphore *sem)
+noinline void __sched __up(struct semaphore *sem)
 {
         struct semaphore_waiter *waiter = list_first_entry(&sem->wait_list,
                                                 struct semaphore_waiter, list);
@@ -261,3 +265,4 @@ static noinline void __sched __up(struct semaphore *sem)
         waiter->up = 1;
         wake_up_process(waiter->task);
 }
+EXPORT_SYMBOL(__up);
\ No newline at end of file
diff --git a/kernel/softirq.c b/kernel/softirq.c
index 07b4f1b1a73a..be4b8fab3637 100644
--- a/kernel/softirq.c
+++ b/kernel/softirq.c
@@ -29,6 +29,14 @@
 #include <trace/events/irq.h>
 
 #include <asm/irq.h>
+
+#include <litmus/litmus.h>
+#include <litmus/sched_trace.h>
+
+#ifdef CONFIG_LITMUS_NVIDIA
+#include <litmus/nvidia_info.h>
+#endif
+
 /*
    - No shared variables, all the data are CPU local.
    - If a softirq needs serialization, let it serialize itself
@@ -54,7 +62,7 @@ EXPORT_SYMBOL(irq_stat);
 
 static struct softirq_action softirq_vec[NR_SOFTIRQS] __cacheline_aligned_in_smp;
 
-static DEFINE_PER_CPU(struct task_struct *, ksoftirqd);
+static DEFINE_PER_CPU(struct task_struct *, ksoftirqd) = NULL;
 
 char *softirq_to_name[NR_SOFTIRQS] = {
         "HI", "TIMER", "NET_TX", "NET_RX", "BLOCK", "BLOCK_IOPOLL",
@@ -177,6 +185,7 @@ void local_bh_enable_ip(unsigned long ip)
 }
 EXPORT_SYMBOL(local_bh_enable_ip);
 
+
 /*
  * We restart softirq processing MAX_SOFTIRQ_RESTART times,
  * and we fall back to softirqd after that.
@@ -187,34 +196,30 @@ EXPORT_SYMBOL(local_bh_enable_ip);
  * should not be able to lock up the box.
  */
 #define MAX_SOFTIRQ_RESTART 10
-
-asmlinkage void __do_softirq(void)
+static void ____do_softirq(void)
 {
-        struct softirq_action *h;
         __u32 pending;
-        int max_restart = MAX_SOFTIRQ_RESTART;
+        
+        struct softirq_action *h;
         int cpu;
-
+        
         pending = local_softirq_pending();
+        
         account_system_vtime(current);
-
-        __local_bh_disable((unsigned long)__builtin_return_address(0));
-        lockdep_softirq_enter();
-
+        
         cpu = smp_processor_id();
-restart:
-        /* Reset the pending bitmask before enabling irqs */
-        set_softirq_pending(0);
 
+        set_softirq_pending(0);
+        
         local_irq_enable();
-
+        
         h = softirq_vec;
-
+        
         do {
                 if (pending & 1) {
                         int prev_count = preempt_count();
                         kstat_incr_softirqs_this_cpu(h - softirq_vec);
-
+                        
                         trace_softirq_entry(h, softirq_vec);
                         h->action(h);
                         trace_softirq_exit(h, softirq_vec);
@@ -226,26 +231,70 @@ restart:
                                        h->action, prev_count, preempt_count());
                                 preempt_count() = prev_count;
                         }
-
+                        
                         rcu_bh_qs(cpu);
                 }
                 h++;
                 pending >>= 1;
         } while (pending);
-
+        
         local_irq_disable();
+}
+
+static void ___do_softirq(void)
+{
+        __u32 pending;
+
+        //struct softirq_action *h;
+        int max_restart = MAX_SOFTIRQ_RESTART;
+        //int cpu;
+
+        pending = local_softirq_pending();
+
+restart:
+        ____do_softirq();
 
         pending = local_softirq_pending();
         if (pending && --max_restart)
                 goto restart;
 
         if (pending)
+        {
                 wakeup_softirqd();
+        }
+}
 
+asmlinkage void __do_softirq(void)
+{
+#ifdef LITMUS_THREAD_ALL_SOFTIRQ
+        /* Skip straight to wakeup_softirqd() if we're using 
+         LITMUS_THREAD_ALL_SOFTIRQ (unless there's really high prio-stuff waiting.). */
+        struct task_struct *tsk = __get_cpu_var(ksoftirqd);
+        
+        if(tsk)
+        {
+                __u32 pending = local_softirq_pending();
+                const __u32 high_prio_softirq = (1<<HI_SOFTIRQ) | (1<<TIMER_SOFTIRQ) | (1<<HRTIMER_SOFTIRQ);
+                if(pending && !(pending & high_prio_softirq))
+                {
+                        wakeup_softirqd();
+                        return;
+                }
+        }
+#endif
+        
+        /*
+         * 'immediate' softirq execution:
+         */
+        __local_bh_disable((unsigned long)__builtin_return_address(0));
+        lockdep_softirq_enter();
+        
+        ___do_softirq();
+        
         lockdep_softirq_exit();
-
+        
         account_system_vtime(current);
-        _local_bh_enable();
+        _local_bh_enable();     
 }
 
 #ifndef __ARCH_HAS_DO_SOFTIRQ
@@ -357,8 +406,64 @@ struct tasklet_head
 static DEFINE_PER_CPU(struct tasklet_head, tasklet_vec);
 static DEFINE_PER_CPU(struct tasklet_head, tasklet_hi_vec);
 
+
 void __tasklet_schedule(struct tasklet_struct *t)
 {
+#ifdef CONFIG_LITMUS_NVIDIA
+        if(is_nvidia_func(t->func))
+        {
+                u32 nvidia_device = get_tasklet_nv_device_num(t);       
+                //              TRACE("%s: Handling NVIDIA tasklet for device\t%u\tat\t%llu\n",
+                //                        __FUNCTION__, nvidia_device,litmus_clock());
+
+                unsigned long flags;
+                struct task_struct* device_owner;
+
+                lock_nv_registry(nvidia_device, &flags);
+
+                device_owner = get_nv_device_owner(nvidia_device);
+
+                if(device_owner==NULL)
+                {
+                        t->owner = NULL;
+                }
+                else
+                {
+                        if(is_realtime(device_owner))
+                        {
+                                TRACE("%s: Handling NVIDIA tasklet for device %u at %llu\n",
+                                          __FUNCTION__, nvidia_device,litmus_clock());                          
+                                TRACE("%s: the owner task %d of NVIDIA Device %u is RT-task\n",
+                                          __FUNCTION__,device_owner->pid,nvidia_device);
+
+                                t->owner = device_owner;
+                                sched_trace_tasklet_release(t->owner);
+                                if(likely(_litmus_tasklet_schedule(t,nvidia_device)))
+                                {
+                                        unlock_nv_registry(nvidia_device, &flags);
+                                        return;
+                                }
+                                else
+                                {
+                                        t->owner = NULL; /* fall through to normal scheduling */
+                                }
+                        }
+                        else
+                        {
+                                t->owner = NULL;
+                        }
+                }
+                unlock_nv_registry(nvidia_device, &flags);
+        }
+#endif
+
+        ___tasklet_schedule(t);
+}
+EXPORT_SYMBOL(__tasklet_schedule);
+
+
+void ___tasklet_schedule(struct tasklet_struct *t)
+{
         unsigned long flags;
 
         local_irq_save(flags);
@@ -368,11 +473,65 @@ void __tasklet_schedule(struct tasklet_struct *t)
         raise_softirq_irqoff(TASKLET_SOFTIRQ);
         local_irq_restore(flags);
 }
+EXPORT_SYMBOL(___tasklet_schedule);
 
-EXPORT_SYMBOL(__tasklet_schedule);
 
 void __tasklet_hi_schedule(struct tasklet_struct *t)
 {
+#ifdef CONFIG_LITMUS_NVIDIA
+        if(is_nvidia_func(t->func))
+        {       
+                u32 nvidia_device = get_tasklet_nv_device_num(t);
+                //              TRACE("%s: Handling NVIDIA tasklet for device\t%u\tat\t%llu\n",
+                //                        __FUNCTION__, nvidia_device,litmus_clock());
+
+                unsigned long flags;
+                struct task_struct* device_owner;
+                
+                lock_nv_registry(nvidia_device, &flags);
+                
+                device_owner = get_nv_device_owner(nvidia_device);
+
+                if(device_owner==NULL) 
+                {
+                        t->owner = NULL;
+                }
+                else
+                {
+                        if( is_realtime(device_owner))
+                        {
+                                TRACE("%s: Handling NVIDIA tasklet for device %u\tat %llu\n",
+                                          __FUNCTION__, nvidia_device,litmus_clock());                          
+                                TRACE("%s: the owner task %d of NVIDIA Device %u is RT-task\n",
+                                          __FUNCTION__,device_owner->pid,nvidia_device);
+                                
+                                t->owner = device_owner;
+                                sched_trace_tasklet_release(t->owner);
+                                if(likely(_litmus_tasklet_hi_schedule(t,nvidia_device)))
+                                {
+                                        unlock_nv_registry(nvidia_device, &flags);
+                                        return;
+                                }
+                                else
+                                {
+                                        t->owner = NULL; /* fall through to normal scheduling */
+                                }
+                        }
+                        else
+                        {
+                                t->owner = NULL;
+                        }
+                }
+                unlock_nv_registry(nvidia_device, &flags);
+        }
+#endif
+
+        ___tasklet_hi_schedule(t);
+}
+EXPORT_SYMBOL(__tasklet_hi_schedule);
+
+void ___tasklet_hi_schedule(struct tasklet_struct* t)
+{
         unsigned long flags;
 
         local_irq_save(flags);
@@ -382,19 +541,72 @@ void __tasklet_hi_schedule(struct tasklet_struct *t)
         raise_softirq_irqoff(HI_SOFTIRQ);
         local_irq_restore(flags);
 }
-
-EXPORT_SYMBOL(__tasklet_hi_schedule);
+EXPORT_SYMBOL(___tasklet_hi_schedule);
 
 void __tasklet_hi_schedule_first(struct tasklet_struct *t)
 {
         BUG_ON(!irqs_disabled());
+#ifdef CONFIG_LITMUS_NVIDIA     
+        if(is_nvidia_func(t->func))
+        {       
+                u32 nvidia_device = get_tasklet_nv_device_num(t);
+                //              TRACE("%s: Handling NVIDIA tasklet for device\t%u\tat\t%llu\n",
+                //                        __FUNCTION__, nvidia_device,litmus_clock());
+                unsigned long flags;
+                struct task_struct* device_owner;
+                
+                lock_nv_registry(nvidia_device, &flags);
+
+                device_owner = get_nv_device_owner(nvidia_device);
+
+                if(device_owner==NULL)
+                {
+                        t->owner = NULL;
+                }
+                else
+                {
+                        if(is_realtime(device_owner))
+                        {
+                                TRACE("%s: Handling NVIDIA tasklet for device %u at %llu\n",
+                                          __FUNCTION__, nvidia_device,litmus_clock());
+                                
+                                TRACE("%s: the owner task %d of NVIDIA Device %u is RT-task\n",
+                                          __FUNCTION__,device_owner->pid,nvidia_device);
+                                
+                                t->owner = device_owner;
+                                sched_trace_tasklet_release(t->owner);
+                                if(likely(_litmus_tasklet_hi_schedule_first(t,nvidia_device)))
+                                {
+                                        unlock_nv_registry(nvidia_device, &flags);
+                                        return;
+                                }
+                                else
+                                {
+                                        t->owner = NULL; /* fall through to normal scheduling */
+                                }
+                        }
+                        else
+                        {
+                                t->owner = NULL;
+                        }
+                }
+                unlock_nv_registry(nvidia_device, &flags);
+        }
+#endif
+
+        ___tasklet_hi_schedule_first(t);
+}
+EXPORT_SYMBOL(__tasklet_hi_schedule_first);
+
+void ___tasklet_hi_schedule_first(struct tasklet_struct* t)
+{
+        BUG_ON(!irqs_disabled());
 
         t->next = __get_cpu_var(tasklet_hi_vec).head;
         __get_cpu_var(tasklet_hi_vec).head = t;
         __raise_softirq_irqoff(HI_SOFTIRQ);
 }
-
-EXPORT_SYMBOL(__tasklet_hi_schedule_first);
+EXPORT_SYMBOL(___tasklet_hi_schedule_first);
 
 static void tasklet_action(struct softirq_action *a)
 {
@@ -450,6 +662,7 @@ static void tasklet_hi_action(struct softirq_action *a)
                         if (!atomic_read(&t->count)) {
                                 if (!test_and_clear_bit(TASKLET_STATE_SCHED, &t->state))
                                         BUG();
+
                                 t->func(t->data);
                                 tasklet_unlock(t);
                                 continue;
@@ -473,8 +686,13 @@ void tasklet_init(struct tasklet_struct *t,
         t->next = NULL;
         t->state = 0;
         atomic_set(&t->count, 0);
+
         t->func = func;
         t->data = data;
+
+#ifdef CONFIG_LITMUS_SOFTIRQD
+        t->owner = NULL;
+#endif
 }
 
 EXPORT_SYMBOL(tasklet_init);
@@ -489,6 +707,7 @@ void tasklet_kill(struct tasklet_struct *t)
                         yield();
                 } while (test_bit(TASKLET_STATE_SCHED, &t->state));
         }
+
         tasklet_unlock_wait(t);
         clear_bit(TASKLET_STATE_SCHED, &t->state);
 }
@@ -694,6 +913,8 @@ void __init softirq_init(void)
 
 static int run_ksoftirqd(void * __bind_cpu)
 {
+        unsigned long flags;
+        
         set_current_state(TASK_INTERRUPTIBLE);
 
         while (!kthread_should_stop()) {
@@ -712,7 +933,11 @@ static int run_ksoftirqd(void * __bind_cpu)
                            don't process */
                         if (cpu_is_offline((long)__bind_cpu))
                                 goto wait_to_die;
-                        do_softirq();
+                        
+                        local_irq_save(flags);
+                        ____do_softirq();
+                        local_irq_restore(flags);
+                        
                         preempt_enable_no_resched();
                         cond_resched();
                         preempt_disable();
@@ -760,6 +985,7 @@ void tasklet_kill_immediate(struct tasklet_struct *t, unsigned int cpu)
         for (i = &per_cpu(tasklet_vec, cpu).head; *i; i = &(*i)->next) {
                 if (*i == t) {
                         *i = t->next;
+
                         /* If this was the tail element, move the tail ptr */
                         if (*i == NULL)
                                 per_cpu(tasklet_vec, cpu).tail = i;
diff --git a/kernel/workqueue.c b/kernel/workqueue.c
index f77afd939229..8139208eaee1 100644
--- a/kernel/workqueue.c
+++ b/kernel/workqueue.c
@@ -47,6 +47,13 @@
 
 #include "workqueue_sched.h"
 
+#ifdef CONFIG_LITMUS_NVIDIA
+#include <litmus/litmus.h>
+#include <litmus/sched_trace.h>
+#include <litmus/nvidia_info.h>
+#endif
+
+
 enum {
         /* global_cwq flags */
         GCWQ_MANAGE_WORKERS     = 1 << 0,       /* need to manage workers */
@@ -1010,9 +1017,7 @@ static void __queue_work(unsigned int cpu, struct workqueue_struct *wq,
                 work_flags |= WORK_STRUCT_DELAYED;
                 worklist = &cwq->delayed_works;
         }
-
         insert_work(cwq, work, worklist, work_flags);
-
         spin_unlock_irqrestore(&gcwq->lock, flags);
 }
 
@@ -2526,10 +2531,69 @@ EXPORT_SYMBOL(cancel_delayed_work_sync);
  */
 int schedule_work(struct work_struct *work)
 {
-        return queue_work(system_wq, work);
+#ifdef CONFIG_LITMUS_NVIDIA
+        if(is_nvidia_func(work->func))
+        {
+                u32 nvidiaDevice = get_work_nv_device_num(work);
+                
+                //1) Ask Litmus which task owns GPU <nvidiaDevice>. (API to be defined.)
+                unsigned long flags;
+                struct task_struct* device_owner;
+                
+                lock_nv_registry(nvidiaDevice, &flags);
+                
+                device_owner = get_nv_device_owner(nvidiaDevice);
+                
+                //2) If there is an owner, set work->owner to the owner's task struct.
+                if(device_owner==NULL) 
+                {
+                        work->owner = NULL;
+                        //TRACE("%s: the owner task of NVIDIA Device %u is NULL\n",__FUNCTION__,nvidiaDevice);
+                }
+                else
+                {
+                        if( is_realtime(device_owner))
+                        {
+                                TRACE("%s: Handling NVIDIA work for device\t%u\tat\t%llu\n",
+                                          __FUNCTION__, nvidiaDevice,litmus_clock());
+                                TRACE("%s: the owner task %d of NVIDIA Device %u is RT-task\n",
+                                          __FUNCTION__,
+                                          device_owner->pid,
+                                          nvidiaDevice);
+                                
+                                //3) Call litmus_schedule_work() and return (don't execute the rest
+                                //      of schedule_schedule()).
+                                work->owner = device_owner;
+                                sched_trace_work_release(work->owner);
+                                if(likely(litmus_schedule_work(work, nvidiaDevice)))
+                                {
+                                        unlock_nv_registry(nvidiaDevice, &flags);
+                                        return 1;
+                                }
+                                else
+                                {
+                                        work->owner = NULL; /* fall through to normal work scheduling */
+                                }
+                        }
+                        else
+                        {
+                                work->owner = NULL;
+                        }
+                }
+                unlock_nv_registry(nvidiaDevice, &flags);
+        }
+#endif
+
+        return(__schedule_work(work));
 }
 EXPORT_SYMBOL(schedule_work);
 
+int __schedule_work(struct work_struct* work)
+{
+        return queue_work(system_wq, work);
+}
+EXPORT_SYMBOL(__schedule_work);
+
 /*
  * schedule_work_on - put work task on a specific cpu
  * @cpu: cpu to put the work task on