2 files changed, 61 insertions, 6 deletions

diff --git a/include/litmus/ikglp_lock.h b/include/litmus/ikglp_lock.h
index 2558a382446c..dd4f14e4eaff 100644
--- a/include/litmus/ikglp_lock.h
+++ b/include/litmus/ikglp_lock.h
@@ -118,7 +118,7 @@ struct ikglp_queue_info
 struct ikglp_affinity_ops
 {
         struct fifo_queue* (*advise_enqueue)(struct ikglp_affinity* aff, struct task_struct* t);        // select FIFO
-        ikglp_wait_state_t* (*advise_steal)(struct ikglp_affinity* aff);        // select steal from FIFO
+        ikglp_wait_state_t* (*advise_steal)(struct ikglp_affinity* aff, struct fifo_queue* dst);        // select steal from FIFO
         ikglp_donee_heap_node_t* (*advise_donee_selection)(struct ikglp_affinity* aff, struct task_struct* t);  // select a donee
         ikglp_wait_state_t* (*advise_doner_to_fq)(struct ikglp_affinity* aff, struct fifo_queue* dst);  // select a donor to move to PQ
 
diff --git a/litmus/ikglp_lock.c b/litmus/ikglp_lock.c
index 9ae71422d813..1ce6572eddcb 100644
--- a/litmus/ikglp_lock.c
+++ b/litmus/ikglp_lock.c
@@ -1281,7 +1281,7 @@ int ikglp_unlock(struct litmus_lock* l)
 
 #ifdef CONFIG_LITMUS_AFFINITY_LOCKING
                 fq_wait = (sem->aff_obs) ?
-                        sem->aff_obs->ops->advise_steal(sem->aff_obs) :
+                        sem->aff_obs->ops->advise_steal(sem->aff_obs, fq) :
                         ikglp_find_hp_waiter_to_steal(sem);
 #else
                 fq_wait = ikglp_find_hp_waiter_to_steal(sem);
@@ -1932,7 +1932,8 @@ struct fifo_queue* gpu_ikglp_advise_enqueue(struct ikglp_affinity* aff, struct t
         //return(sem->shortest_fifo_queue);
 }
 
-ikglp_wait_state_t* gpu_ikglp_advise_steal(struct ikglp_affinity* aff)
+ikglp_wait_state_t* gpu_ikglp_advise_steal(struct ikglp_affinity* aff,
+                                                                                   struct fifo_queue* dst)
 {
         struct ikglp_semaphore *sem = ikglp_from_lock(aff->obs.lock);
 
@@ -2127,11 +2128,64 @@ out:
         return(donee_node);
 }
 
+
+
+static void __find_closest_donor(int target_gpu,
+                                                                 struct binheap_node* donor_node,
+                                                                 ikglp_wait_state_t** cur_closest,
+                                                                 int* cur_dist)
+{
+        ikglp_wait_state_t *this_donor =
+                binheap_entry(donor_node, ikglp_wait_state_t, node);
+        
+        int this_dist =
+                gpu_migration_distance(target_gpu, tsk_rt(this_donor->task)->last_gpu);
+        
+        if(this_dist < *cur_dist) {
+                // take this donor
+                *cur_dist = this_dist;
+                *cur_closest = this_donor;
+        }
+        else if(this_dist == *cur_dist) {
+                // priority tie-break.  Even though this is a pre-order traversal,
+                // this is a heap, not a binary tree, so we still need to do a priority
+                // comparision.
+                if(!(*cur_closest) ||
+                   litmus->compare(this_donor->task, (*cur_closest)->task)) {
+                        *cur_dist = this_dist;
+                        *cur_closest = this_donor;
+                }
+        }
+        
+    if(donor_node->left) __find_closest_donor(target_gpu, donor_node->left, cur_closest, cur_dist);
+    if(donor_node->right) __find_closest_donor(target_gpu, donor_node->right, cur_closest, cur_dist);
+}
+
 ikglp_wait_state_t* gpu_ikglp_advise_doner_to_fq(struct ikglp_affinity* aff, struct fifo_queue* fq)
 {
-        // TODO: MAKE THIS SMARTER
+        // Huristic strategy: Find donor with the closest affinity to fq.
+        // Tie-break on priority.
+        
+        // We need to iterate over all the donors to do this.  Unfortunatly,
+        // our donors are organized in a heap.  We'll visit each node with a
+        // recurisve call.  This is realitively safe since there are only sem->m
+        // donors, at most.  We won't recurse too deeply to have to worry about
+        // our stack.  (even with 128 CPUs, our nest depth is at most 7 deep).
+        
         struct ikglp_semaphore *sem = ikglp_from_lock(aff->obs.lock);
-        ikglp_wait_state_t* donor = binheap_top_entry(&sem->donors, ikglp_wait_state_t, node);
+        ikglp_wait_state_t *donor = NULL;
+        int distance = MIG_NONE;
+        int gpu = replica_to_gpu(aff, ikglp_get_idx(sem, fq));
+        ikglp_wait_state_t* default_donor = binheap_top_entry(&sem->donors, ikglp_wait_state_t, node);
+        
+        __find_closest_donor(gpu, sem->donors.root, &donor, &distance);
+        
+        TRACE_CUR("Selected %s/%d as donor (distance = %d) "
+                          "(non-aff wanted %s/%d)\n",
+                          donor->task->comm, donor->task->pid,
+                          distance,
+                          default_donor->task->comm, default_donor->task->pid);
+        
         return(donor);
 }
 
@@ -2331,7 +2385,8 @@ struct fifo_queue* simple_gpu_ikglp_advise_enqueue(struct ikglp_affinity* aff, s
         return to_enqueue;
 }
 
-ikglp_wait_state_t* simple_gpu_ikglp_advise_steal(struct ikglp_affinity* aff)
+ikglp_wait_state_t* simple_gpu_ikglp_advise_steal(struct ikglp_affinity* aff,
+                                                                                                  struct fifo_queue* dst)
 {
         struct ikglp_semaphore *sem = ikglp_from_lock(aff->obs.lock);
         //      TRACE_CUR("Simple GPU ikglp advise_steal invoked\n");