[PATCH v2 00/11] x86/microcode: Refactor NMI-based rendezvous mechanism to stop-machine

[PATCH v2 00/11] x86/microcode: Refactor NMI-based rendezvous mechanism to stop-machine

[Chang S. Bae]

The last version [0] took off patches for the new NMI stop-machine
facility from the dynamic mitigation series. The review comments prompted me
reconsider it and I ended up reworking on that part for this series.

Here are a few things to call out

 * Introduce a cpumask to ensure NMI delivery to target CPUs, which
   allows removal of the static key and per-CPU flag (Boris)

 * Add a Kconfig to hide the facility to other architectures. There is no
   known use case outside x86, so this is to avoid any misuse (Thomas)

 * Switch to accumulate error code for multi-CPU stop-machine instead of
   returning an arbitrary error (Boris)

 * Ensure the handler write the return value before it is read. In
   practice, it was observed that IPI/NMI was not immediate, but an
   immediate read without any guard isn't reliable.

 * Clarify the behavior when no CPU target is given (@cpus=NULL). Instead
   of executing the stop function on every online CPU, it picks only one
   CPU (CPU0, in fact) to run.

 * Refactor the existing flow to better accommodate the new facility.

Other than that, the rest of x86 side changes are relatively minor (typo,
renaming, rebasing).

---

Below is an updated version of the original cover letter:

David Kaplan previously proposed introducing an NMI-based stop-machine
mechanism as part of his dynamic mitigation series [1]. This "big hammer"
approach was also discussed during the last LPC x86 uconf [2].

This patchset uses the facility for the late-loading of Intel microcode
by refactoring existing NMI control mechanism. There is one bit to call
out for the review here -- offline CPU handling. Before highlighting the
case, let me first walk through the online case:

== Online CPU Handling ==

Currently, the microcode loader uses a NMI-based control flow, similar
to the proposed stop_machine_nmi(). When the NMI option is enabled, the
loader invokes stop-machine, and the callback establishes an IPI back to
the local CPU, which then performs the microcode update from NMI context.

The new facility incorporates this logic:

 * Raise self-NMI, then execute the stop function from the handler

With simplification:

 * Replace the static key and per-CPU flag with a cpumask that ensures
   NMI delivery.

Given this, convert this to use the (new) stop-machine API.

== Offline CPU Handling ==

With NMI-based microcode loading, soft-offlined CPUs are also brought
into the rendezvous to avoid unexpected behavior changes (e.g. related
to MWAIT) introduced by new microcode [3].

Today, the boot CPU explicitly fires IPIs to those offline CPUs, which
establishes yet another NMI control path. However, the new stop-machine
facility does not currently account for this case, as it is primarily
shaped for online CPUs.

This patch set is intended to introduce a new NMI stop-machine along with
an existing use case (Intel ucode late-loading), independent of the
dynamic mitigation series.

The commonality, nonetheless, between them lies on online CPU handling.
So, no new facility is introduced for offline CPUs. Instead, the existing
per-CPU flag in the x86 microcode loader continues to gate the handler
without the (now removed) static key.

== Patch Structure ==

 * PATCH 1–6:  Establish the stop-machine NMI facility
 * PATCH 7–10: Incremental refactoring, with offline handling update.

The patch set is available in this repository:
 git://github.com/intel-staging/microcode.git nmi-stop-machine_v2

Testing was performed using late-loading on Intel CPUs, which default to
the NMI-based update path.

Thanks,
Chang

[0] V1: https://lore.kernel.org/lkml/20260125014224.249901-1-chang.seok.bae@intel.com/
[1]: https://lore.kernel.org/lkml/20251013143444.3999-1-david.kaplan@amd.com/
[2]: https://youtu.be/2eEPYt5XrCE?si=NUGF2pkqk8MUh9GR&t=1739
[3]: 8f849ff63bcb ("x86/microcode: Handle "offline" CPUs correctly")

Chang S. Bae (9):
  stop_machine: Clarify @cpus == NULL semantics
  stop_machine: Accumulate error code rather than overwrite
  stop_machine: Refactor multi-CPU stop glue code
  stop_machine: Add NMI-based execution path
  stop_machine: Introduce stop_machine_nmi_cpuslocked()
  x86/microcode: Distinguish NMI control path on stop-machine callback
  x86/microcode: Use stop-machine NMI facility
  x86/nmi: Simplify offline microcode handler invocation
  x86/microcode: Remove microcode_nmi_handler_enable

David Kaplan (2):
  x86/apic: Implement self-NMI support
  x86/nmi: Support NMI stop-machine handler

 arch/Kconfig                         |   3 +
 arch/x86/Kconfig                     |   1 +
 arch/x86/include/asm/microcode.h     |  11 --
 arch/x86/kernel/apic/ipi.c           |  10 ++
 arch/x86/kernel/cpu/microcode/core.c |  35 +----
 arch/x86/kernel/nmi.c                |   7 +-
 include/linux/stop_machine.h         |  58 +++++++--
 kernel/stop_machine.c                | 188 +++++++++++++++++++++++----
 8 files changed, 241 insertions(+), 72 deletions(-)


base-commit: 6e45429392a5cdd8168619d7558f7bd6cf706394
-- 
2.51.0

[PATCH v2 01/11] stop_machine: Clarify @cpus == NULL semantics

[Chang S. Bae]

