1 files changed, 14 insertions, 12 deletions

diff --git a/include/litmus/gpu_affinity.h b/include/litmus/gpu_affinity.h
index d64a15cbf2a5..a2abc4c1fd7c 100644
--- a/include/litmus/gpu_affinity.h
+++ b/include/litmus/gpu_affinity.h
@@ -31,18 +31,20 @@ static inline lt_t get_gpu_time(struct task_struct* t)
 
 static inline lt_t get_gpu_estimate(struct task_struct* t, gpu_migration_dist_t dist)
 {
-//      int i;
-//      fpbuf_t temp = _fp_to_integer(t->rt_param.gpu_migration_est[dist].est);
-//      lt_t val = (temp >= 0) ? temp : 0;  // never allow negative estimates...
-        lt_t val = t->rt_param.gpu_migration_est[dist].avg;
-
-//      WARN_ON(temp < 0);
-
-        // lower-bound a distant migration to be at least equal to the level
-        // below it.
-//      for(i = dist-1; (val == 0) && (i >= MIG_LOCAL); --i) {
-//              val = _fp_to_integer(t->rt_param.gpu_migration_est[i].est);
-//      }
+        lt_t val = 0;
+
+        switch(t->rt_param.prediction_mode) {
+                case SIMPLE_AVG:
+                case CC_BRUTE_AVG:
+                case CC_MR_AVG:
+                        val = t->rt_param.gpu_avg_est[dist].center_line;
+                        break;
+                case SIMPLE_EWMA:
+                case CC_BRUTE_EWMA:
+                case CC_MR_EWMA:
+                        val = t->rt_param.gpu_ewma_est[dist].center_line;
+                        break;
+        }
 
         return ((val > 0) ? val : dist+1);
 }