diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 45ea62f1b2b5..5bd1105d1a11 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -730,7 +730,6 @@ int stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 struct nmi_stop {
 	struct multi_stop_data	*data;
 	int			ret;
-	bool			done;
 };
 
 static DEFINE_PER_CPU(struct nmi_stop, nmi_stop);
@@ -743,10 +742,10 @@ static DEFINE_PER_CPU(struct nmi_stop, nmi_stop);
 bool noinstr stop_machine_nmi_handler(void)
 {
 	struct multi_stop_data *msdata = raw_cpu_read(nmi_stop.data);
-	unsigned int cpu = smp_processor_id();
+	unsigned int cpu = raw_smp_processor_id();
 	int ret;
 
-	if (!msdata || !cpumask_test_and_clear_cpu(cpu, msdata->nmi_cpus))
+	if (!msdata || !arch_test_and_clear_bit(cpu, cpumask_bits(msdata->nmi_cpus)))
 		return false;
 
 	/*
@@ -759,7 +758,6 @@ bool noinstr stop_machine_nmi_handler(void)
 	instrumentation_end();
 
 	raw_cpu_write(nmi_stop.ret,  ret);
-	raw_cpu_write(nmi_stop.done, true);
 	raw_cpu_write(nmi_stop.data, NULL);
 
 	return true;
@@ -770,10 +768,11 @@ static bool wait_for_nmi_handler(void)
 	/* Conservative timeout */
 	unsigned long timeout = USEC_PER_SEC;
 
-	while (!this_cpu_read(nmi_stop.done) && timeout--)
+	/* The handler clears up at the end */
+	while (this_cpu_read(nmi_stop.data) && timeout--)
 		udelay(1);
 
-	return this_cpu_read(nmi_stop.done);
+	return !this_cpu_read(nmi_stop.data);
 }
 
 static int nmi_stop_run(struct multi_stop_data *msdata)
@@ -783,12 +782,16 @@ static int nmi_stop_run(struct multi_stop_data *msdata)
 	 * self-NMI to execute the stop function from the NMI handler
 	 */
 	this_cpu_write(nmi_stop.data, msdata);
-	this_cpu_write(nmi_stop.done, false);
 	arch_send_self_nmi();
 
-	/* Ensure the handler went through before reading the result */
-	if (!wait_for_nmi_handler())
+	/*
+	 * Ensure the handler went through before reading the result.
+	 * Otherwise, make no stale state left behind.
+	 */
+	if (!wait_for_nmi_handler()) {
+		this_cpu_write(nmi_stop.data, NULL);
 		return -ETIMEDOUT;
+	}
 
 	return this_cpu_read(nmi_stop.ret);
 }