The stop-machine API description currently mentions that @cpus == NULL
means running on "each cpu_online_mask". Previously, it was described as
"any cpu_online_mask" before commit:

  fc6f89dc707 ("stop_machine: Improve kernel-doc function-header comments")

In fact, multi_cpu_stop() selects the first CPU in cpu_online_mask when
@cpus is NULL. Right now cpumask_any() is defined as cpumask_first(). So
the previous description was closer but apparently it was not clear
enough either.

Fix those comments for clarity.

Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: New patch

Considering stop_machine_nmi_cpuslocked() and its another cpumask, I
could realize this nullptr implication is not clear enough. Then, it
ended up with this fix.

I also considered just saying CPU0, but still 'cpumask_first(cpu_online_mask)'
are there. So, leave it like that, instead of converting them
aggressively.
---
 include/linux/stop_machine.h | 6 ++++--
 kernel/stop_machine.c        | 4 +++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index 72820503514c..c753dd53e79d 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -99,7 +99,8 @@ static inline void print_stop_info(const char *log_lvl, struct task_struct *task
  * stop_machine: freeze the machine on all CPUs and run this function
  * @fn: the function to run
  * @data: the data ptr to pass to @fn()
- * @cpus: the cpus to run @fn() on (NULL = run on each online CPU)
+ * @cpus: the CPUs to run @fn() on. If NULL, @fn() runs on a single
+ *        (arbitrary) CPU from cpu_online_mask.
  *
  * Description: This causes a thread to be scheduled on every CPU, which
  * will run with interrupts disabled.  Each CPU specified by @cpus will
@@ -133,7 +134,8 @@ int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus);
  * stop_machine_cpuslocked: freeze the machine on all CPUs and run this function
  * @fn: the function to run
  * @data: the data ptr to pass to @fn()
- * @cpus: the cpus to run @fn() on (NULL = run on each online CPU)
+ * @cpus: the CPUs to run @fn() on. If NULL, @fn() runs on a single
+ *        (arbitrary) CPU from cpu_online_mask.
  *
  * Same as above.  Avoids nested calls to cpus_read_lock().
  *
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 3fe6b0c99f3d..822cf56fdc81 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -655,9 +655,11 @@ EXPORT_SYMBOL_GPL(stop_core_cpuslocked);
 
 /**
  * stop_machine_from_inactive_cpu - stop_machine() from inactive CPU
+ *
  * @fn: the function to run
  * @data: the data ptr for the @fn()
- * @cpus: the cpus to run the @fn() on (NULL = any online cpu)
+ * @cpus: the CPUs to run the @fn() on. If NULL, @fn() runs on a single
+ *        (arbitrary) CPU from cpu_online_mask.
  *
  * This is identical to stop_machine() but can be called from a CPU which
  * is not active.  The local CPU is in the process of hotplug (so no other
-- 
2.51.0

[RFC][PATCH v2 02/11] stop_machine: Accumulate error code rather than overwrite

[Chang S. Bae]

cpu_stopper_thread() invokes a stop function and collects its error code
in struct cpu_stop_done. In the multi stop-machine case, it is shared
data, but currently an arbitrary error is recorded as overwriting.

With different errors, accumulating error code instead can distinguish a
multi-error condition as bits are cumulatively set.

Convert the error recoding to accumulate return values.

Suggested-by: Borislav Petkov <bp@alien8.de>
Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
Link: https://lore.kernel.org/lkml/20260304163335.GDaahe3wdnqxSC2yfw@fat_crate.local
---
V1 -> V2: New patch

While tried to explain its benefit here, I considered this change
deserves more discussions to ensure its impact, so RFC.
---
 include/linux/stop_machine.h | 12 ++++++------
 kernel/stop_machine.c        | 10 +++++-----
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index c753dd53e79d..2f986555113a 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -124,9 +124,9 @@ static inline void print_stop_info(const char *log_lvl, struct task_struct *task
  * the possibility of blocking in cpus_read_lock() means that the caller
  * cannot usefully rely on this serialization.
  *
- * Return: 0 if all invocations of @fn return zero.  Otherwise, the
- * value returned by an arbitrarily chosen member of the set of calls to
- * @fn that returned non-zero.
+ * Return: 0 if all invocations of @fn return zero.  Otherwise, an
+ * accumulated return value from all invocation of @fn that returned
+ * non-zero.
  */
 int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus);
 
@@ -154,9 +154,9 @@ int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data, const struct cpumask *
  *
  * Context: Must be called from within a cpus_read_lock() protected region.
  *
