Now smp_call_function_single() would enable preemption before
csd_lock_wait() to reduce the critical section. To allow callers of
smp_call_function_any() to also benefit from this optimization, remove
get_cpu()/put_cpu() from smp_call_function_any().

Signed-off-by: Chuyi Zhou <zhouchuyi@bytedance.com>
Reviewed-by: Muchun Song <muchun.song@linux.dev>
---
 kernel/smp.c | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/kernel/smp.c b/kernel/smp.c
index b603d4229f95..80daf9dd4a25 100644
--- a/kernel/smp.c
+++ b/kernel/smp.c
@@ -761,16 +761,26 @@ EXPORT_SYMBOL_GPL(smp_call_function_single_async);
 int smp_call_function_any(const struct cpumask *mask,
 			  smp_call_func_t func, void *info, int wait)
 {
+	bool local = true;
 	unsigned int cpu;
 	int ret;
 
-	/* Try for same CPU (cheapest) */
+	/*
+	 * Prevent migration to another CPU after selecting the current CPU
+	 * as the target.
+	 */
 	cpu = get_cpu();
-	if (!cpumask_test_cpu(cpu, mask))
+
+	/* Try for same CPU (cheapest) */
+	if (!cpumask_test_cpu(cpu, mask)) {
 		cpu = sched_numa_find_nth_cpu(mask, 0, cpu_to_node(cpu));
+		local = false;
+		put_cpu();
+	}
 
 	ret = smp_call_function_single(cpu, func, info, wait);
-	put_cpu();
+	if (local)
+		put_cpu();
 	return ret;
 }
 EXPORT_SYMBOL_GPL(smp_call_function_any);
-- 
2.20.1