[PATCH v3 3/9] stop_machine: Introduce stop_machine_nmi_cpuslocked()

["Chang S. Bae" <chang.seok.bae@intel.com>]

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
---
V2 -> V3: Address Sashiko review
* Ensure runtime only via `stop_machine_initialized`
* Ensure target CPUs online
* Add a fallback prototype
* Adjust function description (to align with mainline changes)
---
 include/linux/stop_machine.h | 27 +++++++++++++++++++++
 kernel/stop_machine.c        | 46 ++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index b53b028296e0..b05e65f4bb6d 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -198,8 +198,35 @@ 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:	The CPUs to run @nmisafe_fn() on (NULL = one unspecified online CPU)
+ *
+ * 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 not initialized yet or any target CPU failed to
+ * receive NMI. Otherwise, the value returned by an arbitrarily chosen member of
+ * the set of calls to @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; }
+static inline int
+stop_machine_nmi_cpulocked(cpu_stop_nmisafe_fn_t nmisafe_fn, void *data, conststruct cpumask *cpus)
+{
+	return 0;
+}
 #endif /* CONFIG_STOP_MACHINE_NMI */
 
 #endif	/* _LINUX_STOP_MACHINE */
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index ad01b7c5266c..2f3fe13cd3aa 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -802,6 +802,52 @@ 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 (!stop_machine_initialized)
+		return -EINVAL;
+
+	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_and(msdata.nmi_cpus, msdata.active_cpus, cpu_online_mask);
+
+	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.53.0