- * Return: 0 if all invocations of @fn return zero.  Otherwise, the
- * value returned by an arbitrarily chosen member of the set of calls to
- * @fn that returned non-zero.
+ * Return: 0 if all invocations of @fn return zero.  Otherwise, an
+ * accumulated return value from all invocation of @fn that returned
+ * non-zero.
  */
 int stop_core_cpuslocked(unsigned int cpu, cpu_stop_fn_t fn, void *data);
 
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 822cf56fdc81..15268f1207e9 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -459,7 +459,7 @@ static int __stop_cpus(const struct cpumask *cpumask,
  * RETURNS:
  * -ENOENT if @fn(@arg) was not executed at all because all cpus in
  * @cpumask were offline; otherwise, 0 if all executions of @fn
- * returned 0, any non zero return value if any returned non zero.
+ * returned 0, the accumulated value of all non-zero @fn returns.
  */
 static int stop_cpus(const struct cpumask *cpumask, cpu_stop_fn_t fn, void *arg)
 {
@@ -512,7 +512,7 @@ static void cpu_stopper_thread(unsigned int cpu)
 		ret = fn(arg);
 		if (done) {
 			if (ret)
-				done->ret = ret;
+				done->ret |= ret;
 			cpu_stop_signal_done(done);
 		}
 		preempt_count_dec();
@@ -674,8 +674,8 @@ EXPORT_SYMBOL_GPL(stop_core_cpuslocked);
  * Local CPU is inactive.  Temporarily stops all active CPUs.
  *
  * RETURNS:
- * 0 if all executions of @fn returned 0, any non zero return value if any
- * returned non zero.
+ * 0 if all executions of @fn returned 0, otherwise the accumulated value
+ * of all non-zero @fn returns.
  */
 int stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 				  const struct cpumask *cpus)
@@ -705,5 +705,5 @@ int stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 		cpu_relax();
 
 	mutex_unlock(&stop_cpus_mutex);
-	return ret ?: done.ret;
+	return ret | done.ret;
 }
-- 
2.51.0

[PATCH v2 03/11] stop_machine: Refactor multi-CPU stop glue code

[Chang S. Bae]

stop_machine_cpuslocked() currently configures struct multi_stop_data and
invokes the multi-stop operation. An upcoming stop_machine() variant will
have different configurations but the latter part will be shareable.

Extract the common part into stop_multi_cpus() to highlight each unique
front-end.

No functional change.

Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: New patch
---
 kernel/stop_machine.c | 35 ++++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 15268f1207e9..092c65c002ff 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -584,18 +584,8 @@ static int __init cpu_stop_init(void)
 }
 early_initcall(cpu_stop_init);
 
-int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data,
-			    const struct cpumask *cpus)
+static int stop_multi_cpus(struct multi_stop_data *msdata)
 {
-	struct multi_stop_data msdata = {
-		.fn = fn,
-		.data = data,
-		.num_threads = num_online_cpus(),
-		.active_cpus = cpus,
-	};
-
-	lockdep_assert_cpus_held();
-
 	if (!stop_machine_initialized) {
 		/*
 		 * Handle the case where stop_machine() is called
@@ -605,19 +595,34 @@ int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data,
 		unsigned long flags;
 		int ret;
 
-		WARN_ON_ONCE(msdata.num_threads != 1);
+		WARN_ON_ONCE(msdata->num_threads != 1);
 
 		local_irq_save(flags);
 		hard_irq_disable();
-		ret = (*fn)(data);
+		ret = msdata->fn(msdata->data);
 		local_irq_restore(flags);
 
 		return ret;
 	}
 
 	/* Set the initial state and stop all online cpus. */
-	set_state(&msdata, MULTI_STOP_PREPARE);
-	return stop_cpus(cpu_online_mask, multi_cpu_stop, &msdata);
+	set_state(msdata, MULTI_STOP_PREPARE);
+	return stop_cpus(cpu_online_mask, multi_cpu_stop, msdata);
+}
+
+int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data,
+			    const struct cpumask *cpus)
+{
+	struct multi_stop_data msdata = {
+		.fn		= fn,
+		.data		= data,
+		.num_threads	= num_online_cpus(),
+		.active_cpus	= cpus,
+	};
+
+	lockdep_assert_cpus_held();
+
+	return stop_multi_cpus(&msdata);
 }
 
 int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus)
-- 
2.51.0

[PATCH v2 04/11] stop_machine: Add NMI-based execution path

[Chang S. Bae]

Currently multi_cpu_stop() executes the stop function in the
MULTI_STOP_RUN state. But NMIs can still preempt the execution. For use
cases requiring stricter isolation from them, provide an option to run
the stop function from NMI context.

Then, the NMI stop function must be built without instrumentation, as
exceptions may re-enable NMIs on return via IRET. Annotate the NMI
handler entirely with noinstr.

However, objtool cannot reliably determine whether the indirect call
target is located in a noinstr section, so it may emit false positives.
To avoid this, temporarily lift the instrumentation restriction around
the indirect call site and document the intention.

The x86 microcode loader is currently the primary user for this. But
other architectures are not expected to use it. Add a build option to
make it opt-in.

Originally-by: David Kaplan <david.kaplan@amd.com>
Suggested-by: Borislav Petkov <bp@alien8.de>
Suggested-by: Thomas Gleixner <tglx@kernel.org>
Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
Link: https://lore.kernel.org/lkml/20260129121729.GRaXtP2aeWkQKegxC2@fat_crate.local
Link: https://lore.kernel.org/lkml/87wm0zl8p2.ffs@tglx
---
V1 -> V2: Multiple changes
* Fix racy return value read [**]
* Add Kconfig option (Thomas)
* Add cpumask to track NMI delivery (Boris)
* Rework implementation. Consolidate under ifdef

[**]: Observed delay in IPI/NMI delivery when testing error return paths
---
 arch/Kconfig                 |  3 ++
 include/linux/stop_machine.h | 16 ++++++
 kernel/stop_machine.c        | 96 +++++++++++++++++++++++++++++++++++-
 3 files changed, 114 insertions(+), 1 deletion(-)

