Re: [PATCH v2 1/7] mm: factor out generic PTE batch detection from swap_pte_batch()

On Fri, 2026-08-14 at 13:30 +0530, Shivank Garg wrote:
> On Thu, 2026-08-13 at 11:57 +0200, David Hildenbrand (Arm) wrote:
> > On 8/13/26 06:23, Shivank Garg wrote:
> > > Factor out the generic PTE batch detection logic from swap_pte_batch()
> > > into softleaf_pte_batch() helper, so that it can be reused by upcoming
> > > patch that adds restore-side batching of migration entries.
> > > 
> > > swap_pte_batch() now keeps only its swap-specific concerns and
> > > delegates the batch detection to softleaf_pte_batch().
> > > 
> > > No functional changes intended.
> > > 
> > > Signed-off-by: Shivank Garg <[email protected]>
> > > ---
> > >  mm/internal.h | 36 +++++++++++++++++++++---------------
> > >  1 file changed, 21 insertions(+), 15 deletions(-)
> > > 
> > > diff --git a/mm/internal.h b/mm/internal.h
> > > index a5d0488a54dc..1e1ab53b642e 100644
> > > --- a/mm/internal.h
> > > +++ b/mm/internal.h
> > > @@ -560,6 +560,26 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
> > >  	}
> > >  }
> > >  
> > > +static inline int softleaf_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
> > > +{
> > > +	pte_t expected_pte = pte_next_softleaf_offset(pte);
> > > +	const pte_t *end_ptep = start_ptep + max_nr;
> > > +	pte_t *ptep = start_ptep + 1;
> > > +
> > 
> > Wouldn't we want some way to make sure that the softleaf entries actually carry
> > offsets or pfns of sorts?
> > 
> > > +	VM_WARN_ON(max_nr < 1);
> > > +
> > > +	while (ptep < end_ptep) {
> > > +		pte = ptep_get(ptep);
> > > +
> > > +		if (!pte_same(pte, expected_pte))
> > > +			break;
> > > +		expected_pte = pte_next_softleaf_offset(expected_pte);
> > > +		ptep++;
> > > +	}
> > > +
> > > +	return ptep - start_ptep;
> > > +}
> > > +
> > >  /**
> > >   * swap_pte_batch - detect a PTE batch for a set of contiguous swap entries
> > >   * @start_ptep: Page table pointer for the first entry.
> > > @@ -577,23 +597,9 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
> > >   */
> > >  static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
> > >  {
> > > -	pte_t expected_pte = pte_next_softleaf_offset(pte);
> > > -	const pte_t *end_ptep = start_ptep + max_nr;
> > > -	pte_t *ptep = start_ptep + 1;
> > > -
> > > -	VM_WARN_ON(max_nr < 1);
> > >  	VM_WARN_ON(!softleaf_is_swap(softleaf_from_pte(pte)));
> > 
> > Any reason we cannot simply rename this thing and drop the VM_WARN_ON? I mean, a
> > caller can just do that.
> > 
> 
> Thanks David,
> 
> Agreed on both suggestions.
> 
> I'll add VM_WARN_ON(!softleaf_is_swap(entry) && !softleaf_has_pfn(entry))
> in softleaf_pte_batch() for making sure it has offset/pfn and remove
> swap_pte_batch().
> 
> Best regards,
> Shivank

diff --git a/mm/internal.h b/mm/internal.h
index a5d0488a54dc..31c7c6a04eac 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -561,29 +561,31 @@ static inline void set_softleaf_ptes(struct mm_struct *mm, unsigned long addr,
 }
 
 /**
- * swap_pte_batch - detect a PTE batch for a set of contiguous swap entries
+ * softleaf_pte_batch - detect a PTE batch for a set of contiguous softleaf entries
  * @start_ptep: Page table pointer for the first entry.
  * @max_nr: The maximum number of table entries to consider.
  * @pte: Page table entry for the first entry.
  *
- * Detect a batch of contiguous swap entries: consecutive (non-present) PTEs
- * containing swap entries all with consecutive offsets and targeting the same
- * swap type, all with matching swp pte bits.
+ * Detect consecutive (non-present) PTEs containing softleaf entries with
+ * consecutive offsets and matching PTE bits. @pte must contain a swap entry or
+ * a softleaf entry carrying PFN.
  *
  * max_nr must be at least one and must be limited by the caller so scanning
  * cannot exceed a single page table.
  *
  * Return: the number of table entries in the batch.
  */
-static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
+static inline int softleaf_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {
-	pte_t expected_pte = pte_next_softleaf_offset(pte);
+	const softleaf_t entry = softleaf_from_pte(pte);
+	pte_t expected_pte;
 	const pte_t *end_ptep = start_ptep + max_nr;
 	pte_t *ptep = start_ptep + 1;
 
 	VM_WARN_ON(max_nr < 1);
-	VM_WARN_ON(!softleaf_is_swap(softleaf_from_pte(pte)));
+	VM_WARN_ON(!softleaf_is_swap(entry) && !softleaf_has_pfn(entry));
 
+	expected_pte = pte_next_softleaf_offset(pte);
 	while (ptep < end_ptep) {
 		pte = ptep_get(ptep);
 
diff --git a/mm/madvise.c b/mm/madvise.c
index c179938097bf..8bed1c58c4bb 100644
--- a/mm/madvise.c
+++ b/mm/madvise.c
@@ -692,7 +692,7 @@ static int madvise_free_pte_range(pmd_t *pmd, unsigned long addr,
 
 			if (softleaf_is_swap(entry)) {
 				max_nr = (end - addr) / PAGE_SIZE;
-				nr = swap_pte_batch(pte, max_nr, ptent);
+				nr = softleaf_pte_batch(pte, max_nr, ptent);
 				nr_swap -= nr;
 				swap_put_entries_direct(entry, nr);
 				clear_nonpresent_ptes(mm, addr, pte, nr);
diff --git a/mm/memory.c b/mm/memory.c
index a3d450d225ad..687796e1af84 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -1889,7 +1889,7 @@ static inline int zap_nonpresent_ptes(struct mmu_gather *tlb,
 		if (!should_zap_cows(details))
 			return 1;
 
-		nr = swap_pte_batch(pte, max_nr, ptent);
+		nr = softleaf_pte_batch(pte, max_nr, ptent);
 		rss[MM_SWAPENTS] -= nr;
 		swap_put_entries_direct(entry, nr);
 	} else if (softleaf_is_migration(entry)) {
@@ -4758,7 +4758,7 @@ static bool can_swapin_thp(struct vm_fault *vmf, pte_t *ptep, int nr_pages)
 	 * from different backends. And they are likely corner cases. Similar
 	 * things might be added once zswap support large folios.
 	 */
-	if (swap_pte_batch(ptep, nr_pages, pte) != nr_pages)
+	if (softleaf_pte_batch(ptep, nr_pages, pte) != nr_pages)
 		return false;
 	return true;
 }
@@ -5062,7 +5062,7 @@ vm_fault_t do_swap_page(struct vm_fault *vmf)
 		folio_ptep = vmf->pte - idx;
 		folio_pte = ptep_get(folio_ptep);
 		if (!pte_same(folio_pte, pte_move_softleaf_offset(vmf->orig_pte, -idx)) ||
-		    swap_pte_batch(folio_ptep, nr, folio_pte) != nr)
+		    softleaf_pte_batch(folio_ptep, nr, folio_pte) != nr)
 			goto check_folio;
 
 		page_idx = idx;