[PATCH v2 14/14] arm64/mm: Add ptep_get_and_clear_full() to optimize process teardown

[Barry Song]

> +pte_t contpte_ptep_get_and_clear_full(struct mm_struct *mm,
> +					unsigned long addr, pte_t *ptep)
> +{
> +	/*
> +	 * When doing a full address space teardown, we can avoid unfolding the
> +	 * contiguous range, and therefore avoid the associated tlbi. Instead,
> +	 * just get and clear the pte. The caller is promising to call us for
> +	 * every pte, so every pte in the range will be cleared by the time the
> +	 * tlbi is issued.
> +	 *
> +	 * This approach is not perfect though, as for the duration between
> +	 * returning from the first call to ptep_get_and_clear_full() and making
> +	 * the final call, the contpte block in an intermediate state, where
> +	 * some ptes are cleared and others are still set with the PTE_CONT bit.
> +	 * If any other APIs are called for the ptes in the contpte block during
> +	 * that time, we have to be very careful. The core code currently
> +	 * interleaves calls to ptep_get_and_clear_full() with ptep_get() and so
> +	 * ptep_get() must be careful to ignore the cleared entries when
> +	 * accumulating the access and dirty bits - the same goes for
> +	 * ptep_get_lockless(). The only other calls we might resonably expect
> +	 * are to set markers in the previously cleared ptes. (We shouldn't see
> +	 * valid entries being set until after the tlbi, at which point we are
> +	 * no longer in the intermediate state). Since markers are not valid,
> +	 * this is safe; set_ptes() will see the old, invalid entry and will not
> +	 * attempt to unfold. And the new pte is also invalid so it won't
> +	 * attempt to fold. We shouldn't see this for the 'full' case anyway.
> +	 *
> +	 * The last remaining issue is returning the access/dirty bits. That
> +	 * info could be present in any of the ptes in the contpte block.
> +	 * ptep_get() will gather those bits from across the contpte block. We
> +	 * don't bother doing that here, because we know that the information is
> +	 * used by the core-mm to mark the underlying folio as accessed/dirty.
> +	 * And since the same folio must be underpinning the whole block (that
> +	 * was a requirement for folding in the first place), that information
> +	 * will make it to the folio eventually once all the ptes have been
> +	 * cleared. This approach means we don't have to play games with
> +	 * accumulating and storing the bits. It does mean that any interleaved
> +	 * calls to ptep_get() may lack correct access/dirty information if we
> +	 * have already cleared the pte that happened to store it. The core code
> +	 * does not rely on this though.

even without any other threads running and touching those PTEs, this won't survive
on some hardware. we expose inconsistent CONTPTEs to hardware, this might result
in crashed firmware even in trustzone, strange&unknown faults to trustzone we have
seen on Qualcomm, but for MTK, it seems fine. when you do tlbi on a part of PTEs
with dropped CONT but still some other PTEs have CONT, we make hardware totally
confused.

zap_pte_range() has a force_flush when tlbbatch is full:

                        if (unlikely(__tlb_remove_page(tlb, page, delay_rmap))) {
                                force_flush = 1; 
                                addr += PAGE_SIZE;
                                break;
                        }

this means you can expose partial tlbi/flush directly to hardware while some
other PTEs are still CONT.

on the other hand, contpte_ptep_get_and_clear_full() doesn't need to depend
on fullmm, as long as zap range covers a large folio, we can flush tlbi for
those CONTPTEs all together in your contpte_ptep_get_and_clear_full() rather
than clearing one PTE.

Our approach in [1] is we do a flush for all CONTPTEs and go directly to the end
of the large folio:

#ifdef CONFIG_CONT_PTE_HUGEPAGE
			if (pte_cont(ptent)) {
				unsigned long next = pte_cont_addr_end(addr, end);

				if (next - addr != HPAGE_CONT_PTE_SIZE) {
					__split_huge_cont_pte(vma, pte, addr, false, NULL, ptl);
					/*
					 * After splitting cont-pte
					 * we need to process pte again.
					 */
					goto again_pte;
				} else {
					cont_pte_huge_ptep_get_and_clear(mm, addr, pte);

					tlb_remove_cont_pte_tlb_entry(tlb, pte, addr);
					if (unlikely(!page))
						continue;

					if (is_huge_zero_page(page)) {
						tlb_remove_page_size(tlb, page, HPAGE_CONT_PTE_SIZE);
						goto cont_next;
					}

					rss[mm_counter(page)] -= HPAGE_CONT_PTE_NR;
					page_remove_rmap(page, true);
					if (unlikely(page_mapcount(page) < 0))
						print_bad_pte(vma, addr, ptent, page);

					tlb_remove_page_size(tlb, page, HPAGE_CONT_PTE_SIZE);
				}
cont_next:
				/* "do while()" will do "pte++" and "addr + PAGE_SIZE" */
				pte += (next - PAGE_SIZE - (addr & PAGE_MASK))/PAGE_SIZE;
				addr = next - PAGE_SIZE;
				continue;
			}
#endif

this is our "full" counterpart, which clear_flush CONT_PTES pages directly, and
it never requires tlb->fullmm at all.

static inline pte_t __cont_pte_huge_ptep_get_and_clear_flush(struct mm_struct *mm,
				       unsigned long addr,
				       pte_t *ptep,
				       bool flush)
{
	pte_t orig_pte = ptep_get(ptep);

	CHP_BUG_ON(!pte_cont(orig_pte));
	CHP_BUG_ON(!IS_ALIGNED(addr, HPAGE_CONT_PTE_SIZE));
	CHP_BUG_ON(!IS_ALIGNED(pte_pfn(orig_pte), HPAGE_CONT_PTE_NR));

	return get_clear_flush(mm, addr, ptep, PAGE_SIZE, CONT_PTES, flush);
}

[1] https://github.com/OnePlusOSS/android_kernel_oneplus_sm8550/blob/oneplus/sm8550_u_14.0.0_oneplus11/mm/memory.c#L1539

> +	 */
> +
> +	return __ptep_get_and_clear(mm, addr, ptep);
> +}
> +EXPORT_SYMBOL(contpte_ptep_get_and_clear_full);
> +

Thanks
Barry