}
 
 void __folio_lock(struct folio *folio);
-extern int __lock_page_killable(struct page *page);
+int __folio_lock_killable(struct folio *folio);
 extern int __lock_page_async(struct page *page, struct wait_page_queue *wait);
 extern int __lock_page_or_retry(struct page *page, struct mm_struct *mm,
                                unsigned int flags);
                __folio_lock(folio);
 }
 
+static inline int folio_lock_killable(struct folio *folio)
+{
+       might_sleep();
+       if (!folio_trylock(folio))
+               return __folio_lock_killable(folio);
+       return 0;
+}
+
 /*
  * lock_page_killable is like lock_page but can be interrupted by fatal
  * signals.  It returns 0 if it locked the page and -EINTR if it was
  */
 static inline int lock_page_killable(struct page *page)
 {
-       might_sleep();
-       if (!trylock_page(page))
-               return __lock_page_killable(page);
-       return 0;
+       return folio_lock_killable(page_folio(page));
 }
 
 /*
 
 }
 EXPORT_SYMBOL(__folio_lock);
 
-int __lock_page_killable(struct page *__page)
+int __folio_lock_killable(struct folio *folio)
 {
-       struct page *page = compound_head(__page);
-       wait_queue_head_t *q = page_waitqueue(page);
-       return wait_on_page_bit_common(q, page, PG_locked, TASK_KILLABLE,
+       wait_queue_head_t *q = page_waitqueue(&folio->page);
+       return wait_on_page_bit_common(q, &folio->page, PG_locked, TASK_KILLABLE,
                                        EXCLUSIVE);
 }
-EXPORT_SYMBOL_GPL(__lock_page_killable);
+EXPORT_SYMBOL_GPL(__folio_lock_killable);
 
 int __lock_page_async(struct page *page, struct wait_page_queue *wait)
 {
 int __lock_page_or_retry(struct page *page, struct mm_struct *mm,
                         unsigned int flags)
 {
+       struct folio *folio = page_folio(page);
+
        if (fault_flag_allow_retry_first(flags)) {
                /*
                 * CAUTION! In this case, mmap_lock is not released
        if (flags & FAULT_FLAG_KILLABLE) {
                int ret;
 
-               ret = __lock_page_killable(page);
+               ret = __folio_lock_killable(folio);
                if (ret) {
                        mmap_read_unlock(mm);
                        return 0;
                }
        } else {
-               __folio_lock(page_folio(page));
+               __folio_lock(folio);
        }
 
        return 1;
 
        *fpin = maybe_unlock_mmap_for_io(vmf, *fpin);
        if (vmf->flags & FAULT_FLAG_KILLABLE) {
-               if (__lock_page_killable(&folio->page)) {
+               if (__folio_lock_killable(folio)) {
                        /*
                         * We didn't have the right flags to drop the mmap_lock,
                         * but all fault_handlers only check for fatal signals