diff --git a/arch/Kconfig b/arch/Kconfig
index 102ddbd4298e..f84fd528aae7 100644
--- a/arch/Kconfig
+++ b/arch/Kconfig
@@ -1841,4 +1841,7 @@ config ARCH_WANTS_PRE_LINK_VMLINUX
 config ARCH_HAS_CPU_ATTACK_VECTORS
 	bool
 
+config STOP_MACHINE_NMI
+	bool
+
 endmenu
diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index 2f986555113a..9424d363ab38 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -19,6 +19,12 @@
  */
 typedef int (*cpu_stop_fn_t)(void *arg);
 
+/*
+ * Stop function variant runnable from NMI context. This makes the
+ * noinstr requirement explicit at the type level.
+ */
+typedef int (*cpu_stop_nmisafe_fn_t)(void *arg);
+
 #ifdef CONFIG_SMP
 
 struct cpu_stop_work {
@@ -189,4 +195,14 @@ stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 }
 
 #endif	/* CONFIG_SMP || CONFIG_HOTPLUG_CPU */
+
+#ifdef CONFIG_STOP_MACHINE_NMI
+
+void arch_send_self_nmi(void);
+bool noinstr stop_machine_nmi_handler(void);
+
+#else
+static inline bool stop_machine_nmi_handler(void) { return false; }
+#endif /* CONFIG_STOP_MACHINE_NMI */
+
 #endif	/* _LINUX_STOP_MACHINE */
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 092c65c002ff..45ea62f1b2b5 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -22,6 +22,7 @@
 #include <linux/atomic.h>
 #include <linux/nmi.h>
 #include <linux/sched/wake_q.h>
+#include <linux/delay.h>
 
 /*
  * Structure to determine completion condition and record errors.  May
@@ -174,6 +175,15 @@ struct multi_stop_data {
 
 	enum multi_stop_state	state;
 	atomic_t		thread_ack;
+
+#ifdef CONFIG_STOP_MACHINE_NMI
+	/* Used in the NMI stop_machine variant */
+	bool			use_nmi;
+	/* A separate function type to highlight noinstr requirement */
+	cpu_stop_nmisafe_fn_t	nmisafe_fn;
+	/* cpumask of CPUs on which to raise an NMI */
+	cpumask_var_t		nmi_cpus;
+#endif
 };
 
 static void set_state(struct multi_stop_data *msdata,
@@ -197,6 +207,8 @@ notrace void __weak stop_machine_yield(const struct cpumask *cpumask)
 	cpu_relax();
 }
 
+static int multi_stop_run(struct multi_stop_data *msdata);
+
 /* This is the cpu_stop function which stops the CPU. */
 static int multi_cpu_stop(void *data)
 {
@@ -235,7 +247,7 @@ static int multi_cpu_stop(void *data)
 				break;
 			case MULTI_STOP_RUN:
 				if (is_active)
-					err = msdata->fn(msdata->data);
+					err = multi_stop_run(msdata);
 				break;
 			default:
 				break;
@@ -712,3 +724,85 @@ int stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 	mutex_unlock(&stop_cpus_mutex);
 	return ret | done.ret;
 }
