*
         * So update the flags atomically, and wake up the waiter
         * afterwards to avoid any races. This store-release pairs
-        * with the load-acquire in wait_on_page_bit_common().
+        * with the load-acquire in folio_wait_bit_common().
         */
        smp_store_release(&wait->flags, flags | WQ_FLAG_WOKEN);
        wake_up_state(wait->private, mode);
 }
 
 /*
- * A choice of three behaviors for wait_on_page_bit_common():
+ * A choice of three behaviors for folio_wait_bit_common():
  */
 enum behavior {
        EXCLUSIVE,      /* Hold ref to page and take the bit when woken, like
 };
 
 /*
- * Attempt to check (or get) the page bit, and mark us done
+ * Attempt to check (or get) the folio flag, and mark us done
  * if successful.
  */
-static inline bool trylock_page_bit_common(struct page *page, int bit_nr,
+static inline bool folio_trylock_flag(struct folio *folio, int bit_nr,
                                        struct wait_queue_entry *wait)
 {
        if (wait->flags & WQ_FLAG_EXCLUSIVE) {
-               if (test_and_set_bit(bit_nr, &page->flags))
+               if (test_and_set_bit(bit_nr, &folio->flags))
                        return false;
-       } else if (test_bit(bit_nr, &page->flags))
+       } else if (test_bit(bit_nr, &folio->flags))
                return false;
 
        wait->flags |= WQ_FLAG_WOKEN | WQ_FLAG_DONE;
 /* How many times do we accept lock stealing from under a waiter? */
 int sysctl_page_lock_unfairness = 5;
 
-static inline int wait_on_page_bit_common(wait_queue_head_t *q,
-       struct page *page, int bit_nr, int state, enum behavior behavior)
+static inline int folio_wait_bit_common(struct folio *folio, int bit_nr,
+               int state, enum behavior behavior)
 {
+       wait_queue_head_t *q = page_waitqueue(&folio->page);
        int unfairness = sysctl_page_lock_unfairness;
        struct wait_page_queue wait_page;
        wait_queue_entry_t *wait = &wait_page.wait;
        unsigned long pflags;
 
        if (bit_nr == PG_locked &&
-           !PageUptodate(page) && PageWorkingset(page)) {
-               if (!PageSwapBacked(page)) {
+           !folio_test_uptodate(folio) && folio_test_workingset(folio)) {
+               if (!folio_test_swapbacked(folio)) {
                        delayacct_thrashing_start();
                        delayacct = true;
                }
 
        init_wait(wait);
        wait->func = wake_page_function;
-       wait_page.page = page;
+       wait_page.page = &folio->page;
        wait_page.bit_nr = bit_nr;
 
 repeat:
         * Do one last check whether we can get the
         * page bit synchronously.
         *
-        * Do the SetPageWaiters() marking before that
+        * Do the folio_set_waiters() marking before that
         * to let any waker we _just_ missed know they
         * need to wake us up (otherwise they'll never
         * even go to the slow case that looks at the
         * lock to avoid races.
         */
        spin_lock_irq(&q->lock);
-       SetPageWaiters(page);
-       if (!trylock_page_bit_common(page, bit_nr, wait))
+       folio_set_waiters(folio);
+       if (!folio_trylock_flag(folio, bit_nr, wait))
                __add_wait_queue_entry_tail(q, wait);
        spin_unlock_irq(&q->lock);
 
         * see whether the page bit testing has already
         * been done by the wake function.
         *
-        * We can drop our reference to the page.
+        * We can drop our reference to the folio.
         */
        if (behavior == DROP)
-               put_page(page);
+               folio_put(folio);
 
        /*
         * Note that until the "finish_wait()", or until
                 *
                 * And if that fails, we'll have to retry this all.
                 */
-               if (unlikely(test_and_set_bit(bit_nr, &page->flags)))
+               if (unlikely(test_and_set_bit(bit_nr, folio_flags(folio, 0))))
                        goto repeat;
 
                wait->flags |= WQ_FLAG_DONE;
 
        /*
         * If a signal happened, this 'finish_wait()' may remove the last
-        * waiter from the wait-queues, but the PageWaiters bit will remain
+        * waiter from the wait-queues, but the folio waiters bit will remain
         * set. That's ok. The next wakeup will take care of it, and trying
         * to do it here would be difficult and prone to races.
         */
        return wait->flags & WQ_FLAG_WOKEN ? 0 : -EINTR;
 }
 
-void wait_on_page_bit(struct page *page, int bit_nr)
+void folio_wait_bit(struct folio *folio, int bit_nr)
 {
-       wait_queue_head_t *q = page_waitqueue(page);
-       wait_on_page_bit_common(q, page, bit_nr, TASK_UNINTERRUPTIBLE, SHARED);
+       folio_wait_bit_common(folio, bit_nr, TASK_UNINTERRUPTIBLE, SHARED);
 }
-EXPORT_SYMBOL(wait_on_page_bit);
+EXPORT_SYMBOL(folio_wait_bit);
 
-int wait_on_page_bit_killable(struct page *page, int bit_nr)
+int folio_wait_bit_killable(struct folio *folio, int bit_nr)
 {
-       wait_queue_head_t *q = page_waitqueue(page);
-       return wait_on_page_bit_common(q, page, bit_nr, TASK_KILLABLE, SHARED);
+       return folio_wait_bit_common(folio, bit_nr, TASK_KILLABLE, SHARED);
 }
-EXPORT_SYMBOL(wait_on_page_bit_killable);
+EXPORT_SYMBOL(folio_wait_bit_killable);
 
 /**
  * put_and_wait_on_page_locked - Drop a reference and wait for it to be unlocked
  */
 int put_and_wait_on_page_locked(struct page *page, int state)
 {
-       wait_queue_head_t *q;
-
-       page = compound_head(page);
-       q = page_waitqueue(page);
-       return wait_on_page_bit_common(q, page, PG_locked, state, DROP);
+       return folio_wait_bit_common(page_folio(page), PG_locked, state,
+                       DROP);
 }
 
 /**
  */
 void wait_on_page_private_2(struct page *page)
 {
-       page = compound_head(page);
-       while (PagePrivate2(page))
-               wait_on_page_bit(page, PG_private_2);
+       struct folio *folio = page_folio(page);
+
+       while (folio_test_private_2(folio))
+               folio_wait_bit(folio, PG_private_2);
 }
 EXPORT_SYMBOL(wait_on_page_private_2);
 
  */
 int wait_on_page_private_2_killable(struct page *page)
 {
+       struct folio *folio = page_folio(page);
        int ret = 0;
 
-       page = compound_head(page);
-       while (PagePrivate2(page)) {
-               ret = wait_on_page_bit_killable(page, PG_private_2);
+       while (folio_test_private_2(folio)) {
+               ret = folio_wait_bit_killable(folio, PG_private_2);
                if (ret < 0)
                        break;
        }
  */
 void __folio_lock(struct folio *folio)
 {
-       wait_queue_head_t *q = page_waitqueue(&folio->page);
-       wait_on_page_bit_common(q, &folio->page, PG_locked, TASK_UNINTERRUPTIBLE,
+       folio_wait_bit_common(folio, PG_locked, TASK_UNINTERRUPTIBLE,
                                EXCLUSIVE);
 }
 EXPORT_SYMBOL(__folio_lock);
 
 int __folio_lock_killable(struct folio *folio)
 {
-       wait_queue_head_t *q = page_waitqueue(&folio->page);
-       return wait_on_page_bit_common(q, &folio->page, PG_locked, TASK_KILLABLE,
+       return folio_wait_bit_common(folio, PG_locked, TASK_KILLABLE,
                                        EXCLUSIVE);
 }
 EXPORT_SYMBOL_GPL(__folio_lock_killable);