PSI 任务计数位遍历优化

系列概况

• 标题: [PATCH v2] sched/psi: use __ffs() to walk task-count bitmasks in psi_group_change()
• 作者: Usama Arif usama.arif@linux.dev
• 版本: v2（单 patch，非系列）
• 修改文件: kernel/sched/psi.c
• 代码统计: 12 行新增，8 行删除；psi_group_change() 生成代码在 -O2 -march=x86-64 下减少 67 字节（756 → 689）
• 微基准规模: Zen4c 上 pinned CPU、IRQ关闭、min-of-10 cyc/call，覆盖 6 种 mask 分布
• Message-ID: 主帖 20260717105939.203685-1-usama.arif@linux.dev；回复 20260807095247.1010546-1-usama.arif@linux.dev
• 完整性: commit message、benchmark 数据、版本变化、签名齐备；已获 Acked-by: Johannes Weiner（PSI 维护者）

补丁目的

psi_group_change() 负责在任务状态切换（sleep / iowait / memstall / wake 等）时更新每 CPU 上 groupc->tasks[t]（4 个 NR_PSI_TASK_COUNTS 计数器）。函数需要分别遍历两个 4-bit 位掩码 clear 和 set：

• 对 clear 中每个置位 → 对应计数器自减；
• 对 set 中每个置位 → 对应计数器自增。

目标是把这个遍历从"扫到最高位"改成"只走置位"，从而在 1~2 bit 占绝大多数的真实负载中拿到显著加速，并顺手瘦身生成代码。

旧流程的问题

旧写法：

for (t = 0, m = clear; m; m &= ~(1 << t), t++) {
    if (!(m & (1 << t)))
        continue;
    /* 减一 */
}
for (t = 0, set; set &= ~(1 << t), t++)
    if (set & (1 << t))
        groupc->tasks[t]++;

问题点：

1. 循环上限跟着 m 的最高位置位走，单 bit 但 bit=3 也要走 4 次。
2. 进入循环体后还要 if (!(m & (1 << t))) continue;，绝大多数轮是 skip path。
3. 微基准开销随最高位增长：单 bit sleep 场景 9.60 cycle，memstall 场景（bit2、3）已经爬到 11.87 cycle。
4. 编译器要保留一个"常数 1"scratch 寄存器，位清除要走 SHL + NOT + AND 而不是 LEA + AND。

新流程

clear_orig = clear;
while (clear) {
    t = __ffs(clear);
    clear &= clear - 1;
    /* 用 clear_orig 做诊断打印 */
}
while (set) {
    t = __ffs(set);
    set &= set - 1;
    groupc->tasks[t]++;
}

__ffs(m) 返回最低 set bit 的下标；m &= m - 1 把最低 set bit 清零；循环在 m == 0 时自然退出。clear_orig 仅用于 printk_deferred 输出 task underflow 时的原始 mask。

Patch 概览

整条线索只有一个 patch：

Patch	文件	关键改动
1/1	kernel/sched/psi.c	psi_group_change() 中 clear/set 遍历改成 __ffs() + m &= m - 1

关键实现

@@ -798,7 +798,7 @@ static void psi_group_change(struct psi_group *group, int cpu,
- unsigned int t, m;
+ unsigned int t, clear_orig;
@@ -820,27 +820,31 @@ static void psi_group_change(struct psi_group *group, int cpu,
+ clear_orig = clear;
+ while (clear) {
+     t = __ffs(clear);
+     clear &= clear - 1;
      groupc->tasks[3], clear_orig, set);  /* 仍用 clear_orig 打日志 */
+ while (set) {
+     t = __ffs(set);
+     set &= set - 1;
+     groupc->tasks[t]++;
+ }

值得注意的细节：

• 作者在常见 PSI 路径产出的 mask 分布上跑 noinline 微基准，6 个用例中 5 个非空场景提速 47-63%，empty 场景不变。
• 生成代码小 67 字节：少一个"常数 1"scratch；位清除从 SHL+NOT+AND 变 LEA+AND；无 per-position skip 检查。
• 等价性：旧逻辑里 tasks[t]-- 在 clear 为 0 时不会进入，新版通过 while (clear) 保证同样的不变量；set 同理。

类比

旧写法像停车场管理员从 1 号车位一路走到最远那辆车的位置，每车位都要探头看一眼有没有车。新写法像管理员只去"还有车"的车位，每办完一辆就把记录划掉，全部划完就下班。位掩码"只走 set bit"是内核里非常常见的优化套路，clear_bit / for_each_set_bit / popcount 类工具都建立在这个观察上：低密度位图里，置位数量比位宽更能描述工作量。

+----------------------------------------------------------+
|                   OLD: linear scan                       |
+----------------------------------------------------------+
|                                                          |
|  for t = 0 .. max_bit                                    |
|       |                                                  |
|       v                                                  |
|  +----------+    yes    +-----------------+              |
|  | m & bit? |---------->| continue / skip |--+           |
|  +----------+           +-----------------+  |           |
|       | no                                   |           |
|       v                                       |           |
|  +-----------------+                          |           |
|  | process tasks[t]|                          |           |
|  +-----------------+                          |           |
|       |                                       |           |
|       +---------------------------------------+           |
|                                                          |
|  cost = highest set bit  (max 4 trips for 4-bit mask)     |
+----------------------------------------------------------+

+----------------------------------------------------------+
|                   NEW: iterate set bits                  |
+----------------------------------------------------------+
|                                                          |
|  while (m != 0)                                          |
|       |                                                  |
|       v                                                  |
|  +-----------------+                                      |
|  | t = __ffs(m)    |  (TZCNT / BSF: lowest set bit)      |
|  +-----------------+                                      |
|       |                                                  |
|       v                                                  |
|  +-----------------+                                      |
|  | m &= m - 1      |  (clear lowest set bit)             |
|  +-----------------+                                      |
|       |                                                  |
|       v                                                  |
|  +-----------------+                                      |
|  | process tasks[t]|                                      |
|  +-----------------+                                      |
|       |                                                  |
|       +--- loop back while m != 0                         |
|                                                          |
|  cost = number of set bits  (1 trip for any 1-bit mask)  |
+----------------------------------------------------------+

Highlight：风险与注意点

1. __ffs(0) 未定义：进入循环体时 m != 0，靠 while (m) 保证。如果后续有人把 __ffs 改到循环外（不推荐），必须显式判 0。
2. clear_orig 仅用于诊断：进入 task underflow 才打印，是新增的局部变量，热路径不读它。
3. 微基准只在 Zen4c 上跑：实际多核负载下 cycle 数字不一定能直接复现，但 1-bit 模式几乎贴近 empty floor 的结论应当普遍成立。
4. 缺 picked-up：有 Acked-by: Johannes Weiner，但还需要 scheduler 维护者（Peter Zijlstra / Tejun Heo / Ingo Molnar）显式 picked-up 才会被取走。
5. 第二条 message 是回复引用主帖：内容几乎全是 > 引用主帖 commit message 文字（末尾被截断显示在 "75..." 处），形态像 resend 或自己回复自己的引用块，需要关注后续维护者是否给出正式 ack/取走动作。

版本变化

• v1 → v2：从 for_each_set_bit() 改用 __ffs() + m &= m - 1。作者基于 benchmark 与代码尺寸收益选定 v2 形态。

一句话总结

把 psi_group_change() 里 4-bit 掩码遍历从"扫到最高位"换成"只走置位"，最常见的 1-bit 模式跑出接近空 mask 的开销，生成代码还小了 67 字节。