re-enable klmirqd for workqueues, and grace reboot

11 files changed, 388 insertions, 127 deletions

diff --git a/include/litmus/litmus_softirq.h b/include/litmus/litmus_softirq.h
index cfef08187464..e9590835dc19 100644
--- a/include/litmus/litmus_softirq.h
+++ b/include/litmus/litmus_softirq.h
@@ -153,14 +153,27 @@ static inline int _litmus_tasklet_hi_schedule_first(
 //////////////
 
 extern int __litmus_schedule_work(
-        struct work_struct* w,
+        struct work_struct *w,
         struct task_struct *klmirqd_thread);
 
 static inline int litmus_schedule_work(
-        struct work_struct* w,
+        struct work_struct *w,
+        struct task_struct *klmirqd_thread)
+{
+        int ret = 0;
+        if (!test_and_set_bit(WORK_STRUCT_PENDING_BIT, work_data_bits(w))) {
+                ret = __litmus_schedule_work(w, klmirqd_thread);
+        }
+        return(ret);
+}
+
+static inline int _litmus_schedule_work(
+        struct work_struct *w,
         struct task_struct *klmirqd_thread)
 {
         return(__litmus_schedule_work(w, klmirqd_thread));
 }
 
+
+
 #endif
diff --git a/include/litmus/nvidia_info.h b/include/litmus/nvidia_info.h
index 0ebe9c10c6a2..7db4a32af734 100644
--- a/include/litmus/nvidia_info.h
+++ b/include/litmus/nvidia_info.h
@@ -28,7 +28,15 @@ u32 get_work_nv_device_num(const struct work_struct *t);
 struct task_struct* get_nv_max_device_owner(u32 target_device_id);
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
-struct task_struct* get_nv_klmirqd_thread(u32 target_device_id);
+struct task_struct* get_and_lock_nvklmirqd_thread(u32 target_device_id, unsigned long* flags);
+void unlock_nvklmirqd_thread(u32 target_device_id, unsigned long* flags);
+struct task_struct* get_nvklmirqd_thread(u32 target_device_id);
+
+#if defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON) || defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED)
+struct task_struct* get_and_lock_nvklmworkqd_thread(u32 target_device_id, unsigned long* flags);
+void unlock_nvklmworkqd_thread(u32 target_device_id, unsigned long* flags);
+struct task_struct* get_nvklmworkqd_thread(u32 target_device_id);
+#endif
 #endif
 
 /* call when the GPU-holding task, t, blocks */
diff --git a/include/litmus/sched_trace.h b/include/litmus/sched_trace.h
index 2598cdf6088e..0d22c275dd72 100644
--- a/include/litmus/sched_trace.h
+++ b/include/litmus/sched_trace.h
@@ -103,7 +103,8 @@ struct st_tasklet_end_data {
 
 struct st_work_release_data {
         u64 when;
-        u64 __unused;
+        u32 device;
+        u32 __unused;
 } __attribute__((packed));
 
 struct st_work_begin_data {
@@ -408,8 +409,8 @@ feather_callback void do_sched_trace_migration(unsigned long id,
         SCHED_TRACE2(SCHED_TRACE_BASE_ID + 13, do_sched_trace_tasklet_end, t, flushed)
 
 
-#define sched_trace_work_release(t) \
-        SCHED_TRACE(SCHED_TRACE_BASE_ID + 14, do_sched_trace_work_release, t)
+#define sched_trace_work_release(t, d) \
+        SCHED_TRACE2(SCHED_TRACE_BASE_ID + 14, do_sched_trace_work_release, t, d)
 
 #define sched_trace_work_begin(t, e) \
         SCHED_TRACE2(SCHED_TRACE_BASE_ID + 15, do_sched_trace_work_begin, t, e)
diff --git a/kernel/softirq.c b/kernel/softirq.c
index 3ad0e7d9e782..ef70dbe92a97 100644
--- a/kernel/softirq.c
+++ b/kernel/softirq.c
@@ -455,23 +455,30 @@ typedef int (*klmirqd_tasklet_sched_t)(
 static int __klmirqd_nv_tasklet_schedule(struct tasklet_struct *t,
         klmirqd_tasklet_sched_t klmirqd_func)
 {
-        u32 nvidia_device = get_tasklet_nv_device_num(t);       
-        struct task_struct* klmirqd_th = get_nv_klmirqd_thread(nvidia_device);
+        /* TODO: MOVE THIS TO nvidia_info.c */
+        unsigned long flags;
+        u32 nvidia_device;
+        struct task_struct* klmirqd_th;
 
         BUG_ON(!klmirqd_func);
 
-        TRACE("Handling NVIDIA tasklet for device %u (klmirqd: %s/%d) at %llu\n",
-                nvidia_device,
-                (klmirqd_th) ? klmirqd_th->comm : "nil",
-                (klmirqd_th) ? klmirqd_th->pid : 0,
-                litmus_clock());
+        nvidia_device = get_tasklet_nv_device_num(t);
+        klmirqd_th = get_and_lock_nvklmirqd_thread(nvidia_device, &flags);
 
-        sched_trace_tasklet_release(NULL, nvidia_device);
+        if (likely(klmirqd_th)) {
+                TRACE("Handling NVIDIA tasklet for device %u (klmirqd: %s/%d) at %llu\n",
+                        nvidia_device,
+                        klmirqd_th->comm,
+                        klmirqd_th->pid,
+                        litmus_clock());
 
-        if (klmirqd_th && likely(klmirqd_func(t, klmirqd_th)))
-                return 1;
-        else
-                return 0;
+                sched_trace_tasklet_release(NULL, nvidia_device);
+                if(likely(klmirqd_func(t, klmirqd_th))) {
+                        unlock_nvklmirqd_thread(nvidia_device, &flags);
+                        return 1; /* success */
+                }
+        }
+        return 0;
 }
 #endif
 
diff --git a/kernel/workqueue.c b/kernel/workqueue.c
index 6b59d59ce3cf..9eb58729e1bb 100644
--- a/kernel/workqueue.c
+++ b/kernel/workqueue.c
@@ -44,10 +44,11 @@
 
 #include "workqueue_sched.h"
 
-#ifdef CONFIG_LITMUS_NVIDIA
+#if defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON) || defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED)
 #include <litmus/litmus.h>
-#include <litmus/sched_trace.h>
 #include <litmus/nvidia_info.h>
+#include <litmus/sched_trace.h>
+#include <litmus/trace.h>
 #endif
 
 
@@ -2679,6 +2680,32 @@ bool cancel_delayed_work_sync(struct delayed_work *dwork)
 }
 EXPORT_SYMBOL(cancel_delayed_work_sync);
 
+
+#if defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON) || defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED)
+static int __klmirqd_nv_work_schedule(struct work_struct *work)
+{
+        /* TODO: MOVE THIS TO nvidia_info.c */
+        unsigned long flags;
+        u32 nvidia_device = get_work_nv_device_num(work);
+        struct task_struct* klmirqd_th = get_and_lock_nvklmworkqd_thread(nvidia_device, &flags);
+
+        if (likely(klmirqd_th)) {
+                TRACE("Handling NVIDIA workq for device %u (klmirqd: %s/%d) at %llu\n",
+                        nvidia_device,
+                        klmirqd_th->comm,
+                        klmirqd_th->pid,
+                        litmus_clock());
+
+                sched_trace_work_release(NULL, nvidia_device);
+                if (likely(litmus_schedule_work(work, klmirqd_th))) {
+                        unlock_nvklmworkqd_thread(nvidia_device, &flags);
+                        return 1; /* success */
+                }
+        }
+        return 0;
+}
+#endif
+
 /**
  * schedule_work - put work task in global workqueue
  * @work: job to be done
@@ -2692,60 +2719,14 @@ EXPORT_SYMBOL(cancel_delayed_work_sync);
  */
 int schedule_work(struct work_struct *work)
 {
-#if 0
-#if defined(CONFIG_LITMUS_NVIDIA) && defined(CONFIG_LITMUS_SOFTIRQD)
-        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_max_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);
+#if defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON) || defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED)
+        if(is_nvidia_func(work->func)) {
+                if (__klmirqd_nv_work_schedule(work)) {
+                        return 1;
                 }
-                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
-#endif
+
         return(__schedule_work(work));
 }
 EXPORT_SYMBOL(schedule_work);
diff --git a/litmus/Kconfig b/litmus/Kconfig
index bf802b760e1f..fad4220ae49f 100644
--- a/litmus/Kconfig
+++ b/litmus/Kconfig
@@ -400,7 +400,7 @@ config LITMUS_SOFTIRQD_NONE
           Don't schedule tasklets in Litmus.  Default.
 
 config LITMUS_SOFTIRQD
-        bool "Spawn klmirqd interrupt handling threads."
+        bool "Enable klmirqd interrupt (and workqueue) handling threads."
         help
           Create klmirqd interrupt handling threads.  Work must be
           specifically dispatched to these workers.  (Softirqs for
@@ -423,11 +423,10 @@ endchoice
            
 
 config LITMUS_NVIDIA
-          bool "Litmus handling of NVIDIA interrupts."
+          bool "Litmus handling of NVIDIA driver."
           default n
           help
-            Direct tasklets from NVIDIA devices to Litmus's klmirqd
-                or PAI interrupt handling routines.
+            Enable Litmus control of NVIDIA driver tasklet/workqueues.
 
                 If unsure, say No.
 
@@ -444,6 +443,36 @@ config LITMUS_NVIDIA_NONSPLIT_INTERRUPTS
 
                 If unsure, say No.
 
+choice
+          prompt "Litmus handling of NVIDIA workqueues."
+          depends on LITMUS_NVIDIA
+          default LITMUS_NVIDIA_WORKQ_OFF
+          help
+                Select method for handling NVIDIA workqueues.
+
+config LITMUS_NVIDIA_WORKQ_OFF
+          bool "Use Linux's default work queues."
+          help
+                Let Linux process all NVIDIA work queue items.
+
+config LITMUS_NVIDIA_WORKQ_ON
+          bool "Schedule work with interrupt thread."
+          depends on LITMUS_SOFTIRQD
+          help
+                Direct work queue items from NVIDIA devices Litmus's
+                klmirqd handling routines. Use the same thread
+                as interrupt handling.
+
+config LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+          bool "Sechedule work in dedicated threads."
+          depends on LITMUS_SOFTIRQD
+          help
+                Direct work queue items from NVIDIA devices to Litmus's
+                klmirqd handling routines. Use dedicated thread for
+                work (seperate thread from interrupt handling).
+
+endchoice
+
 config LITMUS_AFFINITY_AWARE_GPU_ASSINGMENT
           bool "Enable affinity-aware heuristics to improve GPU assignment."
           depends on LITMUS_NVIDIA && LITMUS_AFFINITY_LOCKING
diff --git a/litmus/litmus.c b/litmus/litmus.c
index f0d01c34110c..740b5e57ecc1 100644
--- a/litmus/litmus.c
+++ b/litmus/litmus.c
@@ -736,7 +736,7 @@ static int litmus_shutdown_nb(struct notifier_block *unused1,
         if (litmus != &linux_sched_plugin) {
                 int ret = switch_sched_plugin(&linux_sched_plugin);
                 if (ret) {
-                        printk("Auto-shutdown of active Litmus plugin failed.\n");
+                        printk(KERN_EMERG "Auto-shutdown of active Litmus plugin failed.\n");
                 }
         }
         return NOTIFY_DONE;
diff --git a/litmus/litmus_softirq.c b/litmus/litmus_softirq.c
index a5d61afe7952..aa83b363be7c 100644
--- a/litmus/litmus_softirq.c
+++ b/litmus/litmus_softirq.c
@@ -883,9 +883,6 @@ void flush_pending(struct task_struct* tsk)
 
                         work_flushed |= LIT_TASKLET_HI;
 
-//                      t->owner = NULL;
-
-                        // WTF?
                         if(!test_and_set_bit(TASKLET_STATE_SCHED, &t->state))
                         {
                                 atomic_dec(&which->num_hi_pending);
@@ -923,9 +920,6 @@ void flush_pending(struct task_struct* tsk)
 
                         work_flushed |= LIT_TASKLET_LOW;
 
-//                      t->owner = NULL;
-//                      sched_trace_tasklet_end(owner, 1ul);
-
                         if(!test_and_set_bit(TASKLET_STATE_SCHED, &t->state))
                         {
                                 atomic_dec(&which->num_low_pending);
@@ -956,14 +950,10 @@ void flush_pending(struct task_struct* tsk)
                         atomic_dec(&which->num_work_pending);
 
                         work->owner = NULL;
-//                      sched_trace_work_end(owner, current, 1ul);
                         __schedule_work(work);
                 }
         }
 
-        //__dump_state(which, "flush_pending: after (before reeval prio)");
-
-
         mb(); /* commit changes to pending flags */
 
         raw_spin_unlock_irqrestore(&which->lock, flags);
diff --git a/litmus/nvidia_info.c b/litmus/nvidia_info.c
index 29031f741fcf..7f1dc18624e1 100644
--- a/litmus/nvidia_info.c
+++ b/litmus/nvidia_info.c
@@ -141,7 +141,7 @@ typedef struct litmus_nv_linux_state_s {
 } litmus_nv_linux_state_t;
 
 
-
+#ifdef CONFIG_SCHED_DEBUG_TRACE
 static void __attribute__((unused))
 dump_nvidia_info(const struct tasklet_struct *t)
 {
@@ -192,7 +192,6 @@ dump_nvidia_info(const struct tasklet_struct *t)
                 int ns_offset_raw = (void*)(&(linuxstate->device_num)) - (void*)(&(linuxstate->nv_state));
                 int ns_offset_desired = (void*)(&(linuxstate->device_num)) - (void*)(nvstate);
 
-
                 TRACE("LINUX NV State:\n"
                           "\tlinux nv state ptr: %p\n"
                           "\taddress of tasklet: %p\n"
@@ -226,14 +225,39 @@ dump_nvidia_info(const struct tasklet_struct *t)
                 TRACE("INVALID LINUXNVSTATE?????\n");
         }
 }
-
+#endif
 
 
 static struct module* nvidia_mod = NULL;
 
 static int init_nv_device_reg(void);
 static int shutdown_nv_device_reg(void);
+void shutdown_nvidia_info(void);
 
+static int nvidia_going_module_notify(struct notifier_block *self,
+                                unsigned long val, void *data)
+{
+        struct module *mod = data;
+
+        if (nvidia_mod && (mod == nvidia_mod)) {
+                switch (val) {
+                case MODULE_STATE_GOING:
+                        /* just set our mod reference to null to avoid crash */
+                        nvidia_mod = NULL;
+                        mb();
+                        break;
+                default:
+                        break;
+                }
+        }
+
+        return 0;
+}
+
+static struct notifier_block nvidia_going = {
+        .notifier_call = nvidia_going_module_notify,
+        .priority = 1,
+};
 
 int init_nvidia_info(void)
 {
@@ -246,6 +270,9 @@ int init_nvidia_info(void)
                           (void*)(nvidia_mod->module_core),
                           (void*)(nvidia_mod->module_core) + nvidia_mod->core_size);
                 init_nv_device_reg();
+
+                register_module_notifier(&nvidia_going);
+
                 return(0);
         }
         else
@@ -262,6 +289,7 @@ void shutdown_nvidia_info(void)
         nvidia_mod = NULL;
         mb();
 
+        unregister_module_notifier(&nvidia_going);
         shutdown_nv_device_reg();
 }
 
@@ -314,9 +342,15 @@ typedef struct {
         struct binheap  owners;
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
-        klmirqd_callback_t callback;
-        struct task_struct* thread;
-        int ready:1;  /* todo: make threads check for the ready flag */
+        klmirqd_callback_t interrupt_callback;
+        struct task_struct* interrupt_thread;
+        int interrupt_ready:1;  /* todo: make threads check for the ready flag */
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+        klmirqd_callback_t workq_callback;
+        struct task_struct* workq_thread;
+        int workq_ready:1;
+#endif
 #endif
 
 #ifdef CONFIG_LITMUS_NV_KLMIRQD_DEBUG
@@ -330,22 +364,40 @@ static nv_device_registry_t NV_DEVICE_REG[NV_DEVICE_NUM];
 
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
-static int nvidia_klmirqd_cb(void *arg)
+static int nvidia_launch_interrupt_cb(void *arg)
+{
+        unsigned long flags;
+        int reg_device_id = (int)(long long)(arg);
+        nv_device_registry_t *reg = &NV_DEVICE_REG[reg_device_id];
+
+        TRACE("nvklmirqd callback for GPU %d\n", reg_device_id);
+
+        raw_spin_lock_irqsave(&reg->lock, flags);
+        reg->interrupt_thread = current;
+        reg->interrupt_ready = 1;
+        raw_spin_unlock_irqrestore(&reg->lock, flags);
+
+        return 0;
+}
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+static int nvidia_launch_workq_cb(void *arg)
 {
         unsigned long flags;
         int reg_device_id = (int)(long long)(arg);
         nv_device_registry_t *reg = &NV_DEVICE_REG[reg_device_id];
 
-        TRACE("nv klmirqd callback for GPU %d\n", reg_device_id);
+        TRACE("nvklmworkerd callback for GPU %d\n", reg_device_id);
 
         raw_spin_lock_irqsave(&reg->lock, flags);
-        reg->thread = current;
-        reg->ready = 1;
+        reg->workq_thread = current;
+        reg->workq_ready = 1;
         raw_spin_unlock_irqrestore(&reg->lock, flags);
 
         return 0;
 }
 #endif
+#endif
 
 #ifdef CONFIG_LITMUS_NV_KLMIRQD_DEBUG
 struct nv_klmirqd_dbg_timer_struct
@@ -391,9 +443,9 @@ static enum hrtimer_restart nvklmirqd_timer_func(struct hrtimer *timer)
 
         reg = &NV_DEVICE_REG[gpu];
 
-        if (reg->thread && reg->ready) {
+        if (reg->interrupt_thread && reg->interrupt_ready) {
                 TRACE("Adding a tasklet for GPU %d\n", gpu);
-                litmus_tasklet_schedule(&reg->nv_klmirqd_dbg_tasklet, reg->thread);
+                litmus_tasklet_schedule(&reg->nv_klmirqd_dbg_tasklet, reg->interrupt_thread);
         }
         else {
                 TRACE("nv klmirqd is not ready!\n");
@@ -448,15 +500,25 @@ static int init_nv_device_reg(void)
                 {
                         int default_cpu = litmus->map_gpu_to_cpu(i);
 
+                        /* spawn the interrupt thread */
                         snprintf(name, MAX_KLMIRQD_NAME_LEN, "nvklmirqd%d", i);
-
-                        NV_DEVICE_REG[i].callback.func = nvidia_klmirqd_cb;
-                        NV_DEVICE_REG[i].callback.arg = (void*)(long long)(i);
+                        NV_DEVICE_REG[i].interrupt_callback.func = nvidia_launch_interrupt_cb;
+                        NV_DEVICE_REG[i].interrupt_callback.arg = (void*)(long long)(i);
                         mb();
+                        if(launch_klmirqd_thread(name, default_cpu, &NV_DEVICE_REG[i].interrupt_callback) != 0) {
+                                TRACE("Failed to create nvklmirqd thread for GPU %d\n", i);
+                        }
 
-                        if(launch_klmirqd_thread(name, default_cpu, &NV_DEVICE_REG[i].callback) != 0) {
-                                TRACE("Failed to create klmirqd thread for GPU %d\n", i);
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                        /* spawn the workqueue thread */
+                        snprintf(name, MAX_KLMIRQD_NAME_LEN, "nvklmworker%d", i);
+                        NV_DEVICE_REG[i].workq_callback.func = nvidia_launch_workq_cb;
+                        NV_DEVICE_REG[i].workq_callback.arg = (void*)(long long)(i);
+                        mb();
+                        if(launch_klmirqd_thread(name, default_cpu, &NV_DEVICE_REG[i].workq_callback) != 0) {
+                                TRACE("Failed to create nvklmworkqd thread for GPU %d\n", i);
                         }
+#endif
                 }
 #endif
         }
@@ -479,6 +541,7 @@ static int shutdown_nv_device_reg(void)
 
 #ifdef CONFIG_LITMUS_SOFTIRQD
         {
+                unsigned long flags;
                 int i;
                 nv_device_registry_t *reg;
 
@@ -488,12 +551,36 @@ static int shutdown_nv_device_reg(void)
 
                         reg = &NV_DEVICE_REG[i];
 
-                        if (reg->thread && reg->ready) {
-                                kill_klmirqd_thread(reg->thread);
-
-                                /* assume that all goes according to plan... */
-                                reg->thread = NULL;
-                                reg->ready = 0;
+                        if ((reg->interrupt_thread && reg->interrupt_ready)
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                                 || (reg->workq_thread && reg->workq_ready)
+#endif
+                           )
+                        {
+                                raw_spin_lock_irqsave(&reg->lock, flags);
+
+                                if (reg->interrupt_thread && reg->interrupt_ready) {
+                                        struct task_struct* th = reg->interrupt_thread;
+                                        reg->interrupt_thread = NULL;
+                                        mb();
+                                        reg->interrupt_ready = 0;
+                                        mb();
+
+                                        kill_klmirqd_thread(th);
+                                }
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                                if (reg->workq_thread && reg->workq_ready) {
+                                        struct task_struct* th = reg->workq_thread;
+                                        reg->workq_thread = NULL;
+                                        mb();
+                                        reg->workq_ready = 0;
+                                        mb();
+
+                                        kill_klmirqd_thread(th);
+                                }
+#endif
+                                raw_spin_unlock_irqrestore(&reg->lock, flags);
                         }
 
                         while (!binheap_empty(&reg->owners)) {
@@ -528,29 +615,118 @@ struct task_struct* get_nv_max_device_owner(u32 target_device_id)
         return(owner);
 }
 
+
 #ifdef CONFIG_LITMUS_SOFTIRQD
-struct task_struct* get_nv_klmirqd_thread(u32 target_device_id)
+
+typedef enum {
+        INTERRUPT_TH,
+        WORKQ_TH
+} nvklmtype_t;
+
+static struct task_struct* __get_klm_thread(nv_device_registry_t* reg, nvklmtype_t type)
 {
         struct task_struct *klmirqd = NULL;
-        nv_device_registry_t *reg;
+        
+        switch(type)
+        {
+        case INTERRUPT_TH:
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON
+        case WORKQ_TH:
+#endif
+                if(likely(reg->interrupt_ready))
+                        klmirqd = reg->interrupt_thread;
+                break;
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+        case WORKQ_TH:
+                if(likely(reg->workq_ready))
+                        klmirqd = reg->workq_thread;
+                break;
+#endif
+        }
 
-        BUG_ON(target_device_id >= NV_DEVICE_NUM);
+        return klmirqd;
+}
 
-        reg = &NV_DEVICE_REG[target_device_id];
+static struct task_struct* __get_and_lock_klm_thread(nv_device_registry_t* reg, unsigned long* flags, nvklmtype_t type)
+{
+        struct task_struct *klmirqd;
+
+        raw_spin_lock_irqsave(&reg->lock, *flags);
+        klmirqd = __get_klm_thread(reg, type);
 
-        if(likely(reg->ready)) {
-                klmirqd = reg->thread;
+        if (!klmirqd) {
+                /* unlock if thread does not exist or is not ready */
+                raw_spin_unlock_irqrestore(&reg->lock, *flags);
         }
 
         return klmirqd;
 }
-#endif
 
+static void __unlock_klm_thread(nv_device_registry_t* reg, unsigned long* flags, nvklmtype_t type)
+{
+        /* workq and interrupts share a lock per GPU */
+        raw_spin_unlock_irqrestore(&reg->lock, *flags);
+}
+
+struct task_struct* get_and_lock_nvklmirqd_thread(u32 target_device_id, unsigned long* flags)
+{
+        nv_device_registry_t *reg;
+        BUG_ON(target_device_id >= NV_DEVICE_NUM);
+        reg = &NV_DEVICE_REG[target_device_id];
+        return __get_and_lock_klm_thread(reg, flags, INTERRUPT_TH);
+}
+
+void unlock_nvklmirqd_thread(u32 target_device_id, unsigned long* flags)
+{
+        nv_device_registry_t *reg;
+        BUG_ON(target_device_id >= NV_DEVICE_NUM);
+        reg = &NV_DEVICE_REG[target_device_id];
+        __unlock_klm_thread(reg, flags, INTERRUPT_TH);
+}
 
+struct task_struct* get_nvklmirqd_thread(u32 target_device_id)
+{
+        /* should this function be allowed?  who will use klmirqd thread without thread safety? */
+        unsigned long flags;
+        struct task_struct *klmirqd;
+        klmirqd = get_and_lock_nvklmirqd_thread(target_device_id, &flags);
+        if(klmirqd)
+                unlock_nvklmirqd_thread(target_device_id, &flags);
+        return klmirqd;
+}
 
+#if defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON) || defined(CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED)
+
+struct task_struct* get_and_lock_nvklmworkqd_thread(u32 target_device_id, unsigned long* flags)
+{
+        nv_device_registry_t *reg;
+        BUG_ON(target_device_id >= NV_DEVICE_NUM);
+        reg = &NV_DEVICE_REG[target_device_id];
+        return __get_and_lock_klm_thread(reg, flags, WORKQ_TH);
+}
+
+void unlock_nvklmworkqd_thread(u32 target_device_id, unsigned long* flags)
+{
+        nv_device_registry_t *reg;
+        BUG_ON(target_device_id >= NV_DEVICE_NUM);
+        reg = &NV_DEVICE_REG[target_device_id];
+        __unlock_klm_thread(reg, flags, WORKQ_TH);
+}
+
+
+struct task_struct* get_nvklmworkqd_thread(u32 target_device_id)
+{
+        /* should this function be allowed?  who will use klmirqd thread without thread safety? */
+        unsigned long flags;
+        struct task_struct *klmirqd;
+        klmirqd = get_and_lock_nvklmworkqd_thread(target_device_id, &flags);
+        if(klmirqd)
+                unlock_nvklmworkqd_thread(target_device_id, &flags);
+        return klmirqd;
+}
+#endif // end WORKQs
 
 
-#ifdef CONFIG_LITMUS_SOFTIRQD
 static int gpu_klmirqd_increase_priority(struct task_struct *klmirqd, struct task_struct *hp)
 {
         int retval = 0;
@@ -582,7 +758,7 @@ static int gpu_klmirqd_decrease_priority(struct task_struct *klmirqd, struct tas
 
         return retval;
 }
-#endif
+#endif  // end CONFIG_LITMUS_SOFTIRQD
 
 
 
@@ -624,13 +800,27 @@ long enable_gpu_owner(struct task_struct *t)
                                                   struct task_struct, rt_param);
 
         if (hp == t) {
+                int interrupt_success;
+
                 /* we're the new hp */
                 TRACE_CUR("%s/%d (eff_prio = %s/%d) is new hp on GPU %d.\n",
                                 t->comm, t->pid,
                                 effective_priority(t)->comm, effective_priority(t)->pid,
                                 gpu);
 
-                retval = gpu_klmirqd_increase_priority(reg->thread, effective_priority(t));
+                interrupt_success = gpu_klmirqd_increase_priority(reg->interrupt_thread, effective_priority(t));
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                {
+                int workq_success = gpu_klmirqd_increase_priority(reg->workq_thread, effective_priority(t));
+                if(interrupt_success != 1 || workq_success != 1)
+                        retval = (interrupt_success != 1) ? interrupt_success : workq_success;
+                else
+                        retval = 1;
+                }
+#else
+                retval = interrupt_success;
+#endif
         }
 #endif
 
@@ -682,6 +872,9 @@ long disable_gpu_owner(struct task_struct *t)
         }
 
         if (hp == t && new_hp != t) {
+                int interrupt_success;
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+#endif
                 struct task_struct *to_inh = (new_hp) ? effective_priority(new_hp) : NULL;
 
                 TRACE_CUR("%s/%d is no longer hp on GPU %d; new hp = %s/%d (eff_prio = %s/%d).\n",
@@ -692,7 +885,19 @@ long disable_gpu_owner(struct task_struct *t)
                                                 (to_inh) ? to_inh->comm : "null",
                                                 (to_inh) ? to_inh->pid : 0);
 
-                retval = gpu_klmirqd_decrease_priority(reg->thread, to_inh);
+                interrupt_success = gpu_klmirqd_decrease_priority(reg->interrupt_thread, to_inh);
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                {
+                int workq_success = gpu_klmirqd_decrease_priority(reg->workq_thread, to_inh);
+                if(interrupt_success != 1 || workq_success != 1)
+                        retval = (interrupt_success != 1) ? interrupt_success : workq_success;
+                else
+                        retval = 1;
+                }
+#else
+                retval = interrupt_success;
+#endif
         }
 #else
         binheap_delete(&tsk_rt(t)->gpu_owner_node, &reg->owners);
@@ -758,13 +963,28 @@ int gpu_owner_increase_priority(struct task_struct *t)
 
         /* check if the eff. prio. of hp has changed */
         if (increase_klmirqd || (effective_priority(hp) != hp_eff)) {
+                int interrupt_success;
+
                 hp_eff = effective_priority(hp);
                 TRACE_CUR("%s/%d (eff_prio = %s/%d) is new hp on GPU %d.\n",
                                                 t->comm, t->pid,
                                                 hp_eff->comm, hp_eff->pid,
                                                 gpu);
 
-                retval = gpu_klmirqd_increase_priority(reg->thread, hp_eff);
+                interrupt_success = gpu_klmirqd_increase_priority(reg->interrupt_thread, hp_eff);
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                {
+                int workq_success = gpu_klmirqd_increase_priority(reg->workq_thread, hp_eff);
+                if(interrupt_success != 1 || workq_success != 1)
+                        retval = (interrupt_success != 1) ? interrupt_success : workq_success;
+                else
+                        retval = 1;
+                }
+#else
+                retval = interrupt_success;
+#endif
+
         }
 #endif
 
@@ -810,9 +1030,23 @@ int gpu_owner_decrease_priority(struct task_struct *t)
                                          struct task_struct, rt_param);
                 /* if the new_hp is still t, or if the effective priority has changed */
                 if ((new_hp == t) || (effective_priority(new_hp) != hp_eff)) {
+                        int interrupt_success;
+
                         hp_eff = effective_priority(new_hp);
                         TRACE_CUR("%s/%d is no longer hp on GPU %d.\n", t->comm, t->pid, gpu);
-                        retval = gpu_klmirqd_decrease_priority(reg->thread, hp_eff);
+                        interrupt_success = gpu_klmirqd_decrease_priority(reg->interrupt_thread, hp_eff);
+
+#ifdef CONFIG_LITMUS_NVIDIA_WORKQ_ON_DEDICATED
+                        {
+                        int workq_success = gpu_klmirqd_decrease_priority(reg->workq_thread, hp_eff);
+                        if(interrupt_success != 1 || workq_success != 1)
+                                retval = (interrupt_success != 1) ? interrupt_success : workq_success;
+                        else
+                                retval = 1;
+                        }
+#else
+                        retval = interrupt_success;
+#endif
                 }
         }
 #endif
diff --git a/litmus/sched_cedf.c b/litmus/sched_cedf.c
index df931de13e16..e2737bafa9b8 100644
--- a/litmus/sched_cedf.c
+++ b/litmus/sched_cedf.c
@@ -1701,10 +1701,6 @@ static void cleanup_cedf(void)
 {
         int i;
 
-#ifdef CONFIG_LITMUS_NVIDIA
-        shutdown_nvidia_info();
-#endif
-
         if (clusters_allocated) {
                 for (i = 0; i < num_clusters; i++) {
                         kfree(cedf[i].cpus);
diff --git a/litmus/sched_task_trace.c b/litmus/sched_task_trace.c
index 2583ee2bb437..d1df0127cfa4 100644
--- a/litmus/sched_task_trace.c
+++ b/litmus/sched_task_trace.c
@@ -349,13 +349,15 @@ EXPORT_SYMBOL(do_sched_trace_tasklet_end);
 
 
 feather_callback void do_sched_trace_work_release(unsigned long id,
-                                                                                                         unsigned long _owner)
+                                                                                                         unsigned long _owner,
+                                                                                                         unsigned long _device)
 {
         struct task_struct *t = (struct task_struct*) _owner;
         struct st_event_record *rec = get_record(ST_WORK_RELEASE, t);
 
         if (rec) {
                 rec->data.work_release.when = now();
+                rec->data.work_release.device = _device;
                 put_record(rec);
         }
 }