1 files changed, 67 insertions, 3 deletions

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