[PATCH] ARM: Flush the entire TLB for large kernel ranges

[Bart Van Assche <bvanassche@acm.org>]

__purge_vmap_area_lazy() flushes the TLB for a single range that covers
all lazily freed vmap areas. On ARM, BPF JIT images and modules are
allocated below PAGE_OFFSET while other vmalloc allocations are above
it, so that range can span hundreds of megabytes. flush_tlb_kernel_range()
invalidates such a range one page at a time, which takes up to hundreds
of thousands of per-page TLB maintenance operations. That loop has no
voluntary preemption points and runs while vmap_purge_lock is held, so
other tasks that free vmalloc memory are blocked.

Flush the entire TLB on both SMP and UP configurations if the range
touches at least PTRS_PER_PTE (512) pages, matching MAX_DVM_OPS on
arm64.

Signed-off-by: Bart Van Assche <bvanassche@acm.org>
---
 arch/arm/include/asm/tlbflush.h | 30 +++++++++++++++++++++++++++++-
 arch/arm/kernel/smp_tlb.c       | 15 +++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/arch/arm/include/asm/tlbflush.h b/arch/arm/include/asm/tlbflush.h
index 38c6e4a2a0b6..c6a61f4488a9 100644
--- a/arch/arm/include/asm/tlbflush.h
+++ b/arch/arm/include/asm/tlbflush.h
@@ -9,6 +9,7 @@
 
 #ifndef __ASSEMBLY__
 # include <linux/mm_types.h>
+# include <asm/page.h>
 #endif
 
 #ifdef CONFIG_MMU
@@ -597,14 +598,41 @@ static inline void clean_pmd_entry(void *pmd)
 #define local_flush_tlb_range(vma,start,end)	__cpu_flush_user_tlb_range(start,end,vma)
 #define local_flush_tlb_kernel_range(s,e)	__cpu_flush_kern_tlb_range(s,e)
 
+/*
+ * Bound the number of per-page TLB invalidations to avoid soft lock-ups and
+ * long vmap_purge_lock hold times on large ranges (not necessarily a
+ * performance improvement). Matches MAX_DVM_OPS (PTRS_PER_PTE) on arm64.
+ */
+#define MAX_TLB_KERNEL_RANGE_PAGES	512
+
+static inline bool tlb_kernel_range_exceeds_max(unsigned long start,
+						unsigned long end)
+{
+	unsigned long pages;
+
+	if (end <= start)
+		return false;
+
+	pages = ((end - 1) >> PAGE_SHIFT) - (start >> PAGE_SHIFT) + 1;
+	return pages >= MAX_TLB_KERNEL_RANGE_PAGES;
+}
+
 #ifndef CONFIG_SMP
 #define flush_tlb_all		local_flush_tlb_all
 #define flush_tlb_mm		local_flush_tlb_mm
 #define flush_tlb_page		local_flush_tlb_page
 #define flush_tlb_kernel_page	local_flush_tlb_kernel_page
 #define flush_tlb_range		local_flush_tlb_range
-#define flush_tlb_kernel_range	local_flush_tlb_kernel_range
 #define flush_bp_all		local_flush_bp_all
+
+static inline void flush_tlb_kernel_range(unsigned long start,
+					  unsigned long end)
+{
+	if (tlb_kernel_range_exceeds_max(start, end))
+		local_flush_tlb_all();
+	else
+		local_flush_tlb_kernel_range(start, end);
+}
 #else
 extern void flush_tlb_all(void);
 extern void flush_tlb_mm(struct mm_struct *mm);
diff --git a/arch/arm/kernel/smp_tlb.c b/arch/arm/kernel/smp_tlb.c
index d4908b3736d8..957cc2eb662f 100644
--- a/arch/arm/kernel/smp_tlb.c
+++ b/arch/arm/kernel/smp_tlb.c
@@ -4,6 +4,8 @@
  *
  *  Copyright (C) 2002 ARM Limited, All Rights Reserved.
  */
+#include <linux/build_bug.h>
+#include <linux/pgtable.h>
 #include <linux/preempt.h>
 #include <linux/smp.h>
 #include <linux/uaccess.h>
@@ -234,6 +236,19 @@ void flush_tlb_range(struct vm_area_struct *vma,
 
 void flush_tlb_kernel_range(unsigned long start, unsigned long end)
 {
+	static_assert(MAX_TLB_KERNEL_RANGE_PAGES == PTRS_PER_PTE);
+
+	/*
+	 * __purge_vmap_area_lazy() may pass a range that spans from the
+	 * modules area to the end of the vmalloc area. Invalidating such a
+	 * range page by page takes hundreds of thousands of broadcast TLB
+	 * operations, without voluntary preemption points.
+	 */
+	if (tlb_kernel_range_exceeds_max(start, end)) {
+		flush_tlb_all();
+		return;
+	}
+
 	if (tlb_ops_need_broadcast()) {
 		struct tlb_args ta;
 		ta.ta_start = start;