diff --git a/init/main.c b/init/main.c
index badae3bf08f1605bb87c3fdefcf4ea4893350ba3..df58a416dd1da54df4c3f2be807adb5366ec35b2 100644
--- a/init/main.c
+++ b/init/main.c
@@ -414,6 +414,16 @@ static noinline void __ref rest_init(void)
 	rcu_read_lock();
 	kthreadd_task = find_task_by_pid_ns(pid, &init_pid_ns);
 	rcu_read_unlock();
+
+	/*
+	 * Enable might_sleep() and smp_processor_id() checks.
+	 * They cannot be enabled earlier because with CONFIG_PRREMPT=y
+	 * kernel_thread() would trigger might_sleep() splats. With
+	 * CONFIG_PREEMPT_VOLUNTARY=y the init task might have scheduled
+	 * already, but it's stuck on the kthreadd_done completion.
+	 */
+	system_state = SYSTEM_SCHEDULING;
+
 	complete(&kthreadd_done);
 
 	/*
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 877241e9f2b0daad807b55f056e69045010bd32d..c3e50cada84d8e38070857e619900f98b1b1da79 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6238,8 +6238,10 @@ void ___might_sleep(const char *file, int line, int preempt_offset)
 
 	if ((preempt_count_equals(preempt_offset) && !irqs_disabled() &&
 	     !is_idle_task(current)) ||
-	    system_state != SYSTEM_RUNNING || oops_in_progress)
+	    system_state == SYSTEM_BOOTING || system_state > SYSTEM_RUNNING ||
+	    oops_in_progress)
 		return;
+
 	if (time_before(jiffies, prev_jiffy + HZ) && prev_jiffy)
 		return;
 	prev_jiffy = jiffies;
diff --git a/lib/smp_processor_id.c b/lib/smp_processor_id.c
index 690d75b132fa7cfdb41e2f6ae8fd6554db70e9b4..2fb007be02125a17f4c3b3345e656773d49209ae 100644
--- a/lib/smp_processor_id.c
+++ b/lib/smp_processor_id.c
@@ -28,7 +28,7 @@ notrace static unsigned int check_preemption_disabled(const char *what1,
 	/*
 	 * It is valid to assume CPU-locality during early bootup:
 	 */
-	if (system_state != SYSTEM_RUNNING)
+	if (system_state < SYSTEM_SCHEDULING)
 		goto out;
 
 	/*