re-enable klmirqd for workqueues, and grace reboot

6 files changed, 307 insertions, 56 deletions

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);
         }
 }