pypto-optimization

PyPTO 性能优化（持续补充）

规则 1：连续搬运先达阈值，再在归约轴找甜点

对包含归约的算子（softmax、logsoftmax、sum/amax/amin，以及 mean=sum/count 语义）按以下优先级：

• 对固定题型 shape=(16,256,256), dim=1，默认模板直接用 set_vec_tile_shapes(1, 16, 256)。
• 先满足硬约束：prod(tile_shape) <= 16384 且 auto_tiles <= 2048。
• 若 auto_tiles > 2048，优先引入 loop + view/assemble 分块，再做 tile 微调。
• 若编译/验证出现 UB 或 OoOSchedule 相关报错，优先降 prod(tile_shape)（常见从 16384 -> 8192 -> 4096）。
• 先让连续搬运达经验阈值：contiguous_bytes(tile) >= 1KB（经验值，第一性能门槛）。
• 在同等可编译约束下，未达 1KB 的候选默认淘汰，不能仅因“规约轴不分段”直接入选。
• 只有当可编译候选都达不到 1KB 时，才在 <1KB 候选中比较规约轴分段与实测。
• 达到阈值后，不要默认“规约轴越大越快”；对规约轴 tile 做候选实测（默认顺序 16 -> 32 -> 64），按实测选甜点。
• tile 不浪费：优先让每一维 tile[i] <= shape[i]。tile[i] 远大于对应维度通常不会增加有效并行，反而会浪费 tile 预算并抬高 auto-tiling 开销。
• 禁止误读：不是“越连续越好”也不是“归约轴不切分”。连续搬运达标后，目标转为减少规约分段，而不是继续放大非规约轴 tile。

其中（关键，禁止误算）：

• contiguous_bytes(tile) = contiguous_tile_elems * dtype_bytes
• contiguous_tile_elems 指一次连续搬运段里的 tile 元素数，不是原始 shape 元素数。
• Vec 场景默认按最后一维估算：contiguous_tile_elems = tile[last_axis]（不做转置/置换时）。
• 连续搬运阈值判定前，优先满足 tile[last_axis] <= shape[last_axis]；不要通过 tile > shape 做“折算达标”。
• FP32 常用阈值：contiguous_tile_elems >= 256（约 1KB）
• FP16/BF16 常用阈值：contiguous_tile_elems >= 512（约 1KB）
• 反例：shape=(16,256,256), dim=1, tile=(1,256,64) 时，连续搬运按 tile[2]=64 算，仅 64*4=256B，未达到 1KB。

原因

• 连续搬运不足时，访存碎片和跨步开销会先成为瓶颈。
• 连续搬运达到高效区后，进一步放大通常边际收益很低。
• 归约轴 tile 过大时，单 task 可能过胖（局部归约树更重、寄存器/流水压力更高）。
• 归约轴 tile 过小时，分段与合并开销会上升。
• 因此常见是非单调关系（U 型），需要在达标候选中找甜点，而不是单调追大。

示例 A：Softmax (16, 16384), dim=1

• set_vec_tile_shapes(1, 8192)：每行 2 段（优先）
• set_vec_tile_shapes(2, 4096)：每行 4 段
• set_vec_tile_shapes(4, 4096)：每行 4 段，且更多预算给了非归约轴

经验上，(1, 8192) 通常优于 (2, 4096) 和 (4, 4096)。

示例 B：Reduction (16, 64, 256, 256), dim=1

• set_vec_tile_shapes(1, 16, 1, 256) / (1, 32, 1, 256) / (1, 64, 1, 256)：均满足连续搬运达标，需做甜点比较。
• set_vec_tile_shapes(1, 1, 16, 256)：虽然连续搬运达标，但归约轴未被有效利用，通常是劣候选。

这个例子体现了层级规则：连续搬运先达标，达标后做归约轴甜点搜索。

示例 C：TripletMarginLoss Phase-1 (128, 4096), dim=1

• set_vec_tile_shapes(4, 4096)：归约轴完整覆盖，且 batch 轴并行度更高（优先）
• set_vec_tile_shapes(1, 16384)：第二维 tile 明显超过真实维度 4096，存在 tile 预算浪费

经验上，(4, 4096) 通常优于 (1, 16384)。

示例 D：3D Max/Reduction (16, 256, 256), dim=1

