diff --git a/include/linux/swap.h b/include/linux/swap.h
index 1d38d9475c4d0861d7a4597614d0ea225613c568..65a37e5551249f57b8e86fd73b46adeeb0953396 100644
--- a/include/linux/swap.h
+++ b/include/linux/swap.h
@@ -427,7 +427,7 @@ extern int swap_writepage(struct page *page, struct writeback_control *wbc);
 extern void end_swap_bio_write(struct bio *bio);
 extern int __swap_writepage(struct page *page, struct writeback_control *wbc,
 	bio_end_io_t end_write_func);
-extern int swap_set_page_dirty(struct page *page);
+bool swap_dirty_folio(struct address_space *mapping, struct folio *folio);
 
 int add_swap_extent(struct swap_info_struct *sis, unsigned long start_page,
 		unsigned long nr_pages, sector_t start_block);
diff --git a/mm/page_io.c b/mm/page_io.c
index 24c975fb4e211c4c4f6d485950fea771ee44b521..8f20f4dad289b7a9c88cf456eb1085cd91f64b2b 100644
--- a/mm/page_io.c
+++ b/mm/page_io.c
@@ -438,19 +438,21 @@ int swap_readpage(struct page *page, bool synchronous)
 	return ret;
 }
 
-int swap_set_page_dirty(struct page *page)
+bool swap_dirty_folio(struct address_space *mapping, struct folio *folio)
 {
-	struct swap_info_struct *sis = page_swap_info(page);
+	struct swap_info_struct *sis = swp_swap_info(folio_swap_entry(folio));
 
 	if (data_race(sis->flags & SWP_FS_OPS)) {
-		struct address_space *mapping = sis->swap_file->f_mapping;
-		const struct address_space_operations *aops = mapping->a_ops;
+		const struct address_space_operations *aops;
+
+		mapping = sis->swap_file->f_mapping;
+		aops = mapping->a_ops;
 
-		VM_BUG_ON_PAGE(!PageSwapCache(page), page);
+		VM_BUG_ON_FOLIO(!folio_test_swapcache(folio), folio);
 		if (aops->dirty_folio)
-			return aops->dirty_folio(mapping, page_folio(page));
-		return aops->set_page_dirty(page);
+			return aops->dirty_folio(mapping, folio);
+		return aops->set_page_dirty(&folio->page);
 	} else {
-		return __set_page_dirty_no_writeback(page);
+		return __set_page_dirty_no_writeback(&folio->page);
 	}
 }
diff --git a/mm/swap_state.c b/mm/swap_state.c
index 8d410424210007f08ac95b41aa26a57fac6b698f..4772afd08101aaf35ed34e0319052165b1c33cda 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -30,7 +30,7 @@
  */
 static const struct address_space_operations swap_aops = {
 	.writepage	= swap_writepage,
-	.set_page_dirty	= swap_set_page_dirty,
+	.dirty_folio	= swap_dirty_folio,
 #ifdef CONFIG_MIGRATION
 	.migratepage	= migrate_page,
 #endif