Re: [PATCH] mm/hugetlb: avoid recursive i_mmap_rwsem in PMD sharing

["Li Zhe" <lizhe.67@bytedance.com>]

On 10/10/26 2:05 AM, David Hildenbrand (Arm) wrote:
> On 10/9/26 13:52, Lorenzo Stoakes (ARM) wrote:
>> On Mon, Oct 05, 2026 at 11:25:01AM +0200, David Hildenbrand (Arm) wrote:
>>>> David mentioned that Oscar was planning to send a proper fix for this
>>>> issue, so I will wait for that patch instead of moving forward with this
>>>> trylock fallback approach.
>>> @Lorenzo, if Oscar is too busy, I guess we can paste the overall idea for the
>>> fix here as well (publicly)?
>> Yeah I don't see why not! I don't think there's anything that needs to be
>> private here.
> Li, see below, can you work with the below?
>
>
> ----8<----
>  From 9e1c0676ea6c9cbcabf91cd8badf81a9d88fd3ef Mon Sep 17 00:00:00 2001
> From: "Lorenzo Stoakes (ARM)" <ljs@kernel.org>
> Date: Wed, 29 Jul 2026 18:33:08 +0100
> Subject: [PATCH] fix
>
> ---
>   mm/hugetlb.c | 103 +++++++++++++++++++++++++++++++++++++++------------
>   1 file changed, 79 insertions(+), 24 deletions(-)
>
> diff --git a/mm/hugetlb.c b/mm/hugetlb.c
> index e93c4d2456aa..88ff2c4ad483 100644
> --- a/mm/hugetlb.c
> +++ b/mm/hugetlb.c
> @@ -5077,6 +5077,79 @@ static void move_huge_pte(struct vm_area_struct *vma, unsigned long old_addr,
>   	spin_unlock(dst_ptl);
>   }
>
> +/* Returns last address processed */
> +static unsigned long prealloc_or_move_page_tables(struct vm_area_struct *vma,
> +		struct vm_area_struct *new_vma, unsigned long old_addr,
> +		unsigned long new_addr, unsigned long len, unsigned long sz,
> +		unsigned long last_addr_mask, struct mmu_gather *tlb,
> +		bool is_prealloc)
> +{
> +	unsigned long old_end = old_addr + len;
> +	struct hstate *h = hstate_vma(vma);
> +	struct mm_struct *mm = vma->vm_mm;
> +	pte_t *src_pte, *dst_pte;
> +
> +	hugetlb_vma_assert_locked(vma);
> +
> +	for (; old_addr < old_end; old_addr += sz, new_addr += sz) {
> +		src_pte = hugetlb_walk(vma, old_addr, sz);
> +		if (!src_pte) {
> +			old_addr |= last_addr_mask;
> +			new_addr |= last_addr_mask;
> +			continue;
> +		}
> +		if (huge_pte_none(huge_ptep_get(mm, old_addr, src_pte)))
> +			continue;
> +
> +		if (is_prealloc) {
> +			if (!huge_pte_alloc(mm, new_vma, new_addr, sz))
> +				break;
> +			continue;
> +		}
> +
> +		if (huge_pmd_unshare(tlb, vma, old_addr, src_pte)) {
> +			old_addr |= last_addr_mask;
> +			new_addr |= last_addr_mask;
> +			continue;
> +		}
> +
> +		dst_pte = hugetlb_walk(new_vma, new_addr, sz);
> +		if (!dst_pte)
> +			break;
> +
> +		move_huge_pte(vma, old_addr, new_addr, src_pte, dst_pte, sz);
> +		tlb_remove_huge_tlb_entry(h, tlb, src_pte, old_addr);
> +	}
> +
> +	return old_addr;
> +}
> +
> +/* Returns true if preallocation succeeded across range, otherwise false. */
> +static bool prealloc_hugetlb_page_tables(struct vm_area_struct *vma,
> +		struct vm_area_struct *new_vma, unsigned long old_addr,
> +		unsigned long new_addr, unsigned long len, unsigned long sz,
> +		unsigned long last_addr_mask)
> +{
> +	unsigned long old_end = old_addr + len;
> +	unsigned long addr_end;
> +
> +	addr_end = prealloc_or_move_page_tables(vma, new_vma, old_addr, new_addr,
> +			len, sz, last_addr_mask, NULL, /*is_prealloc=*/true);
> +
> +	return addr_end >= old_end;
> +}
> +
> +/* Returns last processed address. */
> +static unsigned long __move_hugetlb_page_tables(struct vm_area_struct *vma,
> +		struct vm_area_struct *new_vma, unsigned long old_addr,
> +		unsigned long new_addr, unsigned long len, unsigned long sz,
> +		unsigned long last_addr_mask, struct mmu_gather *tlb)
> +{
> +	i_mmap_assert_write_locked(vma->vm_file->f_mapping);
> +	return prealloc_or_move_page_tables(vma, new_vma, old_addr, new_addr,
> +			len, sz, last_addr_mask, tlb, /*is_prealloc=*/false);
> +}
> +
>   int move_hugetlb_page_tables(struct vm_area_struct *vma,
>   			     struct vm_area_struct *new_vma,
>   			     unsigned long old_addr, unsigned long new_addr,
> @@ -5088,9 +5161,9 @@ int move_hugetlb_page_tables(struct vm_area_struct *vma,
>   	struct mm_struct *mm = vma->vm_mm;
>   	unsigned long old_end = old_addr + len;
>   	unsigned long last_addr_mask;
> -	pte_t *src_pte, *dst_pte;
>   	struct mmu_notifier_range range;
>   	struct mmu_gather tlb;
> +	bool preallocated;
>
>   	mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm, old_addr,
>   				old_end);
> @@ -5106,30 +5179,12 @@ int move_hugetlb_page_tables(struct vm_area_struct *vma,
>   	last_addr_mask = hugetlb_mask_last_page(h);
>   	/* Prevent race with file truncation */
>   	hugetlb_vma_lock_write(vma);
> +	preallocated = prealloc_hugetlb_page_tables(vma, new_vma, old_addr,
> +			new_addr, len, sz, last_addr_mask);
>   	i_mmap_lock_write(mapping);
> -	for (; old_addr < old_end; old_addr += sz, new_addr += sz) {
> -		src_pte = hugetlb_walk(vma, old_addr, sz);
> -		if (!src_pte) {
> -			old_addr |= last_addr_mask;
> -			new_addr |= last_addr_mask;
> -			continue;
> -		}
> -		if (huge_pte_none(huge_ptep_get(mm, old_addr, src_pte)))
> -			continue;
> -
> -		if (huge_pmd_unshare(&tlb, vma, old_addr, src_pte)) {
> -			old_addr |= last_addr_mask;
> -			new_addr |= last_addr_mask;
> -			continue;
> -		}
> -
> -		dst_pte = huge_pte_alloc(mm, new_vma, new_addr, sz);
> -		if (!dst_pte)
> -			break;
> -
> -		move_huge_pte(vma, old_addr, new_addr, src_pte, dst_pte, sz);
> -		tlb_remove_huge_tlb_entry(h, &tlb, src_pte, old_addr);
> -	}
> +	if (preallocated)
> +		old_addr = __move_hugetlb_page_tables(vma, new_vma, old_addr,
> +			new_addr, len, sz, last_addr_mask, &tlb);
>
>   	tlb_flush_mmu_tlbonly(&tlb);
>   	huge_pmd_unshare_flush(&tlb, vma);
> --
> 2.55.0

Hi David,

I thought more about the case where the destination page table was
preallocated as a shared PMD.

I do not think the current move_huge_pte() handles that case very well. The
reason is that it assumes the destination PTE is not part of a shared PMD
page table: it clears the source PTE and then installs that PTE into the
destination. That works when the destination PMD page table is not shared,
because the source removal and destination install have a net-zero effect on
rmap and hugetlb accounting.

If the destination PTE belongs to a shared PMD, that assumption does not
hold. In that case, I think we should leave the shared destination page
table untouched and treat the operation as a source-side unmap. Otherwise
move_huge_pte() either writes through a shared page table or skips the
source-side unmap accounting. That can make the rmap, hugetlb count and
folio release handling inconsistent with the actual PTE update.

So I think the change should be roughly:


---
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
@@
  static void move_huge_pte(struct vm_area_struct *vma, unsigned long 
old_addr,
-              unsigned long new_addr, pte_t *src_pte,
-              pte_t *dst_pte, unsigned long sz)
+              unsigned long new_addr, pte_t *src_pte,
+              pte_t *dst_pte, struct mmu_gather *tlb,
+              unsigned long sz)
  {
      struct hstate *h = hstate_vma(vma);
+    bool dst_shared = hugetlb_pmd_shared(dst_pte);
      struct mm_struct *mm = vma->vm_mm;
+    struct folio *folio = NULL;
      pte_t pte;

      pte = huge_ptep_get_and_clear(mm, old_addr, src_pte, sz);
+    tlb_remove_huge_tlb_entry(h, tlb, src_pte, old_addr);
+
+    if (dst_shared) {
+        if (pte_present(pte)) {
+            folio = page_folio(pte_page(pte));
+            if (huge_pte_dirty(pte))
+                folio_mark_dirty(folio);
+            hugetlb_count_sub(pages_per_huge_page(h), mm);
+            hugetlb_remove_rmap(folio);
+        }
+        goto unlock;
+    }

      /* Existing non-shared destination path remains unchanged. */
      set_huge_pte_at(mm, new_addr, dst_pte, pte, sz);

+unlock:
      ...
+
+    if (folio)
+        tlb_remove_page_size(tlb, folio_page(folio, 0),
+                     folio_size(folio));
  }
---


Does this sound like the right direction to you? If so, I will fold this
into the next version.

Thanks,
Zhe