diff --git a/include/linux/sched/idle.h b/include/linux/sched/idle.h
index d739ab810e00..c22312087c30 100644
--- a/include/linux/sched/idle.h
+++ b/include/linux/sched/idle.h
@@ -58,8 +58,8 @@ static __always_inline bool __must_check current_set_polling_and_test(void)
 	__current_set_polling();
 
 	/*
-	 * Polling state must be visible before we test NEED_RESCHED,
-	 * paired by resched_curr()
+	 * Polling state must be visible before we test NEED_RESCHED or
+	 * NOTIFY_IPI paired by resched_curr() or notify_ipi_if_polling()
 	 */
 	smp_mb__after_atomic();
 
@@ -71,8 +71,8 @@ static __always_inline bool __must_check current_clr_polling_and_test(void)
 	__current_clr_polling();
 
 	/*
-	 * Polling state must be visible before we test NEED_RESCHED,
-	 * paired by resched_curr()
+	 * Polling state must be visible before we test NEED_RESCHED or
+	 * NOTIFY_IPI paired by resched_curr() or notify_ipi_if_polling()
 	 */
 	smp_mb__after_atomic();
 
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index db4be4921e7f..6fb6e5b75724 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -909,12 +909,30 @@ static inline bool set_nr_and_not_polling(struct task_struct *p)
 }
 
 /*
- * Atomically set TIF_NEED_RESCHED if TIF_POLLING_NRFLAG is set.
+ * Certain architectures that support TIF_POLLING_NRFLAG may not support
+ * TIF_NOTIFY_IPI to notify an idle CPU in TIF_POLLING mode of a pending
+ * IPI. On such architectures, set TIF_NEED_RESCHED instead to wake the
+ * idle CPU and process the pending IPI.
+ */
+#ifdef _TIF_NOTIFY_IPI
+#define _TIF_WAKE_FLAG _TIF_NOTIFY_IPI
+#else
+#define _TIF_WAKE_FLAG _TIF_NEED_RESCHED
+#endif
+
+/*
+ * Atomically set TIF_WAKE_FLAG when TIF_POLLING_NRFLAG is set.
+ *
+ * On architectures that define TIF_NOTIFY_IPI, the same is set in the
+ * idle task's thread_info to pull the CPU out of idle and process
+ * the pending interrupt. On architectures that don't support
+ * TIF_NOTIFY_IPI, TIF_NEED_RESCHED is set instead to notify the
+ * pending IPI.
  *
- * If this returns true, then the idle task promises to call
- * sched_ttwu_pending() and reschedule soon.
+ * If this returns true, then the idle task promises to process the
+ * call function soon.
  */
-static bool set_nr_if_polling(struct task_struct *p)
+static bool notify_ipi_if_polling(struct task_struct *p)
 {
 	struct thread_info *ti = task_thread_info(p);
 	typeof(ti->flags) val = READ_ONCE(ti->flags);
@@ -922,9 +940,16 @@ static bool set_nr_if_polling(struct task_struct *p)
 	do {
 		if (!(val & _TIF_POLLING_NRFLAG))
 			return false;
-		if (val & _TIF_NEED_RESCHED)
+		/*
+		 * If TIF_NEED_RESCHED flag is set in addition to
+		 * TIF_POLLING_NRFLAG, the CPU will soon fall out of
+		 * idle. Since flush_smp_call_function_queue() is called
+		 * soon after the idle exit, setting TIF_WAKE_FLAG is
+		 * not necessary.
+		 */
+		if (val & (_TIF_NEED_RESCHED | _TIF_WAKE_FLAG))
 			return true;
-	} while (!try_cmpxchg(&ti->flags, &val, val | _TIF_NEED_RESCHED));
+	} while (!try_cmpxchg(&ti->flags, &val, val | _TIF_WAKE_FLAG));
 
 	return true;
 }
@@ -937,7 +962,7 @@ static inline bool set_nr_and_not_polling(struct task_struct *p)
 }
 
 #ifdef CONFIG_SMP
-static inline bool set_nr_if_polling(struct task_struct *p)
+static inline bool notify_ipi_if_polling(struct task_struct *p)
 {
 	return false;
 }
@@ -3918,7 +3943,7 @@ void sched_ttwu_pending(void *arg)
  */
 bool call_function_single_prep_ipi(int cpu)
 {
-	if (set_nr_if_polling(cpu_rq(cpu)->idle)) {
+	if (notify_ipi_if_polling(cpu_rq(cpu)->idle)) {
 		trace_sched_wake_idle_without_ipi(cpu);
 		return false;
 	}
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index fcc734f45a2a..b91dc1f62a56 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -315,13 +315,13 @@ static void do_idle(void)
 	}
 
 	/*
-	 * Since we fell out of the loop above, we know TIF_NEED_RESCHED must
-	 * be set, propagate it into PREEMPT_NEED_RESCHED.
+	 * Since we fell out of the loop above, TIF_NEED_RESCHED may be set.
+	 * Propagate it into PREEMPT_NEED_RESCHED.
 	 *
 	 * This is required because for polling idle loops we will not have had
 	 * an IPI to fold the state for us.
 	 */
-	preempt_set_need_resched();
+	preempt_fold_need_resched();
 	tick_nohz_idle_exit();
 	__current_clr_polling();
 
@@ -338,7 +338,15 @@ static void do_idle(void)
 	 */
 	current_clr_notify_ipi();
 	flush_smp_call_function_queue();
-	schedule_idle();
+
+	/*
+	 * When NEED_RESCHED is set, the idle thread promises to call
+	 * schedule_idle(). schedule_idle() can be skipped when an idle CPU
+	 * was woken up to process an IPI that does not queue a task on the
+	 * idle CPU, facilitating faster idle re-entry.
+	 */
+	if (need_resched())
+		schedule_idle();
 
 	if (unlikely(klp_patch_pending(current)))
 		klp_update_patch_state(current);
