Re: [PATCH v2 7/7] mm/rmap: batch the unmap of large folios in try_to_migrate_one()

[Lance Yang <lance.yang@linux.dev>]

On Wed, Aug 19, 2026 at 05:42:17PM +0800, Huang, Ying wrote:
>Lance Yang <lance.yang@linux.dev> writes:
>
>> On Tue, Aug 18, 2026 at 04:55:06PM +0800, Miaohe Lin wrote:
>>>On 2026/8/17 17:14, Lance Yang wrote:
>>>> +Cc Miaohe
>>>> 
>>>> On Thu, Aug 13, 2026 at 04:23:18AM +0000, Shivank Garg wrote:
>>>>> try_to_migrate_one() converts present PTEs to migration entries one at a
>>>>> time. For a PTE-mapped large folio, this repeat calls to ptep clear+flush,
>>>>> the migration entry build and set, folio_remove_rmap_pte() and folio_put(),
>>>>> each re-entering page_vma_mapped_walk() once per base page (256 times for
>>>>> 1M folio).
>>>>>
>>>>> Mirror try_to_unmap_one() to introduce folio_migrate_pte_batch() to detect
>>>>> eligible batch for PTEs mapping conseuctive subpages of a large folios,
>>>>> and convert the whole batch in one shot using the batched helpers.
>>>>>
>>>>> A side-effect of this change is trace_set_migration_pte() will record
>>>>> one event per batched run instead of earlier behavior of one per base page.
>>>>>
>>>>> Signed-off-by: Shivank Garg <shivankg@amd.com>
>>>>> ---
>>>>> mm/rmap.c | 115 ++++++++++++++++++++++++++++++++++++++++++++++----------------
>>>>> 1 file changed, 86 insertions(+), 29 deletions(-)
>>>>>
>>>>> diff --git a/mm/rmap.c b/mm/rmap.c
>>>>> index 35752a70f3a0..63b885c0b7ef 100644
>>>>> --- a/mm/rmap.c
>>>>> +++ b/mm/rmap.c
>>>>> @@ -2675,6 +2675,44 @@ static bool try_to_migrate_hugetlb_one(struct folio *folio,
>>>>> 	return ret;
>>>>> }
>>>>>
>>>>> +static inline unsigned int folio_migrate_pte_batch(struct folio *folio,
>>>>> +		struct page_vma_mapped_walk *pvmw, pte_t pte,
>>>>> +		struct page *subpage, bool anon_exclusive)
>>>>> +{
>>>>> +	unsigned long end_addr, addr = pvmw->address;
>>>>> +	struct vm_area_struct *vma = pvmw->vma;
>>>>> +	unsigned int max_nr, nr;
>>>>> +
>>>>> +#ifdef __HAVE_ARCH_UNMAP_ONE
>>>>> +	/* Cannot batch unmap if arch_unmap_one() is defined. */
>>>>> +	return 1;
>>>>> +#endif
>>>>> +
>>>>> +	if (!folio_test_large(folio))
>>>>> +		return 1;
>>>>> +	if (folio_is_zone_device(folio) || folio_test_has_hwpoisoned(folio))
>>>>> +		return 1;
>>>>> +	if (pte_unused(pte))
>>>>> +		return 1;
>>>>> +
>>>>> +	/* We may only batch within a single VMA and a single page table. */
>>>>> +	end_addr = pmd_addr_end(addr, vma->vm_end);
>>>>> +	max_nr = (end_addr - addr) >> PAGE_SHIFT;
>>>> 
>>>> Hmm ... can this still batch over a poisoned tail page?
>>>> 
>>>> memory_failure() sets PageHWPoison() before taking folio lock, but
>>>> cannot set PG_has_hwpoisoned until it acquires and releases that lock.
>>>> 
>>>> So tail page can already be poisoned while folio_test_has_hwpoisoned()
>>>> still returns false ... no?
>>>
>>>When memory error hits thp pages, memory_failure() first set PG_has_hwpoisoned and
>>>then tries to split thp pages. And try_to_migrate() will be called to set migration
>>>entries for anon pages. Does folio_migrate_pte_batch() work on this case? If so, the
>>>folio_test_has_hwpoisoned() check above could catch the bad pages?
>>>
>>>Or do you worry about the scene that meory error hits a thp while it's under migration?
>>
>> Yeah, latter case is exactly what I meant.
>>
>> try_to_migrate() is called with folio lock held. If migration already
>> owns the lock, memory_failure() can set PageHWPoison() on a tail page and
>> then block in folio_lock(), before reaching folio_set_has_hwpoisoned().
>> try_to_migrate_one() may meanwhile start from a healthy subpage, see
>> PageHWPoison(subpage) clear and PG_has_hwpoisoned still clear, then batch
>> across poisoned tail and install a normal migration entry for it.
>>
>> Once memory_failure() publishes PG_has_hwpoisoned, the folio-level check
>> does stop batching, as you said. It's just this window before publication
>> that worries me ...
>
>I agree with you that we should check PageHWPoison() for each subpage
>before batching.
>
>However, IIUC, your suggestion can only reduce the race window instead
>of avoiding it completely. PageHWPoison() may be set even after checking
>it anyway.

Yeah, agreed. Still racy if PageHWPoison() gets set right after the
check.

IIUC, the real window is earlier than folio_lock(): between
TestSetPageHWPoison() and get_hwpoison_page(). No extra ref pins the old
folio there, so migration can finish and move the mappings to the new
folio. memory_failure() then only handles the old page, leaving the
migrated mapping without HWPoison handling.

Thing is, batching makes that window a whole lot wider ... old loop at
least checks each subpage, while one healthy first page now gives the
whole batch a green light. So it can miss a tail page already poisoned
before batching even starts.

Per-subpage check won't make the race go away, but at least keeps
batching from making it worse.

No clever idea from me yet on how to close it completely ...

Cheers, Lance

>
>---
>Best Regards,
>Huang, Ying
>
>> Thanks, Lance
>>
>>>Thanks both.
>>>.
>>>
>>>> 
>>>> Starting from a healthy first subpage, folio_migrate_pte_batch() can
>>>> then batch across poisoned tail page. hwpoison only describes first
>>>> subpage, so set_softleaf_ptes() installs a normal migration entry for
>>>> poisoned page instead of an HWPoison entry ...
>>>> 
>>>> Should folio_migrate_pte_batch() check PageHWPoison() on every candidate
>>>> subpage and stop before a poisoned one?
>>>> 
>>>> Cheers, Lance
>>>> 
>>>> 
>>>>> +	/*
>>>>> +	 * If unmap fails, we need to restore the ptes. To avoid accidentally
>>>>> +	 * upgrading write permissions for ptes that were not originally writable,
>>>>> +	 * and to avoid losing the soft-dirty bit, use the appropriate FPB flags.
>>>>> +	 */
>>>>> +	nr = folio_pte_batch_flags(folio, vma, pvmw->pte, &pte, max_nr,
>>>>> +				   FPB_RESPECT_WRITE | FPB_RESPECT_SOFT_DIRTY);
>>>>> +
>>>>> +	/* Limit possible batch count to a uniform PageAnonExclusive value */
>>>>> +	if (folio_test_anon(folio))
>>>>> +		nr = page_anon_exclusive_batch(0, nr, subpage, anon_exclusive);
>>>>> +
>>>>> +	return nr;
>>>>> +}
>>>>> +
>>>>> /*
>>>>>  * @arg: enum ttu_flags will be passed to this argument.
>>>>>  *
>>>>> @@ -2686,12 +2724,12 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> {
>>>>> 	struct mm_struct *mm = vma->vm_mm;
>>>>> 	DEFINE_FOLIO_VMA_WALK(pvmw, folio, vma, address, 0);
>>>>> -	bool anon_exclusive, writable, ret = true;
>>>>> +	bool anon_exclusive, hwpoison, writable, ret = true;
>>>>> 	pte_t pteval;
>>>>> 	struct page *subpage;
>>>>> 	struct mmu_notifier_range range;
>>>>> 	enum ttu_flags flags = (enum ttu_flags)(long)arg;
>>>>> -	unsigned long pfn;
>>>>> +	unsigned long pfn, end_addr, nr_pages;
>>>>>
>>>>> 	/*
>>>>> 	 * When racing against e.g. zap_pte_range() on another cpu,
>>>>> @@ -2744,11 +2782,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> 			VM_BUG_ON_FOLIO(folio_test_hugetlb(folio) ||
>>>>> 					!folio_test_pmd_mappable(folio), folio);
>>>>>
>>>>> -			if (set_pmd_migration_entry(&pvmw, subpage)) {
>>>>> -				ret = false;
>>>>> -				page_vma_mapped_walk_done(&pvmw);
>>>>> -				break;
>>>>> -			}
>>>>> +			if (set_pmd_migration_entry(&pvmw, subpage))
>>>>> +				goto walk_abort;
>>>>> 			continue;
>>>>> #endif
>>>>> 		}
>>>>> @@ -2773,10 +2808,25 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> 		subpage = folio_page(folio, pfn - folio_pfn(folio));
>>>>> 		anon_exclusive = folio_test_anon(folio) &&
>>>>> 				 PageAnonExclusive(subpage);
>>>>> +		/*
>>>>> +		 * memory_failure() can set PageHWPoison concurrently without holding
>>>>> +		 * the folio lock. Snapshot the flag here to decide whether to batch
>>>>> +		 * PTEs or install hwpoison entry.
>>>>> +		 */
>>>>> +		hwpoison = PageHWPoison(subpage);
>>>>>
>>>>> +		nr_pages = 1;
>>>>> 		if (likely(pte_present(pteval))) {
>>>>> -			flush_cache_page(vma, address, pfn);
>>>>> -			/* Nuke the page table entry. */
>>>>> +			if (!hwpoison)
>>>>> +				nr_pages = folio_migrate_pte_batch(folio, &pvmw,
>>>>> +								   pteval, subpage,
>>>>> +								   anon_exclusive);
>>>>> +
>>>>> +			end_addr = address + nr_pages * PAGE_SIZE;
>>>>> +			flush_cache_range(vma, address, end_addr);
>>>>> +
>>>>> +			/* Nuke the page table entries. */
>>>>> +			pteval = get_and_clear_ptes(mm, address, pvmw.pte, nr_pages);
>>>>> 			if (should_defer_flush(mm, flags)) {
>>>>> 				/*
>>>>> 				 * We clear the PTE but do not flush so potentially
>>>>> @@ -2786,11 +2836,9 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> 				 * transition on a cached TLB entry is written through
>>>>> 				 * and traps if the PTE is unmapped.
>>>>> 				 */
>>>>> -				pteval = ptep_get_and_clear(mm, address, pvmw.pte);
>>>>> -
>>>>> -				set_tlb_ubc_flush_pending(mm, pteval, address, address + PAGE_SIZE);
>>>>> +				set_tlb_ubc_flush_pending(mm, pteval, address, end_addr);
>>>>> 			} else {
>>>>> -				pteval = ptep_clear_flush(vma, address, pvmw.pte);
>>>>> +				flush_tlb_range(vma, address, end_addr);
>>>>> 			}
>>>>> 			if (pte_dirty(pteval))
>>>>> 				folio_mark_dirty(folio);
>>>>> @@ -2809,7 +2857,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> 		/* Update high watermark before we lower rss */
>>>>> 		update_hiwater_rss(mm);
>>>>>
>>>>> -		if (PageHWPoison(subpage)) {
>>>>> +		if (hwpoison) {
>>>>> +			VM_WARN_ON_ONCE(nr_pages != 1);
>>>>> 			VM_WARN_ON_FOLIO(folio_is_device_private(folio), folio);
>>>>>
>>>>> 			pteval = swp_entry_to_pte(make_hwpoison_entry(subpage));
>>>>> @@ -2837,19 +2886,15 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> 			 * so we'll not check/care.
>>>>> 			 */
>>>>> 			if (arch_unmap_one(mm, vma, address, pteval) < 0) {
>>>>> -				set_pte_at(mm, address, pvmw.pte, pteval);
>>>>> -				ret = false;
>>>>> -				page_vma_mapped_walk_done(&pvmw);
>>>>> -				break;
>>>>> +				set_ptes(mm, address, pvmw.pte, pteval, nr_pages);
>>>>> +				goto walk_abort;
>>>>> 			}
>>>>>
>>>>> -			/* See folio_try_share_anon_rmap_pte(): clear PTE first. */
>>>>> +			/* See folio_try_share_anon_rmap_ptes(): clear PTE first. */
>>>>> 			if (anon_exclusive &&
>>>>> -			    folio_try_share_anon_rmap_pte(folio, subpage)) {
>>>>> -				set_pte_at(mm, address, pvmw.pte, pteval);
>>>>> -				ret = false;
>>>>> -				page_vma_mapped_walk_done(&pvmw);
>>>>> -				break;
>>>>> +			    folio_try_share_anon_rmap_ptes(folio, subpage, nr_pages)) {
>>>>> +				set_ptes(mm, address, pvmw.pte, pteval, nr_pages);
>>>>> +				goto walk_abort;
>>>>> 			}
>>>>>
>>>>> 			/*
>>>>> @@ -2859,19 +2904,31 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
>>>>> 			 */
>>>>> 			swp_pte = make_migration_pte(subpage, pteval,
>>>>> 						     writable, anon_exclusive);
>>>>> -			set_pte_at(mm, address, pvmw.pte, swp_pte);
>>>>> 			trace_set_migration_pte(address, pte_val(swp_pte),
>>>>> 						folio_order(folio));
>>>>> +
>>>>> +			/* Set nr_pages migration entries, advancing the PFN. */
>>>>> +			set_softleaf_ptes(mm, address, pvmw.pte, swp_pte, nr_pages);
>>>>> 			/*
>>>>> 			 * No need to invalidate here it will synchronize on
>>>>> 			 * against the special swap migration pte.
>>>>> 			 */
>>>>> 		}
>>>>>
>>>>> -		folio_remove_rmap_pte(folio, subpage, vma);
>>>>> -		if (vma->vm_flags & VM_LOCKED)
>>>>> -			mlock_drain_local();
>>>>> -		folio_put(folio);
>>>>> +		finish_folio_unmap(vma, folio, subpage, nr_pages);
>>>>> +
>>>>> +		/*
>>>>> +		 * If we batched the entire folio, there is nothing left to
>>>>> +		 * walk; stop right here.
>>>>> +		 */
>>>>> +		if (nr_pages == folio_nr_pages(folio))
>>>>> +			goto walk_done;
>>>>> +		continue;
>>>>> +walk_abort:
>>>>> +		ret = false;
>>>>> +walk_done:
>>>>> +		page_vma_mapped_walk_done(&pvmw);
>>>>> +		break;
>>>>> 	}
>>>>>
>>>>> 	mmu_notifier_invalidate_range_end(&range);
>>>>>
>>>>> -- 
>>>>> 2.43.0
>>>>>
>>>>>
>>>> .
>>>> 
>>>
>>>
>