[PATCH v4 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
---
V3 -> V4: Remove the fallback prototype

  Tracking back the previous feedback on V1,
	  https://lore.kernel.org/lkml/87wm0zl8p2.ffs@tglx
  tglx gave this comment "That function should not be exposed at all when
  the config switch is off."

  So, the x86 microcode loader is the only user and the callsite should be
  compiled out when CONFIG_STOP_MACHINE_NMI=n.
---
 include/linux/stop_machine.h | 22 +++++++++++++++++
 kernel/stop_machine.c        | 46 ++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index b53b028296e0..af3218987f53 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -198,6 +198,28 @@ 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; }
 #endif /* CONFIG_STOP_MACHINE_NMI */
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