[PATCH 1/7] stop_machine: Introduce stop_machine_nmi()

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

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

stop_machine_nmi() is a variant 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>
Signed-off-by: Chang S. Bae <chang.seok.bae@intel.com>
---
Update from the original version:
* Move static key handling into stop_machine_cpuslocked_nmi() to support
  core-code users that already hold cpu hotplug locks
* Tweak the subject to better reflect the new interface and changelog a
  bit as well
---
 include/linux/stop_machine.h | 50 +++++++++++++++++++++
 kernel/stop_machine.c        | 84 ++++++++++++++++++++++++++++++++++--
 2 files changed, 130 insertions(+), 4 deletions(-)

diff --git a/include/linux/stop_machine.h b/include/linux/stop_machine.h
index 72820503514c..86113084456a 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,14 @@ 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);
+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,
@@ -186,5 +217,24 @@ stop_machine_from_inactive_cpu(cpu_stop_fn_t fn, void *data,
 	return stop_machine(fn, data, cpus);
 }
 
+/* stop_machine_nmi() is only supported in SMP systems. */
+static __always_inline int stop_machine_nmi(cpu_stop_fn_t fn, void *data,
+					const struct cpumask *cpus)
+{
+	return -EINVAL;
+}
+
+static __always_inline bool stop_machine_nmi_handler_enabled(void)
+{
+	return false;
+}
+
+static __always_inline bool stop_machine_nmi_handler(void)
+{
+	return false;
+}
+
 #endif	/* CONFIG_SMP || CONFIG_HOTPLUG_CPU */
+
+void arch_send_self_nmi(void);
 #endif	/* _LINUX_STOP_MACHINE */
diff --git a/kernel/stop_machine.c b/kernel/stop_machine.c
index 3fe6b0c99f3d..189b4b108d13 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,24 @@ 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)
+{
+	int ret;
+
+	static_branch_enable_cpuslocked(&stop_machine_nmi_handler_enable);
+	ret = __stop_machine_cpuslocked(fn, data, cpus, true);
+	static_branch_disable_cpuslocked(&stop_machine_nmi_handler_enable);
+
+	return ret;
+}
+
 int stop_machine(cpu_stop_fn_t fn, void *data, const struct cpumask *cpus)
 {
 	int ret;
@@ -632,6 +696,18 @@ 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();
+	ret = stop_machine_cpuslocked_nmi(fn, data, cpus);
+	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.51.0