#include <linux/kgdb.h>
 #include <linux/init.h>
 #include <linux/smp.h>
+#include <linux/nmi.h>
 
 #include <asm/apicdef.h>
 #include <asm/system.h>
        return NOTIFY_STOP;
 }
 
+static int was_in_debug_nmi[NR_CPUS];
+
 static int __kgdb_notify(struct die_args *args, unsigned long cmd)
 {
        struct pt_regs *regs = args->regs;
                if (atomic_read(&kgdb_active) != -1) {
                        /* KGDB CPU roundup */
                        kgdb_nmicallback(raw_smp_processor_id(), regs);
+                       was_in_debug_nmi[raw_smp_processor_id()] = 1;
+                       touch_nmi_watchdog();
                        return NOTIFY_STOP;
                }
                return NOTIFY_DONE;
 
        case DIE_NMI_IPI:
                if (atomic_read(&kgdb_active) != -1) {
-                       /* KGDB CPU roundup: */
-                       if (kgdb_nmicallback(raw_smp_processor_id(), regs))
-                               return NOTIFY_DONE;
+                       /* KGDB CPU roundup */
+                       kgdb_nmicallback(raw_smp_processor_id(), regs);
+                       was_in_debug_nmi[raw_smp_processor_id()] = 1;
+                       touch_nmi_watchdog();
+               }
+               return NOTIFY_DONE;
+
+       case DIE_NMIUNKNOWN:
+               if (was_in_debug_nmi[raw_smp_processor_id()]) {
+                       was_in_debug_nmi[raw_smp_processor_id()] = 0;
                        return NOTIFY_STOP;
                }
                return NOTIFY_DONE;
 
 static __kprobes void
 unknown_nmi_error(unsigned char reason, struct pt_regs * regs)
 {
+       if (notify_die(DIE_NMIUNKNOWN, "nmi", regs, reason, 2, SIGINT) == NOTIFY_STOP)
+               return;
 #ifdef CONFIG_MCA
        /* Might actually be able to figure out what the guilty party
        * is. */
 
 static __kprobes void
 unknown_nmi_error(unsigned char reason, struct pt_regs * regs)
 {
+       if (notify_die(DIE_NMIUNKNOWN, "nmi", regs, reason, 2, SIGINT) == NOTIFY_STOP)
+               return;
        printk(KERN_EMERG "Uhhuh. NMI received for unknown reason %02x.\n",
                reason);
        printk(KERN_EMERG "Do you have a strange power saving mode enabled?\n");