#include <linux/random.h>
 #include <linux/sched/sysctl.h>
 #include <linux/memory-tiers.h>
+#include <linux/pagewalk.h>
 
 #include <asm/tlbflush.h>
 
        for (i = 0; i < nr_pages; i++) {
                unsigned long addr = (unsigned long)(*pages);
                struct vm_area_struct *vma;
-               struct page *page;
+               struct folio_walk fw;
+               struct folio *folio;
                int err = -EFAULT;
 
                vma = vma_lookup(mm, addr);
                if (!vma)
                        goto set_status;
 
-               /* FOLL_DUMP to ignore special (like zero) pages */
-               page = follow_page(vma, addr, FOLL_GET | FOLL_DUMP);
-
-               err = PTR_ERR(page);
-               if (IS_ERR(page))
-                       goto set_status;
-
-               err = -ENOENT;
-               if (!page)
-                       goto set_status;
-
-               if (!is_zone_device_page(page))
-                       err = page_to_nid(page);
-
-               put_page(page);
+               folio = folio_walk_start(&fw, vma, addr, FW_ZEROPAGE);
+               if (folio) {
+                       if (is_zero_folio(folio) || is_huge_zero_folio(folio))
+                               err = -EFAULT;
+                       else if (folio_is_zone_device(folio))
+                               err = -ENOENT;
+                       else
+                               err = folio_nid(folio);
+                       folio_walk_end(&fw, vma);
+               } else {
+                       err = -ENOENT;
+               }
 set_status:
                *status = err;