1 files changed, 5 insertions, 0 deletions

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index b073ffda2bf2..b3252d165e22 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -87,6 +87,7 @@
 #include "../smpboot.h"
 
 #include <litmus/trace.h>
+#include <litmus/sched_trace.h>
 
 #define CREATE_TRACE_POINTS
 #include <trace/events/sched.h>
@@ -1968,6 +1969,8 @@ asmlinkage void schedule_tail(struct task_struct *prev)
         rq = this_rq();
         finish_task_switch(rq, prev);
 
+        sched_trace_task_switch_to(current);
+
         /*
          * FIXME: do we need to worry about rq being invalidated by the
          * task_switch?
@@ -2987,6 +2990,7 @@ need_resched:
          */
 litmus_need_resched_nonpreemptible:
         TS_SCHED_START;
+        sched_trace_task_switch_away(prev);
 
         schedule_debug(prev);
 
@@ -3052,6 +3056,7 @@ litmus_need_resched_nonpreemptible:
         }
 
         TS_SCHED2_START(prev);
+        sched_trace_task_switch_to(current);
 
         post_schedule(rq);