Call the stop_machine_nmi_handler() if the appropriate static branch is
enabled in order to support stop_machine_nmi().

Signed-off-by: David Kaplan <david.kaplan@amd.com>
---
 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 be93ec7255bf..c0d8e80029a0 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 <asm/cpu_entry_area.h>
 #include <asm/traps.h>
@@ -381,6 +382,9 @@ static noinstr void default_do_nmi(struct pt_regs *regs)
 
 	instrumentation_begin();
 
+	if (stop_machine_nmi_handler_enabled() && stop_machine_nmi_handler())
+		goto out;
+
 	if (microcode_nmi_handler_enabled() && microcode_nmi_handler())
 		goto out;
 
-- 
2.34.1