2 files changed, 8 insertions, 1 deletions

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index fac6492f0b98..94732d1ab00a 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2781,6 +2781,10 @@ static struct rq *finish_task_switch(struct task_struct *prev)
                  * task and put them back on the free list.
                  */
                 kprobe_flush_task(prev);
+
+                /* Task is done with its stack. */
+                put_task_stack(prev);
+
                 put_task_struct(prev);
         }
 
@@ -3403,7 +3407,6 @@ static void __sched notrace __schedule(bool preempt)
 
         balance_callback(rq);
 }
-STACK_FRAME_NON_STANDARD(__schedule); /* switch_to() */
 
 void __noreturn do_task_dead(void)
 {
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 58df5590d028..055f935d4421 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1021,7 +1021,11 @@ static inline void __set_task_cpu(struct task_struct *p, unsigned int cpu)
          * per-task data have been completed by this moment.
          */
         smp_wmb();
+#ifdef CONFIG_THREAD_INFO_IN_TASK
+        p->cpu = cpu;
+#else
         task_thread_info(p)->cpu = cpu;
+#endif
         p->wake_cpu = cpu;
 #endif
 }