@@ -12,6 +12,11 @@ int set_handle_irq(void (*handle_irq)(struct pt_regs *));
 #define set_handle_irq	set_handle_irq
 int set_handle_fiq(void (*handle_fiq)(struct pt_regs *));
 
+extern void (*handle_arch_irq)(struct pt_regs *regs);
+extern void (*handle_arch_fiq)(struct pt_regs *regs);
+extern void (*handle_arch_nmi)(struct pt_regs *regs);
+extern bool (*interrupt_is_nmi)(void);
+
 static inline int nr_legacy_irqs(void)
 {
 	return 0;

@@ -24,6 +24,7 @@
 #include <asm/stacktrace.h>
 #include <asm/sysreg.h>
 #include <asm/system_misc.h>
+#include <asm/irq.h>
 
 /*
  * Handle IRQ/context state management when entering from kernel mode.

@@ -219,17 +220,28 @@ static void noinstr arm64_exit_el1_dbg(struct pt_regs *regs)
 		lockdep_hardirqs_on(CALLER_ADDR0);
 }
 
-static void noinstr enter_el1_irq_or_nmi(struct pt_regs *regs)
+static inline bool arm64_in_nmi(struct pt_regs *regs)
 {
-	if (IS_ENABLED(CONFIG_ARM64_PSEUDO_NMI) && !interrupts_enabled(regs))
+	if (!interrupts_enabled(regs) || (*interrupt_is_nmi)())
+		return true;
+	return false;
+}
+
+/* return true if in irq, otherwise in nmi */
+static bool noinstr enter_el1_irq_or_nmi(struct pt_regs *regs)
+{
+	if (IS_ENABLED(CONFIG_ARM64_PSEUDO_NMI) && arm64_in_nmi(regs)) {
 		arm64_enter_nmi(regs);
-	else
+		return false;
+	} else {
 		enter_from_kernel_mode(regs);
+		return true;
+	}
 }
 
-static void noinstr exit_el1_irq_or_nmi(struct pt_regs *regs)
+static void noinstr exit_el1_irq_or_nmi(struct pt_regs *regs, bool in_irq)
 {
-	if (IS_ENABLED(CONFIG_ARM64_PSEUDO_NMI) && !interrupts_enabled(regs))
+	if (IS_ENABLED(CONFIG_ARM64_PSEUDO_NMI) && !in_irq)
 		arm64_exit_nmi(regs);
 	else
 		exit_to_kernel_mode(regs);

@@ -269,9 +281,6 @@ static void do_interrupt_handler(struct pt_regs *regs,
 		handler(regs);
 }
 
-extern void (*handle_arch_irq)(struct pt_regs *);
-extern void (*handle_arch_fiq)(struct pt_regs *);
-
 static void noinstr __panic_unhandled(struct pt_regs *regs, const char *vector,
 				      unsigned int esr)
 {

@@ -433,12 +442,20 @@ asmlinkage void noinstr el1h_64_sync_handler(struct pt_regs *regs)
 }
 
 static void noinstr el1_interrupt(struct pt_regs *regs,
-				  void (*handler)(struct pt_regs *))
+				  void (*handler)(struct pt_regs *),
+				  void (*nmi_handler)(struct pt_regs *))
 {
+	bool in_irq;
+	void (*h)(struct pt_regs *regs);
+
 	write_sysreg(DAIF_PROCCTX_NOIRQ, daif);
 
-	enter_el1_irq_or_nmi(regs);
-	do_interrupt_handler(regs, handler);
+	in_irq = enter_el1_irq_or_nmi(regs);
+	if (in_irq)
+		h = handler;
+	else
+		h = nmi_handler;
+	do_interrupt_handler(regs, h);
 
 	/*
 	 * Note: thread_info::preempt_count includes both thread_info::count

@@ -449,17 +466,17 @@ static void noinstr el1_interrupt(struct pt_regs *regs,
 	    READ_ONCE(current_thread_info()->preempt_count) == 0)
 		arm64_preempt_schedule_irq();
 
-	exit_el1_irq_or_nmi(regs);
+	exit_el1_irq_or_nmi(regs, in_irq);
 }
 
 asmlinkage void noinstr el1h_64_irq_handler(struct pt_regs *regs)
 {
-	el1_interrupt(regs, handle_arch_irq);
+	el1_interrupt(regs, handle_arch_irq, handle_arch_nmi);
 }
 
 asmlinkage void noinstr el1h_64_fiq_handler(struct pt_regs *regs)
 {
-	el1_interrupt(regs, handle_arch_fiq);
+	el1_interrupt(regs, handle_arch_fiq, handle_arch_nmi);
 }
 
 asmlinkage void noinstr el1h_64_error_handler(struct pt_regs *regs)

@@ -81,8 +81,19 @@ static void default_handle_fiq(struct pt_regs *regs)
 	panic("FIQ taken without a root FIQ handler\n");
 }
 
+static void default_handle_nmi(struct pt_regs *unused)
+{
+}
+
+static bool default_nmi_discriminator(void)
+{
+	return false;
+}
+
 void (*handle_arch_irq)(struct pt_regs *) __ro_after_init = default_handle_irq;
 void (*handle_arch_fiq)(struct pt_regs *) __ro_after_init = default_handle_fiq;
+void (*handle_arch_nmi)(struct pt_regs *) __ro_after_init = default_handle_nmi;
+bool (*interrupt_is_nmi)(void) __ro_after_init = default_nmi_discriminator;
 
 int __init set_handle_irq(void (*handle_irq)(struct pt_regs *))
 {

@@ -104,6 +115,24 @@ int __init set_handle_fiq(void (*handle_fiq)(struct pt_regs *))
 	return 0;
 }
 
+int __init set_handle_nmi(void (*handle_nmi)(struct pt_regs *))
+{
+	if (handle_arch_nmi != default_handle_nmi)
+		return -EBUSY;
+
+	handle_arch_nmi = handle_nmi;
+	return 0;
+}
+
+int __init set_nmi_discriminator(bool (*discriminator)(void))
+{
+	if (interrupt_is_nmi != default_nmi_discriminator)
+		return -EBUSY;
+
+	interrupt_is_nmi = discriminator;
+	return 0;
+}
+
 void __init init_IRQ(void)
 {
 	init_irq_stacks();