Re: [PATCH v2 8/9] mm/huge_memory: deduplicate zap_huge_pmd() further by tracking state

[Lorenzo Stoakes (Oracle)]

On Fri, Mar 20, 2026 at 11:49:18AM +0800, Baolin Wang wrote:
>
>
> On 3/19/26 9:00 PM, Lorenzo Stoakes (Oracle) wrote:
> > The flush_needed boolean is really tracking whether a PMD entry is present,
> > so use it that way directly and rename it to is_present.
> >
> > Deduplicate the folio_remove_rmap_pmd() and folio map count warning between
> > present and device private by tracking where we need to remove the rmap.
> >
> > We can also remove the comment about using flush_needed to track whether a
> > PMD entry is present as it's now irrelevant.
> >
> > Signed-off-by: Lorenzo Stoakes (Oracle) <ljs@xxxxxxxxxx>
> > ---
> >   mm/huge_memory.c | 28 +++++++++++++---------------
> >   1 file changed, 13 insertions(+), 15 deletions(-)
> >
> > diff --git a/mm/huge_memory.c b/mm/huge_memory.c
> > index c4e00c645e58..22715027e56c 100644
> > --- a/mm/huge_memory.c
> > +++ b/mm/huge_memory.c
> > @@ -2430,9 +2430,10 @@ static inline void zap_deposited_table(struct mm_struct *mm, pmd_t *pmd)
> >   bool zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
> >   		 pmd_t *pmd, unsigned long addr)
> >   {
> > +	bool needs_remove_rmap = false;
> >   	struct folio *folio = NULL;
> >   	bool needs_deposit = false;
> > -	bool flush_needed = false;
> > +	bool is_present = false;
> >   	spinlock_t *ptl;
> >   	pmd_t orig_pmd;
> > @@ -2449,6 +2450,7 @@ bool zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
> >   	 */
> >   	orig_pmd = pmdp_huge_get_and_clear_full(vma, addr, pmd,
> >   						tlb->fullmm);
> > +
> >   	arch_check_zapped_pmd(vma, orig_pmd);
> >   	tlb_remove_pmd_tlb_entry(tlb, pmd, addr);
> >   	if (vma_is_special_huge(vma))
> > @@ -2458,17 +2460,15 @@ bool zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
> >   		goto out;
> >   	}
> > -	if (pmd_present(orig_pmd)) {
> > +	is_present = pmd_present(orig_pmd);
> > +	if (is_present) {
> >   		folio = pmd_folio(orig_pmd);
> > -
> > -		flush_needed = true;
> > -		folio_remove_rmap_pmd(folio, &folio->page, vma);
> > -		WARN_ON_ONCE(folio_mapcount(folio) < 0);
> > +		needs_remove_rmap = true;
> >   	} else if (pmd_is_valid_softleaf(orig_pmd)) {
> >   		const softleaf_t entry = softleaf_from_pmd(orig_pmd);
> >   		folio = softleaf_to_folio(entry);
> > -
> > +		needs_remove_rmap = folio_is_device_private(folio);
> >   		if (!thp_migration_supported())
> >   			WARN_ONCE(1, "Non present huge pmd without pmd migration enabled!");
> >   	} else {
> > @@ -2483,27 +2483,25 @@ bool zap_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
> >   		add_mm_counter(tlb->mm, mm_counter_file(folio),
> >   			       -HPAGE_PMD_NR);
> > -		/*
> > -		 * Use flush_needed to indicate whether the PMD entry
> > -		 * is present, instead of checking pmd_present() again.
> > -		 */
> > -		if (flush_needed && pmd_young(orig_pmd) &&
> > +		if (is_present && pmd_young(orig_pmd) &&
> >   		    likely(vma_has_recency(vma)))
> >   			folio_mark_accessed(folio);
> >   	}
> > -	if (folio_is_device_private(folio)) {
> > +	if (needs_remove_rmap) {
> >   		folio_remove_rmap_pmd(folio, &folio->page, vma);
> >   		WARN_ON_ONCE(folio_mapcount(folio) < 0);
> > -		folio_put(folio);
> >   	}
> >   out:
> >   	if (arch_needs_pgtable_deposit() || needs_deposit)
> >   		zap_deposited_table(tlb->mm, pmd);
> > +	if (needs_remove_rmap && !is_present)
> > +		folio_put(folio);
> > +
>
> This kind of deduplication splits the device folio handling into three
> places, which is not easy to understand (at least for me), since the device
> folio has some special handling.

I think open-coded the exact same thing over and over again is FAR worse.

It's also actually 2 places for softleaf, because we were duplicating #2 below:

1. how do I get a folio?

2. do I need to remove this folio from the rmap? (yes for device private)

3. Do I need to put the folio (yes for device private)

You're maybe now just seeing exactly what happens here because the code is
clearer? Because before it was an unfathomable open coded mess with no
explanation.

Now you explicitly see what's happening :)

>
> Especially here, without looking closely at the if condition, it is unclear
> why we need to call folio_put(). Maybe some comments?

I can add one, the original didn't. This is an existing issue :)

>
> Anyway, I don't have a strong opinion. Let's wait for others' preferences.

Thanks, Lorenzo