+
+#ifdef CONFIG_STOP_MACHINE_NMI
+
+struct nmi_stop {
+	struct multi_stop_data	*data;
+	int			ret;
+	bool			done;
+};
+
+static DEFINE_PER_CPU(struct nmi_stop, nmi_stop);
+
+/*
+ * Instrumentation may trigger nested exceptions such as #INT3, #DB,
+ * or #PF. IRET from those would re-enable NMIs, which opposes the goal
+ * of this NMI stop-machine facility.
+ */
+bool noinstr stop_machine_nmi_handler(void)
+{
+	struct multi_stop_data *msdata = raw_cpu_read(nmi_stop.data);
+	unsigned int cpu = smp_processor_id();
+	int ret;
+
+	if (!msdata || !cpumask_test_and_clear_cpu(cpu, msdata->nmi_cpus))
+		return false;
+
+	/*
+	 * The indirect call to @nmisafe_fn() is indistinguishable at
+	 * post-compilation. Temporarily enabling instrumentation avoids
+	 * objtool false positives.
+	 */
+	instrumentation_begin();
+	ret = msdata->nmisafe_fn(msdata->data);
+	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;
+}
+
+static bool wait_for_nmi_handler(void)
+{
+	/* Conservative timeout */
+	unsigned long timeout = USEC_PER_SEC;
+
+	while (!this_cpu_read(nmi_stop.done) && timeout--)
+		udelay(1);
+
+	return this_cpu_read(nmi_stop.done);
+}
+
+static int nmi_stop_run(struct multi_stop_data *msdata)
+{
+	/*
+	 * Save per-CPU state accessible from NMI context and raise a
+	 * 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())
+		return -ETIMEDOUT;
+
+	return this_cpu_read(nmi_stop.ret);
+}
+
+static int multi_stop_run(struct multi_stop_data *msdata)
+{
+	return msdata->use_nmi ? nmi_stop_run(msdata) : msdata->fn(msdata->data);
+}
+
+#else
+
+static int multi_stop_run(struct multi_stop_data *msdata)
+{
+	return msdata->fn(msdata->data);
+}
+
+#endif /* CONFIG_STOP_MACHINE_NMI */
-- 
2.51.0

Re: [PATCH v2 04/11] stop_machine: Add NMI-based execution path

[Chang S. Bae]

[-- Attachment #1: Type: text/plain, Size: 1242 bytes --]

Like others, I also checked the review bot:

  https://sashiko.dev/#/patchset/20260331014251.86353-1-chang.seok.bae@intel.com

I thought all of comments could be converged in this patch. My take is 
below.

On 3/30/2026 6:42 PM, Chang S. Bae wrote:
> 
> +struct nmi_stop {
> +	struct multi_stop_data	*data;
> +	int			ret;
> +	bool			done;

The intention was to make it clear at the waiting loop. But ->data == 
NULL check can substitute it and a single variable looks to make it more 
robust and simple.

> +bool noinstr stop_machine_nmi_handler(void)
> +{
> +	struct multi_stop_data *msdata = raw_cpu_read(nmi_stop.data);
> +	unsigned int cpu = smp_processor_id();
> +	int ret;
> +
> +	if (!msdata || !cpumask_test_and_clear_cpu(cpu, msdata->nmi_cpus))
> +		return false;

smp_processor_id() and cpumask_test_and_clear_cpu() are wrappers that 
could include instrumentation, so not suitable here. Instead, 
raw_smp_processor_id() and arch_test_and_clear_bit() are inner functions.


> +	/* Ensure the handler went through before reading the result */
> +	if (!wait_for_nmi_handler())
> +		return -ETIMEDOUT;

On error exit, the stop_data pointer should be cleaned up as well before 
it is freed later.


Attached is the diff addressing them.

[-- Attachment #2: patch4.diff --]
[-- Type: text/plain, Size: 2107 bytes --]

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);
 }

[PATCH v2 05/11] stop_machine: Introduce stop_machine_nmi_cpuslocked()

[Chang S. Bae]

With the NMI control logic in place, introduce an API to run the target
function from NMI context.

Originally-by: David Kaplan <david.kaplan@amd.com>
Suggested-by: Borislav Petkov <bp@alien8.de>
Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
Link: https://lore.kernel.org/lkml/20260202105411.GVaYCCUygtEUNrMUtG@fat_crate.local
---
V1 -> V2:
* Support nmi_cpus mask (Boris), including @cpus=NULL cases
* Split out API introduction
* Drop out stop_machine_nmi() [**]

[**] could be added but no user yet in this series
---
 include/linux/stop_machine.h | 24 ++++++++++++++++++++
 kernel/stop_machine.c        | 43 ++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index 9424d363ab38..2da9aa0ec3d3 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -201,6 +201,30 @@ stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 void arch_send_self_nmi(void);
 bool noinstr stop_machine_nmi_handler(void);
 
+/**
+ * stop_machine_nmi_cpuslocked() - Freeze CPUs and run a function in NMI context
+ *
+ * @nmisafe_fn:	The function to run
+ * @data:	The data pointer for @nmisafe_fn()
+ * @cpus:	A cpumask containing the CPUs to run @nmisafe_fn() on. If NULL,
+ *		@nmisafe_fn() runs on a single (arbitrary) CPU from
+ *		cpu_online_mask.
+ *
+ * Description: This stop_machine() variant runs @nmisafe_fn() from NMI context
+ * to prevent preemption by other NMIs. The callback must be built with noinstr.
+ * Other than that, the semantics match stop_machine_cpuslocked().
+ *
+ * Context: Must be called from within a cpus_read_lock() protected region.
+ * Avoid nested calls to cpus_read_lock().
+ *
+ * Return: 0 if all invocations of @nmisafe_fn return zero, -ENOMEM if cpumask
+ * allocation fails, -EINVAL if any target CPU failed to receive NMI. Otherwise,
+ * an accumulated return value from all invocation of @nmisafe_fn that returned
+ * non-zero.
+ */
+int stop_machine_nmi_cpuslocked(cpu_stop_nmisafe_fn_t nmisafe_fn, void *data,
+				const struct cpumask *cpus);
+
 #else
 static inline bool stop_machine_nmi_handler(void) { return false; }
 #endif /* CONFIG_STOP_MACHINE_NMI */
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 45ea62f1b2b5..e20e4d3e7b16 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -798,6 +798,49 @@ static int multi_stop_run(struct multi_stop_data *msdata)
 	return msdata->use_nmi ? nmi_stop_run(msdata) : msdata->fn(msdata->data);
 }
 
+int stop_machine_nmi_cpuslocked(cpu_stop_nmisafe_fn_t nmisafe_fn, void *data,
+				const struct cpumask *cpus)
+{
+	struct multi_stop_data msdata = {
+		.nmisafe_fn	= nmisafe_fn,
+		.data		= data,
+		.num_threads	= num_online_cpus(),
+		.active_cpus	= cpus,
+		.use_nmi	= true,
+	};
+	int ret;
+
+	if (!zalloc_cpumask_var(&msdata.nmi_cpus, GFP_KERNEL))
+		return -ENOMEM;
+
+	/*
+	 * NMI CPUs should be exactly those 'active' CPUs executing the
+	 * stop function. Follow the selection logic in multi_cpu_stop()
+	 * if not provided.
+	 */
+	if (!msdata.active_cpus)
+		cpumask_set_cpu(cpumask_first(cpu_online_mask), msdata.nmi_cpus);
+	else
+		cpumask_copy(msdata.nmi_cpus, msdata.active_cpus);
+
+	lockdep_assert_cpus_held();
+
+	ret = stop_multi_cpus(&msdata);
+
+	/*
+	 * The NMI handler clears each CPU bit. If any of those NMIs were
+	 * ever missed out, return error clearly.
+	 */
+	if (!cpumask_empty(msdata.nmi_cpus)) {
+		pr_err("CPUs %*pbl didn't run the stop_machine NMI handler.\n",
+		       cpumask_pr_args(msdata.nmi_cpus));
+		ret = -EINVAL;
+	}
+
+	free_cpumask_var(msdata.nmi_cpus);
+	return ret;
+}
+
 #else
 
 static int multi_stop_run(struct multi_stop_data *msdata)
-- 
2.51.0

[PATCH v2 06/11] x86/apic: Implement self-NMI support

[Chang S. Bae]

From: David Kaplan <david.kaplan@amd.com>

Add a function to send an NMI-to-self which is needed for stop-machine
NMI facility. Note that send_IPI_self() cannot be used to send an NMI so
send_IPI() is used.

Signed-off-by: David Kaplan <david.kaplan@amd.com>
Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: Fix a typo (Thomas) and wrap it under the new build option
---
 arch/x86/kernel/apic/ipi.c | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/arch/x86/kernel/apic/ipi.c b/arch/x86/kernel/apic/ipi.c
index 98a57cb4aa86..597f03192a49 100644
--- a/arch/x86/kernel/apic/ipi.c
+++ b/arch/x86/kernel/apic/ipi.c
@@ -4,6 +4,7 @@
 #include <linux/delay.h>
 #include <linux/smp.h>
 #include <linux/string_choices.h>
+#include <linux/stop_machine.h>
 
 #include <asm/io_apic.h>
 
@@ -248,6 +249,15 @@ void default_send_IPI_self(int vector)
 	__default_send_IPI_shortcut(APIC_DEST_SELF, vector);
 }
 
+#ifdef CONFIG_STOP_MACHINE_NMI
+
+void arch_send_self_nmi(void)
+{
+	apic->send_IPI(smp_processor_id(), NMI_VECTOR);
+}
+
+#endif
+
 #ifdef CONFIG_X86_32
 void default_send_IPI_mask_sequence_logical(const struct cpumask *mask, int vector)
 {
-- 
2.51.0

[PATCH v2 07/11] x86/nmi: Support NMI stop-machine handler

[Chang S. Bae]

From: David Kaplan <david.kaplan@amd.com>

Call the stop-machine handler from the NMI path in order to support the
NMI stop-machine.

Signed-off-by: David Kaplan <david.kaplan@amd.com>
Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: Switch away from static key reference
---
 arch/x86/kernel/nmi.c | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/arch/x86/kernel/nmi.c b/arch/x86/kernel/nmi.c
index 3d239ed12744..b7ea2907142c 100644
--- a/arch/x86/kernel/nmi.c
+++ b/arch/x86/kernel/nmi.c
@@ -24,6 +24,7 @@
 #include <linux/export.h>
 #include <linux/atomic.h>
 #include <linux/sched/clock.h>
+#include <linux/stop_machine.h>
 #include <linux/kvm_types.h>
 
 #include <asm/cpu_entry_area.h>
@@ -382,6 +383,9 @@ static noinstr void default_do_nmi(struct pt_regs *regs)
 
 	instrumentation_begin();
 
+	if (stop_machine_nmi_handler())
+		goto out;
+
 	if (microcode_nmi_handler_enabled() && microcode_nmi_handler())
 		goto out;
 
-- 
2.51.0

[PATCH v2 08/11] x86/microcode: Distinguish NMI control path on stop-machine callback

[Chang S. Bae]

load_cpus_stopped() currently centralizes the stop_machine() callback for
both NMI and NMI-less rendezvous. microcode_update_handler() alone is
enough for the latter.

While the NMI-based rendezvous finally reaches the same update handler,
it requires additional logic to trigger and process NMIs. That machinery
will be replaced by stop-machine facility.

As preparation for that conversion, split the callback path to make
NMI-specific steps explicit and clear. Rename the function to align with
the change.

Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: Rename load_cpus_stopped() (Thomas)
---
 arch/x86/kernel/cpu/microcode/core.c | 30 ++++++++++++----------------
 1 file changed, 13 insertions(+), 17 deletions(-)

diff --git a/arch/x86/kernel/cpu/microcode/core.c b/arch/x86/kernel/cpu/microcode/core.c
index 651202e6fefb..abd640b1d286 100644
--- a/arch/x86/kernel/cpu/microcode/core.c
+++ b/arch/x86/kernel/cpu/microcode/core.c
@@ -532,7 +532,7 @@ void noinstr microcode_offline_nmi_handler(void)
 	wait_for_ctrl();
 }
 
-static noinstr bool microcode_update_handler(void)
+static noinstr int microcode_update_handler(void *unused)
 {
 	unsigned int cpu = raw_smp_processor_id();
 
@@ -548,7 +548,7 @@ static noinstr bool microcode_update_handler(void)
 	touch_nmi_watchdog();
 	instrumentation_end();
 
-	return true;
+	return 0;
 }
 
 /*
@@ -569,19 +569,15 @@ bool noinstr microcode_nmi_handler(void)
 		return false;
 
 	raw_cpu_write(ucode_ctrl.nmi_enabled, false);
-	return microcode_update_handler();
+	return microcode_update_handler(NULL) == 0;
 }
 
-static int load_cpus_stopped(void *unused)
+static int stop_cpu_in_nmi(void *unused)
 {
-	if (microcode_ops->use_nmi) {
-		/* Enable the NMI handler and raise NMI */
-		this_cpu_write(ucode_ctrl.nmi_enabled, true);
-		apic->send_IPI(smp_processor_id(), NMI_VECTOR);
-	} else {
-		/* Just invoke the handler directly */
-		microcode_update_handler();
-	}
+	/* Enable the NMI handler and raise NMI */
+	this_cpu_write(ucode_ctrl.nmi_enabled, true);
+	apic->send_IPI(smp_processor_id(), NMI_VECTOR);
+
 	return 0;
 }
 
