[Feature] 支持Loop Unroll Hint

[Zhendong404]

Summary

在PTOAS/PTODSL上支持loop unroll hint以支持Tilelang-PTO后端将T.unroll(..., explicit=False) lower成对等语义

Motivation / use case

Tilelang中T.unroll(..., explicit=False)语法在CUDA/ASC后端会生成一个loop unroll hint (#pragma unroll)，表示建议编译器进行unroll，实际unroll行为仍由编译器内部cost model决定。而PTOAS/PTODSL当前仅提供了pto.static_range语法，该语法在DSL前端对循环做强制展开，未考虑寄存器压力等cost model，对循环次数较多的loop使用会导致编译报错。因此亟需在PTOAS/PTODSL侧也提供与#pragma unroll类似功能的语法，实现编译器侧unroll。

Proposed API / behavior

PTOAS/PTODSL Loop Unroll Hint 开发计划

总体目标

在 PTODSL 中扩展 pto.for_，并新增与 Python range 对应的 pto.range。用户可以继续使用 Python 原生循环语法，同时通过 pto.range 携带 loop unroll hint：

for i in pto.range(0, N, unroll="full"):
    ...

for i in pto.range(0, N, unroll_factor=4):
    ...

普通 Python 循环保持不变：

for i in range(0, N):
    ...

阶段一由 Bisheng/CCE 执行展开，阶段二由 PTOAS 原生展开。

阶段一：PTODSL 接口与 Python AST Rewrite

pto.for_ 扩展

在现有接口上增加可选参数：

pto.for_(
    start,
    stop,
    *,
    step,
    unroll=None,
    unroll_factor=None,
)

约束：

• unroll="full" 表示完全展开。
• unroll_factor=N 表示指定展开因子。
• 两者互斥，factor 必须为正整数。
• 未指定 hint 时保持现有行为和 IR 不变。
• .carry(...) 继续附着于同一个 loop builder；hint 不能改变 loop-carried value 语义。

pto.range 设计

新增 Python range-like 前端接口：

pto.range(stop, *, unroll=None, unroll_factor=None)
pto.range(start, stop, *, step=1, unroll=None, unroll_factor=None)

它只用于 PTODSL tracing/AST rewrite，不作为普通 Python runtime iterable 使用。

设计要求：

• 保持 Python range 的 start/stop/step 语义。
• 支持整数常量、PTODSL 编译期值和 runtime loop bound。
• 支持 unroll="full" 和 unroll_factor=N。
• 参数校验规则与 pto.for_ 完全一致。
• pto.range(...) 不引入独立的 loop IR；它只是为原生 for 携带 hint 的 frontend marker。

AST Rewrite

扩展现有 Python 原生循环改写逻辑：

• for i in range(...) 继续按现有规则转换为无 hint 的 pto.for_。
• for i in pto.range(...) 被识别为带 hint 的循环。
• AST rewrite 解析 pto.range 的参数，规范化为：

with pto.for_(
    start,
    stop,
    step=step,
    unroll=unroll,
    unroll_factor=unroll_factor,
) as i:
    ...

• 两种写法必须生成相同的 loop bounds、step、attribute 和 SSA 语义：

for i in pto.range(..., unroll_factor=4):
    ...

with pto.for_(..., unroll_factor=4) as i:
    ...

• break、continue、loop-carried value 等现有原生控制流限制保持不变，并提供针对 pto.range 的明确诊断。
• 嵌套循环只影响直接使用 pto.range 的层级。
• pto.range 不得被误识别为 pto.static_range；后者仍保持 trace-time 展开语义。

Loop Hint 的 IR 表达

统一将 pto.for_ / pto.range 的 hint 编码到目标 scf.for：

{pto.unroll = "full"}
{pto.unroll_factor = 4}

需要实现：

• PTODSL builder 设置 attribute。
• verifier 校验类型、互斥关系和 factor 范围。
• attribute 在 PTOAS pipeline 中保持不丢失。
• 对不支持的循环形态保留原循环并产生明确诊断，不得静默改变语义。

阶段一：Bisheng/CCE Bypass

在 PTOAS lowering 阶段将 loop attribute 转换为 Bisheng/CCE 可识别的 pragma 或等价 metadata：

• full hint 映射到 CCE full-unroll hint。
• factor hint 映射到 CCE factor-unroll hint。
• 确保 hint 出现在目标循环的有效作用域中，并保留到 Bisheng 输入。
• 对无法映射的 hint 提供 warning/error 策略；默认保留原循环语义。
• 阶段一不复制循环体，不生成 remainder loop。

阶段二：PTOAS 原生 Loop Unroll

扩展现有 PTOUnrollSIMTFor，或实现通用 loop-unroll pass：

• full unroll：复制静态循环体并替换 induction variable。
• factor unroll：生成展开主循环和 remainder loop。
• 首版限制为静态 lower/upper/step、正 step 和可计算 trip count。
• 正确处理 loop-carried SSA、嵌套循环、条件分支和 memory side effects。
• 明确 barrier、pipeline 和其他同步操作在展开时的复制与顺序规则。
• 在 scf.for 转 CFG/LLVM 前执行。
• 展开后运行 canonicalize、SCCP、CSE 和必要的 dead-code elimination。
• 通过配置保证阶段一 CCE bypass 与阶段二 native unroll 不会对同一循环重复展开。

测试计划

PTODSL/API

• pto.for_(..., unroll="full") 生成 full attribute。
• pto.for_(..., unroll_factor=4) 生成 factor attribute。
• for i in pto.range(...) 可被 AST rewrite 正确改写。
• range、pto.range、pto.for_ 三者在无 hint 时行为兼容。
• pto.range 的 start/stop/step 组合与 Python range 一致。
• .carry(...)、嵌套循环和 runtime bound 测试。
• 非法参数、冲突参数、动态 factor 和不支持控制流产生稳定诊断。

PTOAS/Bisheng

• PTODSL frontend 编译通过。
• attribute verifier 和 lowering 测试通过。
• full/factor hint 生成预期 CCE pragma/metadata。
• Bisheng 编译通过，并确认 hint 未在中间产物中丢失。
• 使用 SIMT/VF kernel 完成 PTODSL -> PTOAS -> Bisheng 闭环验证。

Native Unroll

• full unroll 展开次数和 induction variable 替换。
• factor unroll 的整除、不整除、trip count 小于 factor、0/1 次迭代。
• 嵌套循环、条件分支、loop-carried value 和同步操作。
• 动态边界、负 step 和不支持场景的诊断。
• 阶段一/阶段二互斥配置回归测试。

默认假设

• 本计划只覆盖 PTOAS/PTODSL，不包含 TileLang codegen 修改。
• pto.for_ 是实际 loop builder，pto.range 是 Python 原生 for 的 hint carrier。
• 普通 range 与现有 pto.static_range 的语义保持不变。
• 阶段一由 Bisheng/CCE 完成循环展开；阶段二才由 PTOAS 修改循环结构。
• 首版支持 full unroll 和指定 factor，不承诺任意动态循环均可展开。

Alternatives considered

No response

Additional context

No response

[Zhendong404]

PR #838 已经支持了SIMT内的Unroll，当时作为临时方案规避BUG，本次可以一起重构
#838 (comment)