[Bug] 动态 if 中的嵌套 while 编译失败：局部变量被提前读取 #1589

在编译 TileLang 生成的 PTODSL 代码时遇到了 UnboundLocalError。简化后发现，只用 PTODSL 也能复现：外层动态 if 中有两层 while，内层循环定义并使用一个局部条件变量，编译时却要求这个变量在进入外层 if 前就已经定义。

环境

• Python 3.11.16
• CANN 9.3.0
• ptoas==0.65
• target：a5

复现代码

保存为 repro.py，执行 python repro.py 即可复现，不需要在 NPU 上运行。

from ptodsl import pto


@pto.jit(target="a5")
def nested_while_predicate(limit: pto.i32):
    outer = pto.const(0, dtype=pto.i32)
    total = pto.const(0, dtype=pto.i32)

    if limit > 0:
        while outer < limit:
            inner = pto.const(0, dtype=pto.i32)

            while inner < limit:
                if inner >= 2:
                    break

                stop = inner >= 1
                if stop:
                    break

                inner = inner + 1

            total = total + inner
            outer = outer + 1

    _ = total


print(nested_while_predicate.compile().mlir_text())

编译报错，位置指向外层的 if limit > 0：

UnboundLocalError: cannot access local variable 'stop' where it is not associated with a value

这段代码中的 stop 每次使用前都会赋值。如果执行了前面的 break，后面也不会再读取 stop。循环结束后没有使用它，因此预期这段代码可以正常编译。

两个对照

下面两种改法都已验证可以编译通过：

1. 去掉外层的 if limit > 0，保留两层循环。

2. 在外层 while 的开头给 stop 加一个初值：

   while outer < limit:
       stop = pto.const(0, dtype=pto.i1)
       inner = pto.const(0, dtype=pto.i32)
       # 后面的代码保持原样

第二种改法保留了原来的 stop = inner >= 1。新增的初值不会用于条件判断，只是为了让编译通过。

可能的原因

查看当前安装版本的实现，问题似乎来自 _ast_rewrite_analysis.py 中 _live_before_stmt() 对 while 的分析：

loop_live = set(live_after) | set(test_info.loads) | set(body_info.loads)

这里的 body_info.loads 包含循环体里所有被读取的变量，也包括 stop 这种在当前迭代中先赋值、后使用的临时变量。

在嵌套循环中，这可能让 stop 被误判为进入外层分支时就需要的值。随后，_prepare_if_branch_bodies() 将它加入 branch_entry_names，生成保存分支进入值的代码，效果类似：

saved_stop = stop

这次读取发生在原始代码给 stop 赋值之前，于是报了 UnboundLocalError。

修复建议

建议检查循环活跃变量的计算，区分“循环体里读取过的变量”和“需要上一轮或进入循环前的值的变量”。

对于 stop，本轮赋值后才会读取，不需要保存旧值；对于 count = count + 1 这样的状态，则仍然需要保留跨迭代的值。嵌套循环、break、continue 和 while ... else 的路径也需要分别处理。

可以将上面的代码加入回归测试，并保留现有 #1256 相关测试，确认修复后既能排除局部临时变量，也不会漏掉真正需要传递的循环状态。

目前已验证上述复现和两种规避方式，尚未实施 PTODSL 修复。