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

[Baolin Wang <baolin.wang@linux.alibaba.com>]

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@kernel.org>
> ---
>   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.

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

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