@@ -618,13 +614,13 @@ static int load_late_stop_cpus(bool is_safe)
 	 */
 	store_cpu_caps(&prev_info);
 
-	if (microcode_ops->use_nmi)
+	if (microcode_ops->use_nmi) {
 		static_branch_enable_cpuslocked(&microcode_nmi_handler_enable);
-
-	stop_machine_cpuslocked(load_cpus_stopped, NULL, cpu_online_mask);
-
-	if (microcode_ops->use_nmi)
+		stop_machine_cpuslocked(stop_cpu_in_nmi, NULL, cpu_online_mask);
 		static_branch_disable_cpuslocked(&microcode_nmi_handler_enable);
+	} else {
+		stop_machine_cpuslocked(microcode_update_handler, NULL, cpu_online_mask);
+	}
 
 	/* Analyze the results */
 	for_each_cpu_and(cpu, cpu_present_mask, &cpus_booted_once_mask) {
-- 
2.51.0

[PATCH v2 09/11] x86/microcode: Use stop-machine NMI facility

[Chang S. Bae]

The existing NMI-based loading logic explicitly sends NMIs to online CPUs
and invokes microcode_update_handler() from the NMI context. The
stop-machine NMI variant already provides the mechanism on x86.

Replace the custom NMI control logic with stop_machine_nmi_cpuslocked().

Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: Select that stop-machine build option
---
 arch/x86/Kconfig                     |  1 +
 arch/x86/include/asm/microcode.h     |  1 -
 arch/x86/kernel/cpu/microcode/core.c | 19 +++----------------
 arch/x86/kernel/nmi.c                |  3 ---
 4 files changed, 4 insertions(+), 20 deletions(-)

diff --git a/arch/x86/Kconfig b/arch/x86/Kconfig
index 0b5f30d769ff..0f7e88ba7433 100644
--- a/arch/x86/Kconfig
+++ b/arch/x86/Kconfig
@@ -1338,6 +1338,7 @@ config MICROCODE_LATE_LOADING
 	bool "Late microcode loading (DANGEROUS)"
 	default n
 	depends on MICROCODE && SMP
+	select STOP_MACHINE_NMI
 	help
 	  Loading microcode late, when the system is up and executing instructions
 	  is a tricky business and should be avoided if possible. Just the sequence
diff --git a/arch/x86/include/asm/microcode.h b/arch/x86/include/asm/microcode.h
index 3c317d155771..62d10c43da9c 100644
--- a/arch/x86/include/asm/microcode.h
+++ b/arch/x86/include/asm/microcode.h
@@ -79,7 +79,6 @@ static inline u32 intel_get_microcode_revision(void)
 }
 #endif /* !CONFIG_CPU_SUP_INTEL */
 
