[PATCH v2 5/5] arm64: mm: Override arch_wants_pte_order()

[Ryan Roberts <ryan.roberts@arm.com>]

Define an arch-specific override of arch_wants_pte_order() so that when
FLEXIBLE_THP is enabled, large folios will be allocated for anonymous
memory with an order that is compatible with arm64's contpte mappings.

arch_wants_pte_order() returns an order according to the following
policy: For the unhinted case, when THP is not requested for the vma,
don't allow anything bigger than 64K. This means we don't waste too much
memory. Additionally, for 4K pages this is the contpte size, and for
16K, this is (usually) the HPA size when the uarch feature is
implemented. For the hinted case, when THP is requested for the vma,
allow the contpte size for all page size configurations; 64K for 4K, 2M
for 16K and 2M for 64K.

Additionally, the THP and NOTHP order constants are defined using
Kconfig so it is possible to override them at build time.

Signed-off-by: Ryan Roberts <ryan.roberts@arm.com>
---
 arch/arm64/Kconfig               | 12 ++++++++++++
 arch/arm64/include/asm/pgtable.h |  4 ++++
 arch/arm64/mm/mmu.c              |  8 ++++++++
 3 files changed, 24 insertions(+)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 343e1e1cae10..689c5bf13dc1 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -281,6 +281,18 @@ config ARM64_CONT_PMD_SHIFT
 	default 5 if ARM64_16K_PAGES
 	default 4
 
+config ARM64_PTE_ORDER_NOTHP
+	int
+	default 0 if ARM64_64K_PAGES	# 64K (1 page)
+	default 2 if ARM64_16K_PAGES	# 64K (4 pages; benefits from HPA where HW supports it)
+	default 4 if ARM64_4K_PAGES	# 64K (16 pages; eligible for contpte-mapping)
+
+config ARM64_PTE_ORDER_THP
+	int
+	default 5 if ARM64_64K_PAGES	# 2M  (32 pages; eligible for contpte-mapping)
+	default 7 if ARM64_16K_PAGES	# 2M  (128 pages; eligible for contpte-mapping)
+	default 4 if ARM64_4K_PAGES	# 64K (16 pages; eligible for contpte-mapping)
+
 config ARCH_MMAP_RND_BITS_MIN
 	default 14 if ARM64_64K_PAGES
 	default 16 if ARM64_16K_PAGES
diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index 6fd012663a01..8463d5f9f307 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -1117,6 +1117,10 @@ extern pte_t ptep_modify_prot_start(struct vm_area_struct *vma,
 extern void ptep_modify_prot_commit(struct vm_area_struct *vma,
 				    unsigned long addr, pte_t *ptep,
 				    pte_t old_pte, pte_t new_pte);
+
+#define arch_wants_pte_order arch_wants_pte_order
+extern int arch_wants_pte_order(struct vm_area_struct *vma);
+
 #endif /* !__ASSEMBLY__ */
 
 #endif /* __ASM_PGTABLE_H */
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index af6bc8403ee4..8556c4a9b507 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -1481,3 +1481,11 @@ void ptep_modify_prot_commit(struct vm_area_struct *vma, unsigned long addr, pte
 {
 	set_pte_at(vma->vm_mm, addr, ptep, pte);
 }
+
+int arch_wants_pte_order(struct vm_area_struct *vma)
+{
+	if (hugepage_vma_check(vma, vma->vm_flags, false, true, true))
+		return CONFIG_ARM64_PTE_ORDER_THP;
+	else
+		return CONFIG_ARM64_PTE_ORDER_NOTHP;
+}
-- 
2.25.1