• 候选 set_vec_tile_shapes(1, 256, 64)：
  • 归约轴不分段，但连续搬运仅 64*4=256B（FP32），未达 1KB 阈值。
• 达标候选：
  • set_vec_tile_shapes(1, 16, 256)：连续搬运达标，归约轴分 16 段（本形状默认优先）。
  • set_vec_tile_shapes(1, 32, 256)：连续搬运达标，归约轴分 8 段（备选）。
  • set_vec_tile_shapes(1, 64, 256)：连续搬运达标，归约轴分 4 段（不建议默认使用）。

这个例子说明：连续搬运阈值检查必须基于 tile。对该固定形状，默认直接采用 (1,16,256)，不以“中段起步”替代实证结论。

示例 E：拒绝“折算达标” (16, 256, 256), dim=1

• 候选 set_vec_tile_shapes(1, 32, 512)：
  • 写法上可编译，但 tile[2]=512 > shape[2]=256，属于预算浪费候选，不应拿“effective=256”当达标理由。
• 结论：
  • 阈值判断以 tile 的实际连续段为准；先选 tile[last_axis] <= shape[last_axis] 的候选，再做 16/32/64 甜点比较。

适用边界

• 这是性能优先规则，不是语义规则，最终以实测为准。
• 当非归约轴很大且后端实现特殊时，例外可能存在。
• 连续搬运阈值 1KB 是经验起点，不是硬约束；可按 profile 结果微调。
• 阈值判断使用 tile 的实际连续段，不使用 shape 的全维长度。
• 对静态 shape 的 benchmark，优先避免 tile[last_axis] > shape[last_axis] 的“折算达标”写法。
• 对两段 tile 的流水线（如 CrossEntropy 的 per-sample 阶段与 batch 归约阶段），要分阶段分别应用本规则。
• 对 RMSNorm/BatchNorm 这类 3D 归约，若归约轴是 C 且 C 中等（如 32/64/128），优先让最后一维连续搬运先达 1KB，再对 C 轴 tile 做甜点比较。
• tile[i] > shape[i] 不是语法错误，但通常是性能信号，除非有明确实测收益。
• 无法也不需要手动精确计算 UB 占用；以 prod(tile_shape) 分级降档 + 编译实测为准。

规则 2：loop count 取中间甜点，不取极端

适用于固定总工作量、通过 loop + view/assemble 沿外轴分块的场景（matmul、norm、reduction、loss 的分块实现都常见）。

• 总量固定时，通常满足：total_batch = loop_count * main_batch。
• 调 loop_count 本质是在调“任务粒度”：loop_count 越大，单任务越小；loop_count 越小，单任务越大。
• 经验规律通常是 U 型：存在中间甜点，性能最好；两端都可能变慢。

原因（通用）

• loop_count 过大（任务过碎）：
  • task 数量和依赖边增多，ready-queue/调度/唤醒开销放大。
  • 常见现象：Wait Schedule Time、Wait Predecessor Time 显著上升。
• loop_count 过小（任务过胖）：
  • 单 task 内部子图变重，单次执行时间明显上升。
  • 并行粒度变粗，负载均衡变差，尾核拖慢总时延。

选型方法（通用可复用）

1. 粗扫候选：loop_count ∈ {8, 16, 32, 64, 128}（或按问题规模等比扩展）。
2. 记录同一 profile 流程下的统一指标：
   • gen_time_us（主目标）
   • avg Wait Schedule Time / core
   • avg Execute task num / core
   • 主耗时阶段（dominant PSG）的 sum_dur 与 avg_dur
3. 判定逻辑：
   • 若增大 loop_count 后 gen_time_us 上升且 wait 指标陡增，则进入“过碎区”。
   • 若减小 loop_count 后 gen_time_us 上升且主阶段 avg_dur/sum_dur抬升，则进入“过胖区”。
4. 在“过碎区”和“过胖区”之间选最优点，再做小范围微调。

常见误区

• 误区 A：“loop_count 越小越快”。
• 误区 B：“loop_count 越大越快（并行更多）”。
• 误区 C（Matmul 高频）：“BASIC_BATCH=128 是固定最佳值”。
• 正解：先扫点找甜点，不做单调假设。

Matmul 提醒（高频）

