Re: [PATCH v4 06/16] x86/virt/tdx: Improve PAMT refcounts allocation for sparse memory

["Edgecombe, Rick P" <rick.p.edgecombe@intel.com>]

Kiryl, curious if you have any comments on the below...

On Wed, 2025-11-26 at 16:45 +0200, Nikolay Borisov wrote:
> > +static int pamt_refcount_populate(pte_t *pte, unsigned long addr, void
> > *data)
> > +{
> > +	struct page *page;
> > +	pte_t entry;
> > +
> > +	page = alloc_page(GFP_KERNEL | __GFP_ZERO);
> > +	if (!page)
> >    		return -ENOMEM;
> >    
> > +	entry = mk_pte(page, PAGE_KERNEL);
> > +
> > +	spin_lock(&init_mm.page_table_lock);
> > +	/*
> > +	 * PAMT refcount populations can overlap due to rounding of the
> > +	 * start/end pfn. Make sure the PAMT range is only populated once.
> > +	 */
> > +	if (pte_none(ptep_get(pte)))
> > +		set_pte_at(&init_mm, addr, pte, entry);
> > +	else
> > +		__free_page(page);
> > +	spin_unlock(&init_mm.page_table_lock);
> 
> nit: Wouldn't it be better to perform the pte_none() check before doing 
> the allocation thus avoiding needless allocations? I.e do the 
> alloc/mk_pte only after we are 100% sure we are going to use this entry.

Yes, but I'm also wondering why it needs init_mm.page_table_lock at all. Here is
my reasoning for why it doesn't:

apply_to_page_range() takes init_mm.page_table_lock internally when it modified
page tables in the address range (vmalloc). It needs to do this to avoid races
with other allocations that share the upper level page tables, which could be on
the ends of area that TDX reserves.

But pamt_refcount_populate() is only operating on the PTE's for the address
range that TDX code already controls. Vmalloc should not free the PMD underneath
the PTE operation because there is an allocation in any page tables it covers.
So we can skip the lock and also do the pte_none() check before the page
allocation as Nikolay suggests.

Same for the depopulate path.