/**
  * munlock_vma_page - munlock a vma page
- * @page - page to be unlocked
+ * @page - page to be unlocked, either a normal page or THP page head
+ *
+ * returns the size of the page as a page mask (0 for normal page,
+ *         HPAGE_PMD_NR - 1 for THP head page)
  *
  * called from munlock()/munmap() path with page supposedly on the LRU.
  * When we munlock a page, because the vma where we found the page is being
  */
 unsigned int munlock_vma_page(struct page *page)
 {
-       unsigned int page_mask = 0;
+       unsigned int nr_pages;
 
        BUG_ON(!PageLocked(page));
 
        if (TestClearPageMlocked(page)) {
-               unsigned int nr_pages = hpage_nr_pages(page);
+               nr_pages = hpage_nr_pages(page);
                mod_zone_page_state(page_zone(page), NR_MLOCK, -nr_pages);
-               page_mask = nr_pages - 1;
                if (!isolate_lru_page(page))
                        __munlock_isolated_page(page);
                else
                        __munlock_isolation_failed(page);
+       } else {
+               nr_pages = hpage_nr_pages(page);
        }
 
-       return page_mask;
+       /*
+        * Regardless of the original PageMlocked flag, we determine nr_pages
+        * after touching the flag. This leaves a possible race with a THP page
+        * split, such that a whole THP page was munlocked, but nr_pages == 1.
+        * Returning a smaller mask due to that is OK, the worst that can
+        * happen is subsequent useless scanning of the former tail pages.
+        * The NR_MLOCK accounting can however become broken.
+        */
+       return nr_pages - 1;
 }
 
 /**
 
        while (start < end) {
                struct page *page = NULL;
-               unsigned int page_mask, page_increm;
+               unsigned int page_mask;
+               unsigned long page_increm;
                struct pagevec pvec;
                struct zone *zone;
                int zoneid;
                                goto next;
                        }
                }
-               page_increm = 1 + (~(start >> PAGE_SHIFT) & page_mask);
+               /* It's a bug to munlock in the middle of a THP page */
+               VM_BUG_ON((start >> PAGE_SHIFT) & page_mask);
+               page_increm = 1 + page_mask;
                start += page_increm * PAGE_SIZE;
 next:
                cond_resched();