diff --git a/mm/filemap.c b/mm/filemap.c
index 0ccdc9e92764..5e8e3fd26b8d 100644
--- a/mm/filemap.c
+++ b/mm/filemap.c
@@ -919,6 +919,33 @@ static int __filemap_lock_store(struct xa_state *xas, struct folio *folio,
 	return xas_error(xas);
 }
 
+int __filemap_add_swapcache(struct address_space *mapping, struct folio *folio,
+			    pgoff_t index, gfp_t gfp, void **shadowp)
+{
+	XA_STATE_ORDER(xas, &mapping->i_pages, index, folio_order(folio));
+	long nr;
+	int ret;
+
+	VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
+	VM_BUG_ON_FOLIO(!folio_test_swapcache(folio), folio);
+	mapping_set_update(&xas, mapping);
+
+	nr = folio_nr_pages(folio);
+	folio_ref_add(folio, nr);
+
+	ret = __filemap_lock_store(&xas, folio, index, gfp, shadowp);
+	if (likely(!ret)) {
+		mapping->nrpages += nr;
+		__node_stat_mod_folio(folio, NR_FILE_PAGES, nr);
+		__lruvec_stat_mod_folio(folio, NR_SWAPCACHE, nr);
+		xas_unlock_irq(&xas);
+	} else {
+		folio_put_refs(folio, nr);
+	}
+
+	return ret;
+}
+
 noinline int __filemap_add_folio(struct address_space *mapping,
 		struct folio *folio, pgoff_t index, gfp_t gfp, void **shadowp)
 {
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 9859aa4f7553..4fd2f74b94a9 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -2886,14 +2886,12 @@ static void __split_huge_page_tail(struct folio *folio, int tail,
 	lru_add_page_tail(head, page_tail, lruvec, list);
 }
 
-static void __split_huge_page(struct page *page, struct list_head *list,
-		pgoff_t end, unsigned int new_order)
+static void __split_huge_page(struct address_space *mapping, struct page *page,
+			      struct list_head *list, pgoff_t end, unsigned int new_order)
 {
 	struct folio *folio = page_folio(page);
 	struct page *head = &folio->page;
 	struct lruvec *lruvec;
-	struct address_space *swap_cache = NULL;
-	unsigned long offset = 0;
 	int i, nr_dropped = 0;
 	unsigned int new_nr = 1 << new_order;
 	int order = folio_order(folio);
@@ -2902,12 +2900,6 @@ static void __split_huge_page(struct page *page, struct list_head *list,
 	/* complete memcg works before add pages to LRU */
 	split_page_memcg(head, order, new_order);
 
-	if (folio_test_anon(folio) && folio_test_swapcache(folio)) {
-		offset = swp_offset(folio->swap);
-		swap_cache = swap_address_space(folio->swap);
-		xa_lock(&swap_cache->i_pages);
-	}
-
 	/* lock lru list/PageCompound, ref frozen by page_ref_freeze */
 	lruvec = folio_lruvec_lock(folio);
 
@@ -2919,18 +2911,18 @@ static void __split_huge_page(struct page *page, struct list_head *list,
 		if (head[i].index >= end) {
 			struct folio *tail = page_folio(head + i);
 
-			if (shmem_mapping(folio->mapping))
+			if (shmem_mapping(mapping))
 				nr_dropped++;
 			else if (folio_test_clear_dirty(tail))
 				folio_account_cleaned(tail,
-					inode_to_wb(folio->mapping->host));
+					inode_to_wb(mapping->host));
 			__filemap_remove_folio(tail, NULL);
 			folio_put(tail);
 		} else if (!PageAnon(page)) {
-			__xa_store(&folio->mapping->i_pages, head[i].index,
+			__xa_store(&mapping->i_pages, head[i].index,
 					head + i, 0);
-		} else if (swap_cache) {
-			__xa_store(&swap_cache->i_pages, offset + i,
+		} else if (folio_test_swapcache(folio)) {
+			__xa_store(&mapping->i_pages, swp_offset(folio->swap) + i,
 					head + i, 0);
 		}
 	}
@@ -2948,23 +2940,17 @@ static void __split_huge_page(struct page *page, struct list_head *list,
 	split_page_owner(head, order, new_order);
 
 	/* See comment in __split_huge_page_tail() */
-	if (folio_test_anon(folio)) {
+	if (mapping) {
 		/* Additional pin to swap cache */
-		if (folio_test_swapcache(folio)) {
-			folio_ref_add(folio, 1 + new_nr);
-			xa_unlock(&swap_cache->i_pages);
-		} else {
-			folio_ref_inc(folio);
-		}
-	} else {
-		/* Additional pin to page cache */
 		folio_ref_add(folio, 1 + new_nr);
-		xa_unlock(&folio->mapping->i_pages);
+		xa_unlock(&mapping->i_pages);
+	} else {
+		folio_ref_inc(folio);
 	}
 	local_irq_enable();
 
 	if (nr_dropped)
-		shmem_uncharge(folio->mapping->host, nr_dropped);
+		shmem_uncharge(mapping->host, nr_dropped);
 	remap_page(folio, nr);
 
 	if (folio_test_swapcache(folio))
@@ -3043,11 +3029,12 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 	struct deferred_split *ds_queue = get_deferred_split_queue(folio);
 	/* reset xarray order to new order after split */
 	XA_STATE_ORDER(xas, &folio->mapping->i_pages, folio->index, new_order);
+	struct address_space *mapping = folio_mapping(folio);;
 	struct anon_vma *anon_vma = NULL;
-	struct address_space *mapping = NULL;
 	int extra_pins, ret;
 	pgoff_t end;
 	bool is_hzp;
+	gfp_t gfp;
 
 	VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
 	VM_BUG_ON_FOLIO(!folio_test_large(folio), folio);
@@ -3079,7 +3066,6 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 		}
 	}
 
-
 	is_hzp = is_huge_zero_page(&folio->page);
 	if (is_hzp) {
 		pr_warn_ratelimited("Called split_huge_page for huge zero page\n");
@@ -3089,6 +3075,17 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 	if (folio_test_writeback(folio))
 		return -EBUSY;
 
+	if (mapping) {
+		gfp = current_gfp_context(mapping_gfp_mask(mapping) &
+					  GFP_RECLAIM_MASK);
+
+		xas_split_alloc(&xas, folio, folio_order(folio), gfp);
+		if (xas_error(&xas)) {
+			ret = xas_error(&xas);
+			goto out;
+		}
+	}
+
 	if (folio_test_anon(folio)) {
 		/*
 		 * The caller does not necessarily hold an mmap_lock that would
@@ -3104,33 +3101,19 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 			goto out;
 		}
 		end = -1;
-		mapping = NULL;
 		anon_vma_lock_write(anon_vma);
 	} else {
-		gfp_t gfp;
-
-		mapping = folio->mapping;
-
 		/* Truncated ? */
 		if (!mapping) {
 			ret = -EBUSY;
 			goto out;
 		}
 
-		gfp = current_gfp_context(mapping_gfp_mask(mapping) &
-							GFP_RECLAIM_MASK);
-
 		if (!filemap_release_folio(folio, gfp)) {
 			ret = -EBUSY;
 			goto out;
 		}
 
-		xas_split_alloc(&xas, folio, folio_order(folio), gfp);
-		if (xas_error(&xas)) {
-			ret = xas_error(&xas);
-			goto out;
-		}
-
 		anon_vma = NULL;
 		i_mmap_lock_read(mapping);
 
@@ -3189,7 +3172,9 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 			int nr = folio_nr_pages(folio);
 
 			xas_split(&xas, folio, folio_order(folio));
-			if (folio_test_pmd_mappable(folio) &&
+
+			if (!folio_test_anon(folio) &&
+			    folio_test_pmd_mappable(folio) &&
 			    new_order < HPAGE_PMD_ORDER) {
 				if (folio_test_swapbacked(folio)) {
 					__lruvec_stat_mod_folio(folio,
@@ -3202,7 +3187,7 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 			}
 		}
 
-		__split_huge_page(page, list, end, new_order);
+		__split_huge_page(mapping, page, list, end, new_order);
 		ret = 0;
 	} else {
 		spin_unlock(&ds_queue->split_queue_lock);
@@ -3218,9 +3203,9 @@ int split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 	if (anon_vma) {
 		anon_vma_unlock_write(anon_vma);
 		put_anon_vma(anon_vma);
-	}
-	if (mapping)
+	} else {
 		i_mmap_unlock_read(mapping);
+	}
 out:
 	xas_destroy(&xas);
 	count_vm_event(!ret ? THP_SPLIT_PAGE : THP_SPLIT_PAGE_FAILED);
diff --git a/mm/internal.h b/mm/internal.h
index 7e486f2c502c..b2bbfd3c2b50 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -1059,6 +1059,8 @@ struct migration_target_control {
  */
 size_t splice_folio_into_pipe(struct pipe_inode_info *pipe,
 			      struct folio *folio, loff_t fpos, size_t size);
+int __filemap_add_swapcache(struct address_space *mapping, struct folio *folio,
+			    pgoff_t index, gfp_t gfp, void **shadowp);
 
 /*
  * mm/vmalloc.c
diff --git a/mm/swap_state.c b/mm/swap_state.c
index b84e7b0ea4a5..caf69696f47c 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -90,48 +90,22 @@ int add_to_swap_cache(struct folio *folio, swp_entry_t entry,
 {
 	struct address_space *address_space = swap_address_space(entry);
 	pgoff_t idx = swp_offset(entry);
-	XA_STATE_ORDER(xas, &address_space->i_pages, idx, folio_order(folio));
-	unsigned long i, nr = folio_nr_pages(folio);
-	void *old;
-
-	xas_set_update(&xas, workingset_update_node);
+	int ret;
 
 	VM_BUG_ON_FOLIO(!folio_test_locked(folio), folio);
 	VM_BUG_ON_FOLIO(folio_test_swapcache(folio), folio);
 	VM_BUG_ON_FOLIO(!folio_test_swapbacked(folio), folio);
 
-	folio_ref_add(folio, nr);
 	folio_set_swapcache(folio);
 	folio->swap = entry;
 
-	do {
-		xas_lock_irq(&xas);
-		xas_create_range(&xas);
-		if (xas_error(&xas))
-			goto unlock;
-		for (i = 0; i < nr; i++) {
-			VM_BUG_ON_FOLIO(xas.xa_index != idx + i, folio);
-			if (shadowp) {
-				old = xas_load(&xas);
-				if (xa_is_value(old))
-					*shadowp = old;
-			}
-			xas_store(&xas, folio);
-			xas_next(&xas);
-		}
-		address_space->nrpages += nr;
-		__node_stat_mod_folio(folio, NR_FILE_PAGES, nr);
-		__lruvec_stat_mod_folio(folio, NR_SWAPCACHE, nr);
-unlock:
-		xas_unlock_irq(&xas);
-	} while (xas_nomem(&xas, gfp));
-
-	if (!xas_error(&xas))
-		return 0;
+	ret = __filemap_add_swapcache(address_space, folio, idx, gfp, shadowp);
+	if (ret) {
+		folio_clear_swapcache(folio);
+		folio->swap.val = 0;
+	}
 
-	folio_clear_swapcache(folio);
-	folio_ref_sub(folio, nr);
-	return xas_error(&xas);
+	return ret;
 }
 
 /*
@@ -142,7 +116,6 @@ void __delete_from_swap_cache(struct folio *folio,
 			swp_entry_t entry, void *shadow)
 {
 	struct address_space *address_space = swap_address_space(entry);
-	int i;
 	long nr = folio_nr_pages(folio);
 	pgoff_t idx = swp_offset(entry);
 	XA_STATE(xas, &address_space->i_pages, idx);
@@ -153,11 +126,9 @@ void __delete_from_swap_cache(struct folio *folio,
 	VM_BUG_ON_FOLIO(!folio_test_swapcache(folio), folio);
 	VM_BUG_ON_FOLIO(folio_test_writeback(folio), folio);
 
-	for (i = 0; i < nr; i++) {
-		void *entry = xas_store(&xas, shadow);
-		VM_BUG_ON_PAGE(entry != folio, entry);
-		xas_next(&xas);
-	}
+	xas_set_order(&xas, idx, folio_order(folio));
+	xas_store(&xas, shadow);
+
 	folio->swap.val = 0;
 	folio_clear_swapcache(folio);
 	address_space->nrpages -= nr;
@@ -252,6 +223,11 @@ void clear_shadow_from_swap_cache(swp_entry_t entry)
 
 	xas_set_update(&xas, workingset_update_node);
 
+	/*
+	 * On unmap, it may delete a larger order shadow here. It's mostly
+	 * fine since not entirely mapped folios are spiltted on swap out
+	 * and leaves shadows with order 0.
+	 */
 	xa_lock_irq(&address_space->i_pages);
 	if (xa_is_value(xas_load(&xas)))
 		xas_store(&xas, NULL);