-bool microcode_nmi_handler(void);
 void microcode_offline_nmi_handler(void);
 
 #ifdef CONFIG_MICROCODE_LATE_LOADING
diff --git a/arch/x86/kernel/cpu/microcode/core.c b/arch/x86/kernel/cpu/microcode/core.c
index abd640b1d286..ebcc73e67af1 100644
--- a/arch/x86/kernel/cpu/microcode/core.c
+++ b/arch/x86/kernel/cpu/microcode/core.c
@@ -563,22 +563,9 @@ static noinstr int microcode_update_handler(void *unused)
  * path which must be NMI safe until the primary thread completed the
  * update.
  */
-bool noinstr microcode_nmi_handler(void)
+static noinstr int microcode_nmi_handler(void *data)
 {
-	if (!raw_cpu_read(ucode_ctrl.nmi_enabled))
-		return false;
-
-	raw_cpu_write(ucode_ctrl.nmi_enabled, false);
-	return microcode_update_handler(NULL) == 0;
-}
-
-static int stop_cpu_in_nmi(void *unused)
-{
-	/* Enable the NMI handler and raise NMI */
-	this_cpu_write(ucode_ctrl.nmi_enabled, true);
-	apic->send_IPI(smp_processor_id(), NMI_VECTOR);
-
-	return 0;
+	return microcode_update_handler(data);
 }
 
 static int load_late_stop_cpus(bool is_safe)
