1 files changed, 12 insertions, 5 deletions

diff --git a/mm/highmem.c b/mm/highmem.c
index 25878cc49daa..9c1e627f282e 100644
--- a/mm/highmem.c
+++ b/mm/highmem.c
@@ -426,16 +426,21 @@ void __init page_address_init(void)
 
 void debug_kmap_atomic(enum km_type type)
 {
-        static unsigned warn_count = 10;
+        static int warn_count = 10;
 
-        if (unlikely(warn_count == 0))
+        if (unlikely(warn_count < 0))
                 return;
 
         if (unlikely(in_interrupt())) {
-                if (in_irq()) {
+                if (in_nmi()) {
+                        if (type != KM_NMI && type != KM_NMI_PTE) {
+                                WARN_ON(1);
+                                warn_count--;
+                        }
+                } else if (in_irq()) {
                         if (type != KM_IRQ0 && type != KM_IRQ1 &&
                             type != KM_BIO_SRC_IRQ && type != KM_BIO_DST_IRQ &&
-                            type != KM_BOUNCE_READ) {
+                            type != KM_BOUNCE_READ && type != KM_IRQ_PTE) {
                                 WARN_ON(1);
                                 warn_count--;
                         }
@@ -452,7 +457,9 @@ void debug_kmap_atomic(enum km_type type)
         }
 
         if (type == KM_IRQ0 || type == KM_IRQ1 || type == KM_BOUNCE_READ ||
-                        type == KM_BIO_SRC_IRQ || type == KM_BIO_DST_IRQ) {
+                        type == KM_BIO_SRC_IRQ || type == KM_BIO_DST_IRQ ||
+                        type == KM_IRQ_PTE || type == KM_NMI ||
+                        type == KM_NMI_PTE ) {
                 if (!irqs_disabled()) {
                         WARN_ON(1);
                         warn_count--;