• 对 m 维 loop 的 matmul，禁止固定 BASIC_BATCH 常量（如 128/256）。
• 先在 loop_count 空间选候选，再反推 BASIC_BATCH = ceil_div(m, loop_count)。
• 推荐默认候选顺序：16/32 -> 8/64 -> 1/max（按需要早停）。
• 例：m=16384 时，loop=16/32 对应 BASIC_BATCH=1024/512；loop=8/64 对应 2048/256。
• 与任何示例常量冲突时，以本条为准。

偷懒起步法（推荐默认）

• 当 loop_count 可行范围较宽（如 1 ~ 128），不要从两端顺序扫。
• 直接从中间甜点候选起步（常用先试 16 或 32）。
• 这里的“中间”是对数刻度中间，不是算术中点：
  • loop 候选通常按 2 倍步长变化（1, 2, 4, ..., 128）。
  • 在该刻度下，1~128 的中段优先落在 16/32，可同时避开“过胖”和“过碎”两端。
• 快速流程：
  1. 先测中点候选（16/32）。
  2. 再向两侧各扩一跳（如 8、64）。
  3. 仅在需要时再补端点（1、128）。
• 目的：先快速避开两端极值，优先命中可用甜点区，再做小范围微调。

非单调示例（仅作证据，不是固定答案）

case35（GroupNorm）一次实测：

• loop=16: 630.30us（更优）
• loop=8: 649.58us
• loop=32: 649.42us
• loop=64: 804.96us
• loop=128: 936.66us

该例表明：16 优于 8/32，而 64/128 明显进入“过碎区”。

规则 3：连续规约轴优先合并（语义等价时）

PyPTO 的归约 API（如 sum/amax/amin）一次只支持单个 dim；mean 语义需用 sum/count 实现。
当 baseline 语义是“对多个连续轴做总规约”且不依赖中间轴结果时，优先在 forward 先 reshape 合并这些连续规约轴，再做一次单轴规约。

为什么

• 连续多次 sum(dim=...) 会引入中间张量与额外调度开销。
• 先合并连续规约轴，再单次规约，通常能减少中间写回与同步。

通用做法

先做判定，必须同时满足以下三条：

1. 规约轴在当前张量布局中是相邻连续轴（可通过 reshape 直接合并，不需要转置/置换）。
2. 中间规约结果不被其他算子复用（仅用于后续下一步规约，不作为分支输入/输出）。
3. 合并后语义合同完全一致（规约定义、缩放方式如 sum/batchmean/sum/count、输出形状一致）。

满足三条时：

1. 在 forward 合并连续规约轴（仅 reshape，不做置换）。
2. kernel 内改为单次规约 + 原有最终规约语义。

示例（仅示例）

• Norm 的空间归约 (N, C, H, W)：
  • 若目标是对 H,W 联合规约，可先变为 (N, C, H*W)，再对单轴规约。

何时不要合并

• 需要中间轴结果（例如中间张量还要参与其他算子或作为可观测输出）。
• 规约轴不连续，或合并会改变语义（并非纯 reshape 等价）。
• 合并后导致更差的 tile/loop 组合，实测无收益。

快速检查清单

• 算子是否含归约轴
• 连续搬运是否先达到经验阈值（约 1KB，按 tile 连续段计算）
• 是否避免用 tile[last_axis] > shape[last_axis] 做“effective 折算达标”
• 在连续搬运达标后，是否按 16 -> 32 -> 64 顺序比较归约轴候选（本形状默认首选 16）
• 在达标候选中，是否做了归约轴甜点比较（如 16/32/64）
• 是否避免明显的 tile[i] > shape[i]（无收益浪费）
• 是否满足 prod(tile_shape) <= 16384 与 auto_tiles <= 2048
• 若 auto_tiles > 2048，是否先改成 loop 分块再调 tile
• 若出现 UB/OoOSchedule 报错，是否先将 prod(tile_shape) 从 16384 降到 8192/4096 再试
• 是否至少实测了两组 tile 候选
• 是否做了 loop_count 粗扫（至少 3 个点，建议 5 个点）
• 是否同时排查了“过碎（调度等待）”与“过胖（单任务过重）”两端
• matmul 是否先由 loop_count 中段候选反推 BASIC_BATCH（而非照抄固定常量）
• 多轴规约是否评估过“连续规约轴先合并再单轴规约”的等价实现

待补充

• 规则 4：两段 tile 的 phase 划分准则
• 规则 5：算子融合与数值稳定性对性能的影响