@@ -616,7 +603,7 @@ static int load_late_stop_cpus(bool is_safe)
 
 	if (microcode_ops->use_nmi) {
 		static_branch_enable_cpuslocked(&microcode_nmi_handler_enable);
-		stop_machine_cpuslocked(stop_cpu_in_nmi, NULL, cpu_online_mask);
+		stop_machine_nmi_cpuslocked(microcode_nmi_handler, NULL, cpu_online_mask);
 		static_branch_disable_cpuslocked(&microcode_nmi_handler_enable);
 	} else {
 		stop_machine_cpuslocked(microcode_update_handler, NULL, cpu_online_mask);
diff --git a/arch/x86/kernel/nmi.c b/arch/x86/kernel/nmi.c
index b7ea2907142c..324f4353be88 100644
--- a/arch/x86/kernel/nmi.c
+++ b/arch/x86/kernel/nmi.c
@@ -386,9 +386,6 @@ static noinstr void default_do_nmi(struct pt_regs *regs)
 	if (stop_machine_nmi_handler())
 		goto out;
 
-	if (microcode_nmi_handler_enabled() && microcode_nmi_handler())
-		goto out;
-
 	/*
 	 * CPU-specific NMI must be processed before non-CPU-specific
 	 * NMI, otherwise we may lose it, because the CPU-specific
-- 
2.51.0

[PATCH v2 10/11] x86/nmi: Simplify offline microcode handler invocation

[Chang S. Bae]

The microcode loader now uses the stop-machine facility for NMI-based
loading. The offline handler is the only remaining user of the microcode
NMI static key.

But the handler already has internal guarding. Offline CPU doesn't matter
how fast it handles anyway, so remove the static key. Instead, check the
build option.

Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
V1 -> V2: Switch away from static key reference
---
 arch/x86/kernel/nmi.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arch/x86/kernel/nmi.c b/arch/x86/kernel/nmi.c
index 324f4353be88..4ca891f7075f 100644
--- a/arch/x86/kernel/nmi.c
+++ b/arch/x86/kernel/nmi.c
@@ -548,7 +548,7 @@ DEFINE_IDTENTRY_RAW(exc_nmi)
 		raw_atomic_long_inc(&nsp->idt_calls);
 
 	if (arch_cpu_is_offline(smp_processor_id())) {
-		if (microcode_nmi_handler_enabled())
+		if (IS_ENABLED(CONFIG_MICROCODE_LATE_LOADING))
 			microcode_offline_nmi_handler();
 		return;
 	}
@@ -711,7 +711,7 @@ DEFINE_FREDENTRY_NMI(exc_nmi)
 	irqentry_state_t irq_state;
 
 	if (arch_cpu_is_offline(smp_processor_id())) {
-		if (microcode_nmi_handler_enabled())
+		if (IS_ENABLED(CONFIG_MICROCODE_LATE_LOADING))
 			microcode_offline_nmi_handler();
 		return;
 	}
-- 
2.51.0

[PATCH v2 11/11] x86/microcode: Remove microcode_nmi_handler_enable

[Chang S. Bae]

Remove it, as there is no more user.

Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
 arch/x86/include/asm/microcode.h     | 10 ----------
 arch/x86/kernel/cpu/microcode/core.c |  8 ++------
 2 files changed, 2 insertions(+), 16 deletions(-)

diff --git a/arch/x86/include/asm/microcode.h b/arch/x86/include/asm/microcode.h
index 62d10c43da9c..a250a8849168 100644
--- a/arch/x86/include/asm/microcode.h
+++ b/arch/x86/include/asm/microcode.h
@@ -81,14 +81,4 @@ static inline u32 intel_get_microcode_revision(void)
 
 void microcode_offline_nmi_handler(void);
 
-#ifdef CONFIG_MICROCODE_LATE_LOADING
-DECLARE_STATIC_KEY_FALSE(microcode_nmi_handler_enable);
-static __always_inline bool microcode_nmi_handler_enabled(void)
-{
-	return static_branch_unlikely(&microcode_nmi_handler_enable);
-}
-#else
-static __always_inline bool microcode_nmi_handler_enabled(void) { return false; }
-#endif
-
 #endif /* _ASM_X86_MICROCODE_H */
diff --git a/arch/x86/kernel/cpu/microcode/core.c b/arch/x86/kernel/cpu/microcode/core.c
index ebcc73e67af1..775233c069c7 100644
--- a/arch/x86/kernel/cpu/microcode/core.c
+++ b/arch/x86/kernel/cpu/microcode/core.c
@@ -324,7 +324,6 @@ struct microcode_ctrl {
 	bool			nmi_enabled;
 };
 
-DEFINE_STATIC_KEY_FALSE(microcode_nmi_handler_enable);
 static DEFINE_PER_CPU(struct microcode_ctrl, ucode_ctrl);
 static atomic_t late_cpus_in, offline_in_nmi;
 static unsigned int loops_per_usec;
@@ -601,13 +600,10 @@ static int load_late_stop_cpus(bool is_safe)
 	 */
 	store_cpu_caps(&prev_info);
 
-	if (microcode_ops->use_nmi) {
-		static_branch_enable_cpuslocked(&microcode_nmi_handler_enable);
+	if (microcode_ops->use_nmi)
 		stop_machine_nmi_cpuslocked(microcode_nmi_handler, NULL, cpu_online_mask);
-		static_branch_disable_cpuslocked(&microcode_nmi_handler_enable);
-	} else {
+	else
 		stop_machine_cpuslocked(microcode_update_handler, NULL, cpu_online_mask);
-	}
 
 	/* Analyze the results */
 	for_each_cpu_and(cpu, cpu_present_mask, &cpus_booted_once_mask) {
-- 
2.51.0