[RFC PATCH 2/2] x86/mm/tlb: Skip stale and lazy CPUs in arch_tlbbatch_flush()

[Bing Jiao <bingjiao@google.com>]

Commit 209954cbc7d0 ("x86/mm/tlb: Update mm_cpumask lazily") and
commit 6db2526c1d69 ("x86/mm/tlb: Only trim the mm_cpumask once a second")
stopped clearing outgoing CPUs from mm_cpumask(prev) on context switch.
For per-mm flushes, should_flush_tlb() skips CPUs where loaded_mm !=
info->mm, and flush_tlb_func() periodically trims stale bits from
mm_cpumask(f->mm).

Batched TLB flushes used by page reclaim, migration, and page demotion do
not benefit from this. arch_tlbbatch_add_pending() ORs the full
mm_cpumask(mm) into batch->cpumask, and arch_tlbbatch_flush() calls
flush_tlb_multi() with info.mm = NULL because a batch can span multiple
mm_structs. With !info->mm, should_flush_tlb() returns true for every
non-lazy CPU, and flush_tlb_func() cannot clear stale bits from
mm_cpumask(mm). As a result, every batch repeatedly sends IPIs to all
CPUs where the process previously ran if they are now running other
userspace tasks.

To fix this, record up to ARCH_TLBBATCH_MAX_MMS (8) unique mm_struct
pointers in arch_tlbbatch_add_pending() and filter batch->cpumask once
per batch in arch_tlbbatch_flush() via should_flush_tlb_batch(). Keeping
the array at 8 entries fits within a single 64-byte cache line in
task_struct->tlb_ubc, and batches spanning more than 8 mms fall back to
the unfiltered cpumask.

