[RFC PATCH 15/15] arm64: mm: Set stack size from the kernel command line

[Mostafa Saleh <smostafa@google.com>]

Select CONFIG_ARCH_HAS_VARIABLE_STACK_SIZE for arm64 and override
arch_vm_stack_pages() to a value set from the command line, this
value have the following properties:
- Aligned to page size (doesn't need to be a power of 2), so it can
  be 12kB for example
- Min value is chosen to be 8kB
- Max value will be the default stack size on arm64 which is used to
  allocate the VA range for the kernel stack.

Signed-off-by: Mostafa Saleh <smostafa@google.com>
---
 .../admin-guide/kernel-parameters.txt         |  7 ++++++
 arch/arm64/Kconfig                            |  1 +
 arch/arm64/include/asm/memory.h               | 12 ++++++++++
 arch/arm64/kernel/entry.S                     | 10 +++++---
 arch/arm64/kernel/setup.c                     | 24 +++++++++++++++++++
 5 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/Documentation/admin-guide/kernel-parameters.txt b/Documentation/admin-guide/kernel-parameters.txt
index 33cd30996e47..e6ef22f01e01 100644
--- a/Documentation/admin-guide/kernel-parameters.txt
+++ b/Documentation/admin-guide/kernel-parameters.txt
@@ -2920,6 +2920,13 @@ Kernel parameters
 
 	keepinitrd	[HW,ARM] See retain_initrd.
 
+	kernel_stack_pages=
+			[ARM64,EARLY]
+			Format: <integer>
+			Number of pages mapped for each task kernel stack.
+			The virtual area stays THREAD_SIZE and only its top
+			pages are mapped.
+
 	kernelcore=	[KNL,X86,PPC,EARLY]
 			Format: nn[KMGTPE] | nn% | "mirror"
 			This parameter specifies the amount of memory usable by
diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 020fcaf36b1b..5e09ae904e34 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -58,6 +58,7 @@ config ARM64
 	select ARCH_HAS_SYSCALL_WRAPPER
 	select ARCH_HAS_TICK_BROADCAST if GENERIC_CLOCKEVENTS_BROADCAST
 	select ARCH_HAS_ZONE_DMA_SET if EXPERT
+	select ARCH_HAS_VARIABLE_STACK_SIZE
 	select ARCH_HAVE_ELF_PROT
 	select ARCH_HAVE_NMI_SAFE_CMPXCHG
 	select ARCH_HAVE_TRACE_MMIO_ACCESS
diff --git a/arch/arm64/include/asm/memory.h b/arch/arm64/include/asm/memory.h
index 4545d9b39bd5..1d300b16032d 100644
--- a/arch/arm64/include/asm/memory.h
+++ b/arch/arm64/include/asm/memory.h
@@ -113,6 +113,13 @@
 #define DIRECT_MAP_PHYSMEM_END	__pa(PAGE_END - 1)
 
 #define MIN_THREAD_SHIFT	(14 + KASAN_THREAD_SHIFT)
+#define MIN_KERNEL_STACK_SHIFT	(13 + KASAN_THREAD_SHIFT)
+
+#if (MIN_KERNEL_STACK_SHIFT < PAGE_SHIFT)
+#define MIN_KERNEL_STACK_SIZE	PAGE_SIZE
+#else
+#define MIN_KERNEL_STACK_SIZE	(UL(1) << MIN_KERNEL_STACK_SHIFT)
+#endif
 
 /*
  * VMAP'd stacks are allocated at page granularity, so we must ensure that such
@@ -130,6 +137,11 @@
 
 #define THREAD_SIZE		(UL(1) << THREAD_SHIFT)
 
+#ifndef __ASSEMBLY__
+extern unsigned long kernel_stack_pages;
+#define arch_vm_stack_pages()	(kernel_stack_pages)
+#endif
+
 #define THREAD_ALIGN		THREAD_SIZE
 
 #define IRQ_STACK_SIZE		THREAD_SIZE
diff --git a/arch/arm64/kernel/entry.S b/arch/arm64/kernel/entry.S
index ea733b673970..39e68aace0a6 100644
--- a/arch/arm64/kernel/entry.S
+++ b/arch/arm64/kernel/entry.S
@@ -76,9 +76,13 @@
 	msr	spsel, #0
 	mov	x1, sp
 	msr	spsel, #1
-	sub	x1, x1, x0
-	cmp	x1, #THREAD_SIZE
-	b.lo	.Lstack_ok\@
+
+	add	x0, x0, #THREAD_SIZE
+	sub	x1, x0, x1
+	adr_l	x0, kernel_stack_pages
+	ldr	x0, [x0]
+	cmp	x1, x0, lsl #PAGE_SHIFT
+	b.ls	.Lstack_ok\@
 
 	/* Check whether we are on the IRQ, SDEI or EFI stacks. */
 	stp	x2, x30, [sp, #-16]!
diff --git a/arch/arm64/kernel/setup.c b/arch/arm64/kernel/setup.c
index 29c6100f0c50..34559048dcca 100644
--- a/arch/arm64/kernel/setup.c
+++ b/arch/arm64/kernel/setup.c
@@ -278,6 +278,30 @@ u64 cpu_logical_map(unsigned int cpu)
 	return __cpu_logical_map[cpu];
 }
 
+unsigned long kernel_stack_pages __ro_after_init = THREAD_SIZE / PAGE_SIZE;
+
+static int __init setup_kernel_stack_pages(char *str)
+{
+	unsigned long min = MIN_KERNEL_STACK_SIZE / PAGE_SIZE;
+	unsigned long max = THREAD_SIZE / PAGE_SIZE;
+	unsigned long val;
+
+	if (!str)
+		return -EINVAL;
+
+	if (kstrtoul(str, 10, &val) || val < min || val > max) {
+		pr_err("Invalid kernel_stack_pages=%s, should be between %lu and %lu\n",
+		       str, min, max);
+		return 0;
+	}
+
+	kernel_stack_pages = val;
+	pr_info("Kernel stack size set to %lu KiB\n", (val * PAGE_SIZE) / SZ_1K);
+
+	return 0;
+}
+early_param("kernel_stack_pages", setup_kernel_stack_pages);
+
 void __init __no_sanitize_address setup_arch(char **cmdline_p)
 {
 	setup_initial_init_mm(_text, _etext, _edata, _end);
-- 
2.56.0.rc1.315.gc6ed9934b7-goog