Add generic reduction domains and Triton row-vector scheduling

[whjthu]

pytest output:

pytest -n 16 -x
====================================================================================== test session starts =======================================================================================
platform linux -- Python 3.12.3, pytest-9.0.3, pluggy-1.6.0                                
rootdir: /home/haojie/ntbackend/.merge-ready/reduction-domain-triton-schedule                                                                                                                     
configfile: pyproject.toml                                                                       
plugins: xdist-3.8.0, anyio-4.13.0, cov-7.1.0
16 workers [387 items]                                                                           
...................................................................................................................................................s...................................... [ 48%] 
.......................................................................................................................................................................................... [ 96%]
s..............                                                                                                                                                                            [100%]
=========================================================================== 385 passed, 2 skipped in 529.47s (0:08:49) ==========================================================================

背景

当前 SSA 编译链路对 reduction 的调度主要依赖启发式判断，缺少对归约域、并行域和结果使用关系的结构化分析。这带来两个问题：

1. Triton 可能使用 block mask 访问 scalar pointer，导致生成代码无法编译。
2. Softmax、LayerNorm 等包含归约与广播的数据流无法稳定生成高效的 row-vector kernel，性能明显低于旧版生成路径。

本 PR 从 SSA 类型、归约轴和 use-def 链中提取通用的 Reduction Domain，并据此完成调度。

主要修改

通用 Reduction Domain

新增不可变的 ReductionDomain，用于描述：

• reduction operand、result 和 sum/max/min 类型；
• 规范化后的 reduction axis；
• 输入、结果及 parallel domain shape；
• reduction extent；
• 同域 reduction 及其广播 consumer；
• row-vector 调度的合法性和拒绝原因。

分析同时覆盖：

• 正轴与负轴归一化；
• 多个同域 reduction；
• reduction 到算术、数学、条件、广播和 store 的数据流；
• scf.if、scf.for、scf.yield 和 loop-carried value；
• mem.store 与 mem.atomic_add 等 effect sink；
• 不同输出域、未知布局和不兼容 consumer 的 fail-closed 处理。

Triton Row-Vector 调度

当 Reduction Domain 满足约束时，Triton 使用 row-vector 调度：

• 一个 Triton program 对应一个 parallel domain 实例；
• tl.arange() 覆盖 reduction axis；
• block pointer、value 和 predicate 保持相同向量域；
• 使用 tl.sum、tl.max 或 tl.min 完成归约；
• reduced output 使用 scalar store；
• 广播后的完整 shape output 使用 vector store；
• 支持任意单轴归约，包括 axis 0、末轴和 rank >= 3 的中间轴。

不满足条件时进入明确的 scalar fallback，无法保证语义时直接报错，不生成可疑代码。

调优接入

复用现有 LaunchPlan 和 AutoTuner，不引入第二套调优系统：

• 默认 num_warps 候选为 4、8、1；
• num_stages=1；
• 用户显式配置仍具有最高优先级；
• max_num_configs=1 使用安全的 4-warps 配置；
• winner 继续由现有运行时 benchmark 和缓存机制管理。

后端隔离

Row-vector value semantics 仅由 Triton 后端消费。CUDA 和 TileLang 保持原有 lowering 行为，后续 PR 再做更新。

正确性覆盖

新增测试集中覆盖以下高价值场景：

• Reduction Domain 的 axis、shape 和同域多归约分析；
• sum/max/min 生成；
• Softmax 和 LayerNorm 形式的通用 SSA 数据流；
• axis 0、末轴及 rank-3 中间轴归约；
• 非二次幂 reduction extent；
• 动态 shape、stride 和非连续输入；
• 不兼容广播及不同输出域的 fail-closed；
• tuning candidate 传播和显式配置优先级。

没有为每个内部 helper 增加独立测试，测试主要保持在 IR 契约、生成源码和公共执行路径三个层次。

测试结果

三个后端分别执行完整并行测试：

默认后端	结果
Triton	385 passed, 2 skipped
CUDA	385 passed, 2 skipped
TileLang	385 passed, 2 skipped

两个 skip 均为已有的多 GPU 测试，当前测试服务器只有一张 GPU：

• AOT multi-device test
• Triton multi-context reload test

性能结果

测试平台为 NVIDIA L20。基准为历史 master：

c9ebd4950a185beed8d4c1db9ff4a1fd133934ae

覆盖：

• Softmax：4096 × {128, 256, 512, 1024, 2048, 4096}
• LayerNorm：4096 × {128, 256, 512, 1024, 2048, 4096}
• Held-out width：127、781、1127、4097

共 12 个主性能点和 8 个 held-out 点，全部数值正确。

下表中的比值为：

本 PR 延迟 / 历史 master 延迟

因此小于 1 表示本 PR 更快。

指标	结果
12 个主性能点延迟比几何平均	0.9659×
相对历史 master 的整体提升	约 3.4%
最差单点延迟比	1.0029×
最差单点回归	约 0.29%

所有主测试点均不超过历史 master 10%，整体性能优于历史 master。

改动范围

本 PR 只完成通用 Reduction Domain 和 Triton row-vector 物化。CUDA、TileLang 的并行归约物化与后端专属性能优化将通过后续独立 PR 完成。