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/20/26 9:51 PM, Lorenzo Stoakes (Oracle) wrote:
> 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@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.
> 
> 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 :)

Yes, I understand your point:) I saw that you changed the code to use 
the 'is_device_private' variable in the new version, which is more 
readable for me. Thanks.

>> 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 :)

Thanks.