[RFC PATCH 24/26] arm64/hugetlb: add support for high-granularity mappings

[James Houghton <jthoughton@google.com>]

This is included in this RFC to demonstrate how an architecture that
doesn't use ARCH_WANT_GENERAL_HUGETLB can be updated to support HugeTLB
high-granularity mappings: an architecture just needs to implement
hugetlb_walk_to.

Signed-off-by: James Houghton <jthoughton@google.com>
---
 arch/arm64/Kconfig          |  1 +
 arch/arm64/mm/hugetlbpage.c | 63 +++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 1652a9800ebe..74108713a99a 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -99,6 +99,7 @@ config ARM64
 	select ARCH_WANT_FRAME_POINTERS
 	select ARCH_WANT_HUGE_PMD_SHARE if ARM64_4K_PAGES || (ARM64_16K_PAGES && !ARM64_VA_BITS_36)
 	select ARCH_WANT_HUGETLB_PAGE_OPTIMIZE_VMEMMAP
+	select ARCH_HAS_SPECIAL_HUGETLB_HGM
 	select ARCH_WANT_LD_ORPHAN_WARN
 	select ARCH_WANTS_NO_INSTR
 	select ARCH_HAS_UBSAN_SANITIZE_ALL
diff --git a/arch/arm64/mm/hugetlbpage.c b/arch/arm64/mm/hugetlbpage.c
index e2a5ec9fdc0d..1901818bed9d 100644
--- a/arch/arm64/mm/hugetlbpage.c
+++ b/arch/arm64/mm/hugetlbpage.c
@@ -281,6 +281,69 @@ void set_huge_swap_pte_at(struct mm_struct *mm, unsigned long addr,
 		set_pte(ptep, pte);
 }
 
+int hugetlb_walk_to(struct mm_struct *mm, struct hugetlb_pte *hpte,
+		    unsigned long addr, unsigned long sz, bool stop_at_none)
+{
+	pgd_t *pgdp;
+	p4d_t *p4dp;
+	pte_t *ptep;
+
+	if (!hpte->ptep) {
+		pgdp = pgd_offset(mm, addr);
+		p4dp = p4d_offset(pgdp, addr);
+		if (!p4dp)
+			return -ENOMEM;
+		hugetlb_pte_populate(hpte, (pte_t *)p4dp, P4D_SHIFT);
+	}
+
+	while (hugetlb_pte_size(hpte) > sz &&
+			!hugetlb_pte_present_leaf(hpte) &&
+			!(stop_at_none && hugetlb_pte_none(hpte))) {
+		if (hpte->shift == PMD_SHIFT) {
+			unsigned long rounded_addr = sz == CONT_PTE_SIZE
+						     ? addr & CONT_PTE_MASK
+						     : addr;
+
+			ptep = pte_offset_kernel((pmd_t *)hpte->ptep,
+						 rounded_addr);
+			if (!ptep)
+				return -ENOMEM;
+			if (sz == CONT_PTE_SIZE)
+				hpte->shift = CONT_PTE_SHIFT;
+			else
+				hpte->shift = pte_cont(*ptep) ? CONT_PTE_SHIFT
+							      : PAGE_SHIFT;
+			hpte->ptep = ptep;
+		} else if (hpte->shift == PUD_SHIFT) {
+			pud_t *pudp = (pud_t *)hpte->ptep;
+
+			ptep = (pte_t *)pmd_alloc(mm, pudp, addr);
+
+			if (!ptep)
+				return -ENOMEM;
+			if (sz == CONT_PMD_SIZE)
+				hpte->shift = CONT_PMD_SHIFT;
+			else
+				hpte->shift = pte_cont(*ptep) ? CONT_PMD_SHIFT
+							      : PMD_SHIFT;
+			hpte->ptep = ptep;
+		} else if (hpte->shift == P4D_SHIFT) {
+			ptep = (pte_t *)pud_alloc(mm, (p4d_t *)hpte->ptep, addr);
+			if (!ptep)
+				return -ENOMEM;
+			hpte->shift = PUD_SHIFT;
+			hpte->ptep = ptep;
+		} else
+			/*
+			 * This also catches the cases of CONT_PMD_SHIFT and
+			 * CONT_PTE_SHIFT. Those PTEs should always be leaves.
+			 */
+			BUG();
+	}
+
+	return 0;
+}
+
 pte_t *huge_pte_alloc(struct mm_struct *mm, struct vm_area_struct *vma,
 		      unsigned long addr, unsigned long sz)
 {
-- 
2.37.0.rc0.161.g10f37bed90-goog