Tested on 112-CPU Intel Skylake, 384-CPU AMD EPYC Genoa, and 512-CPU AMD
EPYC Turin with CXL memory, all using IPI-based TLB shootdowns (5-run
mean). Following the setup in commit ed15468787dd ("x86/mm: Move
flush_tlb_info back to the stack") and commit 7e12beb8ca2a
("migrate_pages: batch flushing TLB"), a task with all CPUs set in its
mm_cpumask runs on 1 CPU (2 CPUs in test 4) while the other CPUs run
unrelated userspace tasks:

1. Pageout reclaim (81,920 4K pages reclaimed to zram):

   Platform          Metric        Baseline       Patched        Change
   --------------------------------------------------------------------
   112-CPU Skylake   Remote IPIs   9,092,156          191       -99.99%
                     Time (ms)       2,974.6        763.0  -74.3% (3.9x)
   384-CPU Genoa     Remote IPIs  31,377,184          157       -99.99%
                     Time (ms)       1,652.6        483.1  -70.8% (3.4x)
   512-CPU Turin     Remote IPIs  41,859,670          119       -99.99%
                     Time (ms)       1,695.7        476.3  -71.9% (3.6x)

2. Page demotion (81,876 4K pages demoted from top-tier to CXL):

   Platform          Metric        Baseline       Patched        Change
   --------------------------------------------------------------------
   512-CPU Turin     Remote IPIs     958,699          526        -99.9%
                     Time (ms)         213.0        168.2        -21.0%

3. Single-threaded move_pages() (331,600 4K pages migrated):

   Platform          Metric        Baseline       Patched        Change
   --------------------------------------------------------------------
   112-CPU Skylake   Remote IPIs      73,370        4,964        -93.2%
                     Time (ms)         702.4        699.7         -0.4%
   384-CPU Genoa     Remote IPIs     252,899       13,899        -94.5%
                     Time (ms)         392.4        347.1        -11.6%
   512-CPU Turin     Remote IPIs     337,426       18,193        -94.6%
                     Time (ms)         370.1        368.1         -0.5%

4. Multi-threaded migrate_pages() + active memory access (656,940 4K
   pages migrated, 2 active threads):

   Platform          Metric        Baseline       Patched        Change
   --------------------------------------------------------------------
   112-CPU Skylake   Remote IPIs     144,610        3,362        -97.7%
                     Pages/sec       386,385      391,050         +1.2%
   384-CPU Genoa     Remote IPIs     493,035        7,730        -98.4%
                     Pages/sec       673,166      725,680         +7.8%
   512-CPU Turin     Remote IPIs     658,193        9,685        -98.5%
                     Pages/sec       670,288      733,020         +9.4%

Fixes: 209954cbc7d0 ("x86/mm/tlb: Update mm_cpumask lazily")
Fixes: 6db2526c1d69 ("x86/mm/tlb: Only trim the mm_cpumask once a second")
Signed-off-by: Bing Jiao <bingjiao@google.com>
---
 arch/x86/include/asm/tlbbatch.h | 14 +++++++++++++
 arch/x86/include/asm/tlbflush.h | 16 ++++++++++++++-
 arch/x86/mm/tlb.c               | 36 +++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/arch/x86/include/asm/tlbbatch.h b/arch/x86/include/asm/tlbbatch.h
index 80aaf64ff25f..aac4fc5cba3a 100644
--- a/arch/x86/include/asm/tlbbatch.h
+++ b/arch/x86/include/asm/tlbbatch.h
@@ -4,12 +4,26 @@

 #include <linux/cpumask.h>

+struct mm_struct;
+
+#define ARCH_TLBBATCH_MAX_MMS	8
+
 struct arch_tlbflush_unmap_batch {
 	/*
 	 * Each bit set is a CPU that potentially has a TLB entry for one of
 	 * the PFNs being flushed..
 	 */
 	struct cpumask cpumask;
+
+	/*
+	 * Unique mm_struct pointers unmapped in this batch. Used only for
+	 * pointer equality against per-CPU loaded_mm in arch_tlbbatch_flush()
+	 * and never dereferenced there. nr_mms > ARCH_TLBBATCH_MAX_MMS
+	 * indicates overflow.
+	 */
+	struct mm_struct *mms[ARCH_TLBBATCH_MAX_MMS];
+	unsigned short nr_mms;
+
 	/*
 	 * Set if pages were unmapped from any MM, even one that does not
 	 * have active CPUs in its cpumask.
diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
index 4ec71ea0943f..0e86ade83ebc 100644
--- a/arch/x86/include/asm/tlbflush.h
+++ b/arch/x86/include/asm/tlbflush.h
@@ -379,8 +379,22 @@ static inline void arch_tlbbatch_add_pending(struct arch_tlbflush_unmap_batch *b
 		struct mm_struct *mm, unsigned long start, unsigned long end)
 {
 	inc_mm_tlb_gen(mm);
-	if (!cpu_feature_enabled(X86_FEATURE_INVLPGB))
+	if (!cpu_feature_enabled(X86_FEATURE_INVLPGB)) {
+		unsigned short i, nr = batch->nr_mms;
+
 		cpumask_or(&batch->cpumask, &batch->cpumask, mm_cpumask(mm));
+
+		if (likely(nr <= ARCH_TLBBATCH_MAX_MMS)) {
+			for (i = nr; i > 0; i--) {
+				if (batch->mms[i - 1] == mm)
+					goto out;
+			}
+			if (likely(nr < ARCH_TLBBATCH_MAX_MMS))
+				batch->mms[nr] = mm;
+			batch->nr_mms = nr + 1;
+		}
+	}
+out:
 	batch->unmapped_pages = true;
 	mmu_notifier_arch_invalidate_secondary_tlbs(mm, 0, -1UL);
 }
diff --git a/arch/x86/mm/tlb.c b/arch/x86/mm/tlb.c
index 0c4da320831c..10b8e839e964 100644
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -1679,12 +1679,48 @@ void __flush_tlb_all(void)
 }
 EXPORT_SYMBOL_FOR_KVM(__flush_tlb_all);

+static bool should_flush_tlb_batch(int cpu,
+				    const struct arch_tlbflush_unmap_batch *batch)
+{
+	struct mm_struct *loaded_mm = READ_ONCE(per_cpu(cpu_tlbstate.loaded_mm, cpu));
+	unsigned short i;
+
+	if (loaded_mm != LOADED_MM_SWITCHING) {
+		for (i = 0; i < batch->nr_mms; i++) {
+			if (loaded_mm == batch->mms[i])
+				goto check_lazy;
+		}
+		return false;
+	}
+
+check_lazy:
+	/*
+	 * Order 'loaded_mm' and 'is_lazy' against their write ordering in
+	 * switch_mm_irqs_off(). Ensure 'is_lazy' is at least as new as
+	 * 'loaded_mm'.
+	 */
+	smp_rmb();
+
+	return !READ_ONCE(per_cpu(cpu_tlbstate_shared.is_lazy, cpu));
+}
+
 void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch)
 {
 	struct flush_tlb_info info;
 	bool remote_flush = false;
 	int cpu = get_cpu();

+	if (likely(!cpu_feature_enabled(X86_FEATURE_INVLPGB) &&
+		   batch->nr_mms && batch->nr_mms <= ARCH_TLBBATCH_MAX_MMS)) {
+		int target_cpu;
+
+		for_each_cpu(target_cpu, &batch->cpumask) {
+			if (!should_flush_tlb_batch(target_cpu, batch))
+				__cpumask_clear_cpu(target_cpu, &batch->cpumask);
+		}
+	}
+	batch->nr_mms = 0;
+
 	init_flush_tlb_info(&info, NULL, 0, TLB_FLUSH_ALL, 0, false,
 			    TLB_GENERATION_INVALID);
 	/*
--
2.56.0.385.gac671c819a-goog