[PATCH] sched/fair: Refactor can_migrate_task() to elimate looping

[I Hsin Cheng <richard120310@gmail.com>]

The function "can_migrate_task()" utilize "for_each_cpu_and" with a
"if" statement inside to find the destination cpu. It's the same logic
to find the first set bit of the result of the bitwise-AND of
"env->dst_grpmask", "env->cpus" and "p->cpus_ptr".

Refactor it by using "cpumask_first_and_and()" to perform bitwise-AND
for "env->dst_grpmask", "env->cpus" and "p->cpus_ptr" and pick the first
cpu within the intersection as the destination cpu, so we can elimate
the need of looping and multiple times of branch.

After the refactoring this part of the code can speed up from ~115ns to
~54ns, according to the test below.

Ran the test for 5 times and the result is showned in the following
table, and the test script is paste in next section.

-------------------------------------------------------
|Old method|  130|  118|  115|  109|  106|  avg ~115ns|
-------------------------------------------------------
|New method|   58|   55|   54|   48|   55|  avg  ~54ns|
-------------------------------------------------------

v1 -> v2:
- Use cpumask_first_and_and()
- Remove additional cpumask

Signed-off-by: I Hsin Cheng <richard120310@gmail.com>
---
Test is done on Linux 6.9.0-0-generic x86_64 with
Intel(R) Core(TM) i7-2600 CPU @ 3.40GHz

Test is executed in the form of kernel module.

Test script:

int init_module(void)
{
    struct cpumask cur_mask, custom_mask;
    struct task_struct *p = current;
    int cpu, cpu1 = nr_cpu_ids, cpu2 = nr_cpu_ids;
    unsigned tmp = 0;

    cpumask_copy(&cur_mask, cpu_online_mask);
    /* Self-implemented function, didn't paste here because the length */
    generate_random_cpumask(&custom_mask);

    ktime_t start_1 = ktime_get();
    for_each_cpu_and(cpu, &cur_mask, &custom_mask) {
        if (cpumask_test_cpu(cpu, p->cpus_ptr)) {
            /* imitate load balance operation */
            tmp |= 0x01010101;
            cpu1 = cpu;
            break;
        }
    }
    ktime_t end_1 = ktime_get();

    ktime_t start_2 = ktime_get();
    cpu = cpumask_first_and_and(&cur_mask, &custom_mask, p->cpus_ptr);

    if (cpu < nr_cpu_ids) {
        /* imitate load balance operation */
        tmp |= 0x01010101;
        cpu2 = cpu;
    }
    ktime_t end_2 = ktime_get();

    if (cpu1 != cpu2) {
        pr_err("Failed Assertion, cpu1 = %d, cpu2 = %d\n", cpu1, cpu2);
        return 0;
    }

    pr_info("Old method spend time : %lld\n", ktime_to_ns(end_1 - start_1));
    pr_info("New method spend time : %lld\n", ktime_to_ns(end_2 - start_2));

    return 0;
}
---
 kernel/sched/fair.c | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 2d16c8545..d49960d50 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -9404,12 +9404,11 @@ int can_migrate_task(struct task_struct *p, struct lb_env *env)
 			return 0;
 
 		/* Prevent to re-select dst_cpu via env's CPUs: */
-		for_each_cpu_and(cpu, env->dst_grpmask, env->cpus) {
-			if (cpumask_test_cpu(cpu, p->cpus_ptr)) {
-				env->flags |= LBF_DST_PINNED;
-				env->new_dst_cpu = cpu;
-				break;
-			}
+		cpu = cpumask_first_and_and(env->dst_grpmask, env->cpus, p->cpus_ptr);
+
+		if (cpu < nr_cpu_ids) {
+			env->flags |= LBF_DST_PINNED;
+			env->new_dst_cpu = cpu;
 		}
 
 		return 0;
-- 
2.43.0