sw64: add support for batched TLB flushing during unmap

This patch enables ARCH_WANT_BATCHED_UNMAP_TLB_FLUSH for sw64 architecture.
By deferring TLB flushes during memory reclamation, we can significantly
reduce the number of Inter-Processor Interrupts (IPIs) sent between cores.
This is highly beneficial in multi-threaded workloads sharing the same
memory space, preventing IPI storms when swapping out pages actively
cached in multiple CPUs.

Signed-off-by: Jinyu Tang <tjytimi@163.com>
Signed-off-by: Gu Yuchen <guyuchen@wxiat.com>
Reviewed-by: He Sheng <hesheng@wxiat.com>
Signed-off-by: Gu Zitao <guzitao@wxiat.com>

Author: tjy-zhu

arch/sw_64/Kconfig

@@ -59,6 +59,7 @@ config SW64
 	select ARCH_USE_CMPXCHG_LOCKREF
 	select ARCH_USE_QUEUED_RWLOCKS
 	select ARCH_USE_QUEUED_SPINLOCKS
+	select ARCH_WANT_BATCHED_UNMAP_TLB_FLUSH
 	select ARCH_WANT_DEFAULT_BPF_JIT
 	select ARCH_WANT_FRAME_POINTERS
 	select ARCH_WANT_IPC_PARSE_VERSION

arch/sw_64/include/asm/tlbbatch.h

@@ -0,0 +1,11 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+#ifndef _ASM_SW_64_TLBBATCH_H
+#define _ASM_SW_64_TLBBATCH_H
+
+#include <linux/cpumask.h>
+
+struct arch_tlbflush_unmap_batch {
+	struct cpumask cpumask;
+};
+
+#endif /* _ASM_SW_64_TLBBATCH_H */

arch/sw_64/include/asm/tlbflush.h

@@ -91,4 +91,28 @@ extern void flush_tlb_kernel_range(unsigned long start, unsigned long end);
 
 #endif /* CONFIG_SMP */
 
+
+static inline bool arch_tlbbatch_should_defer(struct mm_struct *mm)
+{
+	bool should_defer = false;
+
+    /* If remote CPUs need to be flushed then defer batch the flush */
+	if (cpumask_any_but(mm_cpumask(mm), get_cpu()) < nr_cpu_ids)
+		should_defer = true;
+	put_cpu();
+
+	return should_defer;
+}
+
+static inline void arch_flush_tlb_batched_pending(struct mm_struct *mm)
+{
+	flush_tlb_mm(mm);
+}
+
+extern void arch_tlbbatch_add_pending(struct arch_tlbflush_unmap_batch *batch,
+				struct mm_struct *mm,
+					unsigned long uaddr);
+extern void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch);
+
+
 #endif /* _ASM_SW64_TLBFLUSH_H */

arch/sw_64/kernel/smp.c

@@ -10,6 +10,7 @@
 #include <linux/cpu.h>
 #include <linux/acpi.h>
 #include <linux/of.h>
+#include <linux/mmu_notifier.h>
 
 #include <asm/irq_impl.h>
 #include <asm/mmu.h>
@@ -821,6 +822,23 @@ void flush_tlb_kernel_range(unsigned long start, unsigned long end)
 }
 EXPORT_SYMBOL(flush_tlb_kernel_range);
 
+void arch_tlbbatch_add_pending(struct arch_tlbflush_unmap_batch *batch,
+					struct mm_struct *mm,
+					unsigned long uaddr)
+{
+	cpumask_or(&batch->cpumask, &batch->cpumask, mm_cpumask(mm));
+
+	mmu_notifier_arch_invalidate_secondary_tlbs(mm, 0, -1UL);
+}
+
+void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch)
+{
+	if (!cpumask_empty(&batch->cpumask)) {
+		on_each_cpu_mask(&batch->cpumask, ipi_flush_tlb_all, NULL, 1);
+		cpumask_clear(&batch->cpumask);
+	}
+}
+
 #ifdef CONFIG_HOTPLUG_CPU
 extern int can_unplug_cpu(void);
 int __cpu_disable(void)