[RFC PATCH 28/56] stop_machine: Add stop_machine_nmi()

David Kaplan
2025-10-13T14:35:55Z

stop_machine_nmi() is a flavor of stop_machine() that runs the specified
function in NMI context. This is useful for flows that cannot tolerate any
risk of interruption even due to an NMI. Arch-specific code must handle
sending the actual NMI and running the stop_machine_nmi_handler().

Signed-off-by: David Kaplan <david.kaplan@amd.com>
---
include/linux/stop_machine.h | 32 +++++++++++++++
kernel/stop_machine.c | 79 ++++++++++++++++++++++++++++++++++--
2 files changed, 107 insertions(+), 4 deletions(-)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index 72820503514c..452c45640012 100644
--- a/include/linux/stop_machine.h
+++ b/include/linux/stop_machine.h
@@ -141,6 +141,29 @@ int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus);
*/
int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus);

+/**
+ * stop_machine_nmi: freeze the machine and run this function in NMI context
+ * @fn: the function to run
+ * @data: the data ptr for the @fn()
+ * @cpus: the cpus to run the @fn() on (NULL = any online cpu)
+ *
+ * Like stop_machine() but runs the function in NMI context to avoid any risk of
+ * interruption due to NMIs.
+ *
+ * Protects against CPU hotplug.
+ */
+int stop_machine_nmi(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus);
+
+/**
+ * stop_machine_cpuslocked_nmi: freeze and run this function in NMI context
+ * @fn: the function to run
+ * @data: the data ptr for the @fn()
+ * @cpus: the cpus to run the @fn() on (NULL = any online cpu)
+ *
+ * Same as above. Must be called from within a cpus_read_lock() protected
+ * region. Avoids nested calls to cpus_read_lock().
+ */
+int stop_machine_cpuslocked_nmi(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus);
/**
* stop_core_cpuslocked: - stop all threads on just one core
* @cpu: any cpu in the targeted core
@@ -160,6 +183,15 @@ int stop_core_cpuslocked(unsigned int cpu, cpu_stop_fn_t fn, void *data);

int stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
const struct cpumask *cpus);
+
+bool noinstr stop_machine_nmi_handler(void);
+void arch_send_self_nmi(void);
+DECLARE_STATIC_KEY_FALSE(stop_machine_nmi_handler_enable);
+static __always_inline bool stop_machine_nmi_handler_enabled(void)
+{
+ return static_branch_unlikely(&stop_machine_nmi_handler_enable);
+}
+
#else /* CONFIG_SMP || CONFIG_HOTPLUG_CPU */

static __always_inline int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data,
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 3fe6b0c99f3d..d135f32528e8 100644
--- a/kernel/stop_machine.c
+++ b/kernel/stop_machine.c
@@ -174,6 +174,8 @@ struct multi_stop_data {

enum multi_stop_state state;
atomic_t thread_ack;
+
+ bool use_nmi;
};

static void set_state(struct multi_stop_data *msdata,
@@ -197,6 +199,42 @@ notrace void __weak stop_machine_yield(const struct cpumask *cpumask)
cpu_relax();
}

+struct stop_machine_nmi_ctrl {
+ bool nmi_enabled;
+ struct multi_stop_data *msdata;
+ int err;
+};
+
+DEFINE_STATIC_KEY_FALSE(stop_machine_nmi_handler_enable);
+static DEFINE_PER_CPU(struct stop_machine_nmi_ctrl, stop_machine_nmi_ctrl);
+
+static void enable_nmi_handler(struct multi_stop_data *msdata)
+{
+ this_cpu_write(stop_machine_nmi_ctrl.msdata, msdata);
+ this_cpu_write(stop_machine_nmi_ctrl.nmi_enabled, true);
+}
+
+void __weak arch_send_self_nmi(void)
+{
+ /* Arch code must implement this to support stop_machine_nmi() */
+}
+
+bool noinstr stop_machine_nmi_handler(void)
+{
+ struct multi_stop_data *msdata;
+ int err;
+
+ if (!raw_cpu_read(stop_machine_nmi_ctrl.nmi_enabled))
+ return false;
+
+ raw_cpu_write(stop_machine_nmi_ctrl.nmi_enabled, false);
+
+ msdata = raw_cpu_read(stop_machine_nmi_ctrl.msdata);
+ err = msdata->fn(msdata->data);
+ raw_cpu_write(stop_machine_nmi_ctrl.err, err);
+ return true;
+}
+
/* This is the cpu_stop function which stops the CPU. */
static int multi_cpu_stop(void *data)
{
@@ -234,8 +272,15 @@ static int multi_cpu_stop(void *data)
hard_irq_disable();
break;
case MULTI_STOP_RUN:
- if (is_active)
- err = msdata->fn(msdata->data);
+ if (is_active) {
+ if (msdata->use_nmi) {
+ enable_nmi_handler(msdata);
+ arch_send_self_nmi();
+ err = raw_cpu_read(stop_machine_nmi_ctrl.err);
+ } else {
+ err = msdata->fn(msdata->data);
+ }
+ }
break;
default:
break;
@@ -584,14 +629,15 @@ 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_machine_cpuslocked(cpu_stop_fn_t fn, void *data,
+ const struct cpumask *cpus, bool use_nmi)
{
struct multi_stop_data msdata = {
.fn = fn,
.data = data,
.num_threads = num_online_cpus(),
.active_cpus = cpus,
+ .use_nmi = use_nmi,
};

lockdep_assert_cpus_held();
@@ -620,6 +666,18 @@ int stop_machine_cpuslocked(cpu_stop_fn_t fn, void *data,
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)
+{
+ return __stop_machine_cpuslocked(fn, data, cpus, false);
+}
+
+int stop_machine_cpuslocked_nmi(cpu_stop_fn_t fn, void *data,
+ const struct cpumask *cpus)
+{
+ return __stop_machine_cpuslocked(fn, data, cpus, true);
+}
+
int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus)
{
int ret;
@@ -632,6 +690,19 @@ int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus)
}
EXPORT_SYMBOL_GPL(stop_machine);

+int stop_machine_nmi(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus)
+{
+ int ret;
+
+ cpus_read_lock();
+ static_branch_enable_cpuslocked(&stop_machine_nmi_handler_enable);
+ ret = stop_machine_cpuslocked_nmi(fn, data, cpus);
+ static_branch_disable_cpuslocked(&stop_machine_nmi_handler_enable);
+ cpus_read_unlock();
+ return ret;
+}
+EXPORT_SYMBOL_GPL(stop_machine_nmi);
+
#ifdef CONFIG_SCHED_SMT
int stop_core_cpuslocked(unsigned int cpu, cpu_stop_fn_t fn, void *data)
{
--
2.34.1