[PATCH 5/5] x86/mm: Re-enable preemption before waiting for kernel TLB flushes

[Chuyi Zhou]

flush_tlb_kernel_range() uses on_each_cpu() to synchronously flush
kernel mappings on all online CPUs when using the IPI backend.

The synchronous wait can become longer as the number of online CPUs
grows, and a remote CPU with interrupts disabled can further delay
completion. Preemption remains disabled throughout this wait, delaying
higher-priority tasks on the initiating CPU. The outer guard prevents
the SMP layer from making its final completion wait preemptible.

The kernel range descriptor is private stack storage and remains valid
until on_each_cpu() returns. The SMP layer protects CPU selection, IPI
queueing and local callback execution, so the kernel flush caller does
not need to remain pinned while waiting for remote completion.

INVLPGB requires separate protection: TLBSYNC only waits for operations
issued on the same CPU. Disable preemption inside
invlpgb_kernel_range_flush() across the broadcast loop and TLBSYNC.
The full-flush helper invlpgb_flush_all() already provides this
protection.

Remove the outer preemption guard from flush_tlb_kernel_range() so the
IPI completion wait can be preempted when the caller's context allows
it. Keep both flush backends synchronous.

Signed-off-by: Chuyi Zhou <zhouchuyi@xxxxxxxxxxxxx>
---
 arch/x86/mm/tlb.c | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/arch/x86/mm/tlb.c b/arch/x86/mm/tlb.c
index 4f9f0a18dbbc..f7773ba4a152 100644
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -1473,6 +1473,9 @@ static void invlpgb_kernel_range_flush(unsigned long start, unsigned long end)
 {
 	unsigned long addr, nr;
 
+	/* Keep the INVLPGB operations and TLBSYNC on the same CPU. */
+	guard(preempt)();
+
 	for (addr = start; addr < end; addr += nr << PAGE_SHIFT) {
 		nr = (end - addr) >> PAGE_SHIFT;
 
@@ -1521,8 +1524,6 @@ static void kernel_tlb_flush_range(unsigned long start, unsigned long end)
 
 void flush_tlb_kernel_range(unsigned long start, unsigned long end)
 {
-	guard(preempt)();
-
 	if (end == TLB_FLUSH_ALL ||
 	    tlb_range_exceeds_ceiling(start, end, PAGE_SHIFT))
 		kernel_tlb_flush_all();
-- 
2.20.1