return unlikely(swp_type(entry) == SWP_DEVICE_WRITE);
 }
 
+static inline unsigned long device_private_entry_to_pfn(swp_entry_t entry)
+{
+       return swp_offset(entry);
+}
+
 static inline struct page *device_private_entry_to_page(swp_entry_t entry)
 {
        return pfn_to_page(swp_offset(entry));
        return false;
 }
 
+static inline unsigned long device_private_entry_to_pfn(swp_entry_t entry)
+{
+       return 0;
+}
+
 static inline struct page *device_private_entry_to_page(swp_entry_t entry)
 {
        return NULL;
        return unlikely(swp_type(entry) == SWP_MIGRATION_WRITE);
 }
 
+static inline unsigned long migration_entry_to_pfn(swp_entry_t entry)
+{
+       return swp_offset(entry);
+}
+
 static inline struct page *migration_entry_to_page(swp_entry_t entry)
 {
        struct page *p = pfn_to_page(swp_offset(entry));
 {
        return 0;
 }
+
+static inline unsigned long migration_entry_to_pfn(swp_entry_t entry)
+{
+       return 0;
+}
+
 static inline struct page *migration_entry_to_page(swp_entry_t entry)
 {
        return NULL;
 
        return true;
 }
 
+/**
+ * check_pte - check if @pvmw->page is mapped at the @pvmw->pte
+ *
+ * page_vma_mapped_walk() found a place where @pvmw->page is *potentially*
+ * mapped. check_pte() has to validate this.
+ *
+ * @pvmw->pte may point to empty PTE, swap PTE or PTE pointing to arbitrary
+ * page.
+ *
+ * If PVMW_MIGRATION flag is set, returns true if @pvmw->pte contains migration
+ * entry that points to @pvmw->page or any subpage in case of THP.
+ *
+ * If PVMW_MIGRATION flag is not set, returns true if @pvmw->pte points to
+ * @pvmw->page or any subpage in case of THP.
+ *
+ * Otherwise, return false.
+ *
+ */
 static bool check_pte(struct page_vma_mapped_walk *pvmw)
 {
+       unsigned long pfn;
+
        if (pvmw->flags & PVMW_MIGRATION) {
-#ifdef CONFIG_MIGRATION
                swp_entry_t entry;
                if (!is_swap_pte(*pvmw->pte))
                        return false;
 
                if (!is_migration_entry(entry))
                        return false;
-               if (migration_entry_to_page(entry) - pvmw->page >=
-                               hpage_nr_pages(pvmw->page)) {
-                       return false;
-               }
-               if (migration_entry_to_page(entry) < pvmw->page)
-                       return false;
-#else
-               WARN_ON_ONCE(1);
-#endif
-       } else {
-               if (is_swap_pte(*pvmw->pte)) {
-                       swp_entry_t entry;
 
-                       entry = pte_to_swp_entry(*pvmw->pte);
-                       if (is_device_private_entry(entry) &&
-                           device_private_entry_to_page(entry) == pvmw->page)
-                               return true;
-               }
+               pfn = migration_entry_to_pfn(entry);
+       } else if (is_swap_pte(*pvmw->pte)) {
+               swp_entry_t entry;
 
-               if (!pte_present(*pvmw->pte))
+               /* Handle un-addressable ZONE_DEVICE memory */
+               entry = pte_to_swp_entry(*pvmw->pte);
+               if (!is_device_private_entry(entry))
                        return false;
 
-               /* THP can be referenced by any subpage */
-               if (pte_page(*pvmw->pte) - pvmw->page >=
-                               hpage_nr_pages(pvmw->page)) {
-                       return false;
-               }
-               if (pte_page(*pvmw->pte) < pvmw->page)
+               pfn = device_private_entry_to_pfn(entry);
+       } else {
+               if (!pte_present(*pvmw->pte))
                        return false;
+
+               pfn = pte_pfn(*pvmw->pte);
        }
 
+       if (pfn < page_to_pfn(pvmw->page))
+               return false;
+
+       /* THP can be referenced by any subpage */
+       if (pfn - page_to_pfn(pvmw->page) >= hpage_nr_pages(pvmw->page))
+               return false;
+
        return true;
 }