int curr_ret_stack;
        /* Stack of return addresses for return function tracing */
        struct ftrace_ret_stack *ret_stack;
+       /* time stamp for last schedule */
+       unsigned long long ftrace_timestamp;
        /*
         * Number of functions that haven't been traced
         * because of depth overrun.
 
 #include <linux/list.h>
 #include <linux/hash.h>
 
+#include <trace/sched.h>
+
 #include <asm/ftrace.h>
 
 #include "trace.h"
        return ret;
 }
 
+static void
+ftrace_graph_probe_sched_switch(struct rq *__rq, struct task_struct *prev,
+                               struct task_struct *next)
+{
+       unsigned long long timestamp;
+       int index;
+
+       timestamp = trace_clock_local();
+
+       prev->ftrace_timestamp = timestamp;
+
+       /* only process tasks that we timestamped */
+       if (!next->ftrace_timestamp)
+               return;
+
+       /*
+        * Update all the counters in next to make up for the
+        * time next was sleeping.
+        */
+       timestamp -= next->ftrace_timestamp;
+
+       for (index = next->curr_ret_stack; index >= 0; index--)
+               next->ret_stack[index].calltime += timestamp;
+}
+
 /* Allocate a return stack for each task */
 static int start_graph_tracing(void)
 {
                ret = alloc_retstack_tasklist(ret_stack_list);
        } while (ret == -EAGAIN);
 
+       if (!ret) {
+               ret = register_trace_sched_switch(ftrace_graph_probe_sched_switch);
+               if (ret)
+                       pr_info("ftrace_graph: Couldn't activate tracepoint"
+                               " probe to kernel_sched_switch\n");
+       }
+
        kfree(ret_stack_list);
        return ret;
 }
        mutex_lock(&ftrace_lock);
 
        atomic_dec(&ftrace_graph_active);
+       unregister_trace_sched_switch(ftrace_graph_probe_sched_switch);
        ftrace_graph_return = (trace_func_graph_ret_t)ftrace_stub;
        ftrace_graph_entry = ftrace_graph_entry_stub;
        ftrace_shutdown(FTRACE_STOP_FUNC_RET);
                t->curr_ret_stack = -1;
                atomic_set(&t->tracing_graph_pause, 0);
                atomic_set(&t->trace_overrun, 0);
+               t->ftrace_timestamp = 0;
        } else
                t->ret_stack = NULL;
 }