NormalizeReturnOrder Pass¶

将每个 InCore 函数的返回 tuple 重新排序，使 return[i] 对应声明顺序中的第 i 个 Out/InOut 参数，并相应地重映射非 InCore 调用方中的 TupleGetItemExpr 索引。该 Pass 完成后，编排（orchestration）代码生成可以通过位置直接映射 out_indices[i]，无需追踪 tile.store / ForStmt yield 链。

概述¶

用户代码可以以任意顺序写 tile.store —— 先 out_b 后 out_a，或者与计算混排。流水线前期会原样保留 body 顺序，因此 InCore 的 ReturnStmt::value_ 中各输出可能与声明的 Out/InOut 参数顺序不一致。若不规范化，编排代码生成就必须沿着每个 return[i] 反向追踪赋值与 tile.store，才能确定它实际写入哪个参数 —— 这种分析应当属于 Pass，而不是代码生成层（参见 docs/zh/dev/codegen/00-pto_codegen.md）。

本 Pass 把契约规范化为「按位置 return[k] ↔ out_indices[k]」，分两步进行：

Step A0（返回值参数化规范化） —— 对每个 InCore、Group、 Spmd、Graph 函数，把每个属于参数回写（param writeback）的 tensor 返回值改写为直接引用对应参数（指针同一性），追踪由共享的 return_lineage 工具完成。 kernel 内部分配的输出（无法追踪到任何参数）和标量返回不受影响。
Step A（InCore 函数重写） —— 对每个 InCore 函数，计算一个使 ReturnStmt::value_ 与声明的 Out/InOut 参数顺序一致的置换，然后同步重写返回值与 Function::return_types_。
Step B（调用端索引重映射） —— 对每个非 InCore 函数（Orchestration / Group / Spmd / Graph / opaque），重写所有 TupleGetItemExpr.index_，前提是它的 tuple 操作数来源于 Step A 中被重排序的函数调用结果。新索引为 permutation[old_index]，因此调用结果上的观察者仍然把同名 SSA 变量绑定到同一物理输出。

对于 tensor 返回值已经直接指名形参、且返回顺序已经与 Out/InOut 参数声明顺序匹配的函数，本 Pass 是 no-op；对于不含 InCore、Group、Spmd、 Graph 函数的程序也是 no-op —— Step A0 覆盖这四类，所以「不含 InCore 函数」本身已不足以说明本 Pass 什么都不做。

流水线位置: Default 策略中 #20 —— 位于 SplitVectorKernel（#19）之后、LowerPipelineLoops（#21）之前。这样既保证所有 kernel 拆分 / tile 结构决策仍基于原始返回顺序完成，又确保下游 tile 级 Pass（LowerPipelineLoops、 CanonicalizeIOOrder、InitMemRef、MemoryReuse、AllocateMemoryAddr）以及最终的 PTO 编排代码生成都能看到规范化后的顺序。

API¶

C++	Python	级别
`pass::NormalizeReturnOrder()`	`passes.normalize_return_order()`	程序级

from pypto import passes
result = passes.normalize_return_order()(program)

Pass 属性¶

属性	取值
前置（Required）	`SplitIncoreOrch`、`IncoreTileOps`
产出（Produced）	`ReturnParamsExplicit`
失效（Invalidated）	—

SplitIncoreOrch 保证 InCore 工作已经被外提为独立函数；IncoreTileOps 保证函数体使用 tile 操作，从而 Step A 所依赖的 tile.store(_, _, out_param) 信号一定存在。本 Pass 产出 ReturnParamsExplicit（由 verify_return_params_explicit.cpp 校验）：每个 InCore/Group/Spmd/Graph 函数中属于参数回写的 tensor 返回值都以指针同一性引用对应参数，编排代码生成因此只需查表即可建立返回值到实参的映射。本 Pass 不使任何属性失效 —— SSA、规范化语句结构、内存推断等所有上游属性均被保留。

算法¶

Step A0 —— 把返回值规范化为参数引用¶

对每个 InCore / Group / Spmd / Graph 函数，CanonicalizeReturnValues 调用 return_lineage::ReturnedParamIndices（可追踪 Var 到 Var 别名、循环携带、 builtin 回写、tuple 调用的 TupleGetItem，以及 Group/Spmd 包装函数调用），把每个可追踪到参数的 tensor 返回值替换为参数 Var 本身。无法追踪的值（kernel 内部分配的输出）和标量保留原表达式。

张量 IfStmt 合流仅在所有分支都追踪到同一个参数时才能继续追踪。存储来源一致性与合流元数据分别保留，因此两个分支内部的收窄可以在合流后恢复。合流元数据取自追踪后的分支类型，包括已验证调用转发的 caller 类型。实际返回的张量必须与参数类型一致，包括视图（view）和有效形状（valid shape）元数据；合流后的 tensor.set_validshape 结果不能通过规范化丢失这些元数据。来源冲突、新分配张量或改变视图的分支保持原返回表达式。Graph 函数继续拒绝通过合流追踪返回值：其边界分析尚不支持合流张量的来源，接受这种返回值可能导致 replay 时冻结随调用变化的边界切片。变量的参数根和失败结果在分支及多个返回值之间共享缓存，避免连续条件回写重复追踪前驱合流。回写链检查最终结果的类型：临时收窄有效形状后用 tensor.assemble 恢复，仍可将返回值规范化为参数引用。

正是这一步让"返回 → 参数"映射无需分析即可读出。 它运行之后，返回位置 j 写回参数 i 当且仅当 ReturnStmt->value_[j] 就是 params_[i]（指针同一性）——这正是 IRProperty::ReturnParamsExplicit 所断言的。因此位于本 pass 及其之后的消费者（orchestration codegen、ClassifyIterArgCarry）调用 return_lineage::ExplicitReturnedParamIndices(func) 这一"函数局部的结构性读取"，而不再重跑跨函数追踪器。ReturnedParamIndices 只保留给：在该属性建立之前运行的调用方（ExpandMixedKernel、scope outliner）、本 pass 自身，以及必须独立重新推导才能起到校验作用的属性验证器。

Graph 被纳入该集合的理由与包装函数相同：编排代码生成正是通过这张映射把 Graph 调用的结果别名到调用端张量。Graph 的函数体写作 out = core(..., out, ...); return out，outliner 会把它变成 TupleGetItem 重绑定，因此返回值到达 codegen 时是一个 SSA 重命名。在 Graph 被排除期间，这张映射全部返回 nullopt，codegen 退回到位置启发式（返回位置 j ↦ 第 j 个 Out/InOut 参数）；一旦某个 Graph 的返回顺序与参数顺序不一致，它的每个结果都会被静默绑定到错误的张量，且没有任何诊断（#2601）。

由于它是 codegen 的前置条件，手工构造 IR 并直接调用 orchestration codegen 的测试必须先运行本 pass（见 tests/ut/codegen/_orchestration_codegen_common.py），正如它们必须先运行 DeriveCallDirections、MaterializeRuntimeScopes 和 ClassifyIterArgCarry 一样。

Step A —— 计算并应用每个函数的置换¶

对每个 InCore 函数，BuildReturnToParamMapping 单次遍历函数体（不含末尾的 ReturnStmt），通过三条规则维护一个 Var* → out_param_index 的映射：

规则	语句模式	行为
1. `tile.store` 写入 Out/InOut buffer	`lhs = tile.store(tile, offsets, out_param, ...)`	`lhs → param_index_of(out_param)`
2. Var 到 Var 别名传播	`lhs = rhs_var`（且 `rhs_var` 已被映射）	`lhs → lookup(rhs_var)`
3. `ForStmt` iter-arg yield	`for_stmt.iter_args[i].initValue_` 已被映射	`for_stmt.return_vars_[i] → lookup(initValue)`

随后对 ReturnStmt::value_ 中的每个值，先在该映射里查找其 Var，否则回退到与 Function::params_ 的直接身份匹配；若都未命中则对应位置返回 kNoParam，表示该位置「未检测到与 out 参数的关联」，保留其原始下标。

ComputeReturnPermutation 把映射变为 permutation[old_index] = new_index，其中 new_index 是匹配参数在 CollectOutIndices(func) 中的位置。出现以下任一情况都返回空置换（对该函数即为 no-op）：

函数体不含 ReturnStmt（开放 IR），或不含任何 Out/InOut 参数。
out_indices.size() > ret_to_param.size() —— 声明的输出参数数量多于返回值数量，分析不完整，不能构造越界置换。
计算出的置换是恒等置换（已规范）。

当置换非空时，ReorderReturns 通过 MutableCopy 克隆出新的 Function，将末尾的 ReturnStmt 替换为 value_[permutation[i]] = old_value_[i] 的版本，并同步置换 Function::return_types_，使类型列表与值列表始终对齐。

Step B —— 重映射调用端的 `TupleGetItemExpr`¶

对（Step A 已重写后的）程序中的每个非 InCore 函数， TupleIndexPermutationMutator 在单次 SSA 遍历中：

对每个 RHS 是 Call(GlobalVar) 或 Submit(GlobalVar)、且被调函数已被重排的 AssignStmt，同步置换表达式的元组元素类型，并创建使用该类型的新结果 Var。置换的是调用端已实例化的类型，因此保留 caller 的局部 shape 符号。 Submit 末尾的 TASK_ID、依赖及启动元数据保持不变。
在 reordered_tuple_vars_ 中记录旧结果 Var、新绑定和置换。
对 SSA 转换前具有多个定义的绑定，将每次重排调用写入新的物理临时变量，再将逻辑元组写回原绑定。分支和循环因此保留可供 ConvertToSSA 合并的共享身份，也支持重排调用与非重排赋值混用。
一旦该 Var 被重新赋值（赋给非重排序函数的调用、非 Call 表达式等），立即移除其条目，避免基于身份的查找读到失效绑定。
对每个 TupleGetItemExpr(tuple_var, k)，若 tuple_var 在该映射中，把索引重写为 permutation[k]，并使用新的元组绑定。
整元组使用（包括别名和转发）通过投影构造原逻辑顺序的元组；每个定义缓存一次，不重复执行调用。声明多个平铺返回位置的函数会把逻辑元组展开为对应返回值。

由于 Step A 重写函数签名与 Step B 重写调用端索引在同一次 Pass 调用中完成，出口程序状态自洽：每个 tuple 元素仍然绑定到同一个物理输出 buffer，只是用新的下标访问；即使输出的 dtype 或 shape 不同，调用、绑定和投影类型也保持一致。

约束¶

约束	原因
Step A 仅重写 `InCore` 函数	其他函数类型（`Orchestration` / `Group` / `Spmd` / `Graph` / opaque）遵循用户声明的返回形态；它们的调用端在 Step B 中被重映射。`Group`/`Spmd`/`Graph` 的返回值在 Step A0 中仍会被规范化为参数引用，但不会被重排
Step A0 在 Graph 函数之外要求类型兼容的分支一致性（branch consensus）	所有分支必须追踪到同一参数；最终返回张量必须保留参数的类型、视图和有效形状。允许恢复中间的收窄。Graph 合流返回值仍不支持
Step A0 不改动 kernel 内部分配的输出与标量	只有参数回写必须显式化；没有参数血缘的返回值没有可引用的参数
`out_indices.size() > ret_to_param.size()` 时跳过	不完整分析不能产生越界置换 —— 保留原状，让 verifier 捕获不一致
恒等置换 ⇒ 不重写	避免不必要的 `Function` 克隆，使 Pass 幂等
Step B 记录原结果 `Var` 并创建带有新类型的绑定	投影使用新的物理元组，其他消费者保留旧逻辑顺序；不修改现有 IR 节点

示例¶

两个 Out 参数，但 InCore body 写出顺序与参数声明顺序相反；编排函数按 ret[0] / ret[1] 默认对应 out_a / out_b 取出。Pass 完成后，InCore 返回顺序匹配参数声明顺序，编排函数中的 TupleGetItemExpr 下标也被相应重映射，使同一 SSA 值仍流入 a 和 b。

Before:

@pl.program
class Module:
    @pl.function(type=pl.FunctionType.InCore)
    def kernel(self, x: pl.Tensor[[16], pl.FP32],
               out_a: pl.Out[pl.Tensor[[16], pl.FP32]],
               out_b: pl.Out[pl.Tensor[[16], pl.FP32]]) \
            -> tuple[pl.Tensor[[16], pl.FP32], pl.Tensor[[16], pl.FP32]]:
        x_tile = pl.load(x, [0], [16])
        a_tile = pl.tile.add(x_tile, x_tile)
        b_tile = pl.tile.mul(x_tile, x_tile)
        out_b_store = pl.store(b_tile, [0], out_b)
        out_a_store = pl.store(a_tile, [0], out_a)
        return (out_b_store, out_a_store)        # ← 与 (out_a, out_b) 声明顺序不一致

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x, out_a, out_b):
        ret = self.kernel(x, out_a, out_b)
        a = ret[0]                                # ← 当前实际取的是 out_b
        b = ret[1]                                # ← 当前实际取的是 out_a
        return (a, b)

After:

@pl.program
class Module:
    @pl.function(type=pl.FunctionType.InCore)
    def kernel(self, x, out_a, out_b):
        x_tile = pl.load(x, [0], [16])
        a_tile = pl.tile.add(x_tile, x_tile)
        b_tile = pl.tile.mul(x_tile, x_tile)
        out_b_store = pl.store(b_tile, [0], out_b)
        out_a_store = pl.store(a_tile, [0], out_a)
        return (out_a_store, out_b_store)        # ReorderReturns: 置换 [1, 0]

    @pl.function(type=pl.FunctionType.Orchestration)
    def main(self, x, out_a, out_b):
        ret = self.kernel(x, out_a, out_b)
        a = ret[1]                                # TupleIndexPermutationMutator: 0 → 1
        b = ret[0]                                # TupleIndexPermutationMutator: 1 → 0
        return (a, b)

同一条 SSA 赋值（a = ...）仍然绑定到 pl.store(a_tile, ..., out_a) 的输出；变化的只是 tuple 访问路径。InOut 参数的处理方式相同。

完整用例参见 tests/ut/ir/transforms/test_normalize_return_order.py：

test_swapped_returns_reordered —— 上文展示的两个 Out 参数案例
test_already_ordered_noop —— 已规范的 IR 保持不变
test_single_return_noop —— 单个 Out 参数无需置换
test_non_incore_unchanged —— 不含 InCore 函数的程序为 no-op
test_three_returns_scrambled —— 三元置换
test_2d_tensor_reorder —— 2 维 tensor / 多维 offset
test_inout_param_reorder —— InOut 参数同样参与重排
test_branch_writeback_and_wrapper_return_reference_param —— 同一缓冲区的分支回写及 wrapper 转发（#2966）
test_dynamic_branch_call_returns_reference_caller_param —— 生成的 phi 类型保留 callee 局部形状变量
test_branch_return_with_unproven_identity_is_not_canonicalized —— 来源冲突、新张量、视图及有效形状分支
test_writeback_after_branch_merge_preserves_returned_valid_shape —— 两种返回顺序下保留合流后收窄的元数据
test_restored_writeback_metadata_references_param —— 两种返回顺序下，共享血缘的恢复结果与收窄结果
test_sequential_branch_writebacks_share_lineage —— 48 次连续合流共享血缘缓存

实现¶

头文件: include/pypto/ir/transforms/passes.h

Pass NormalizeReturnOrder();

实现文件: src/ir/transforms/normalize_return_order_pass.cpp

CanonicalizeReturnValues —— Step A0 改写器：通过 return_lineage::ReturnedParamIndices 把可追踪的 tensor 返回值替换为参数 Var。
BuildReturnToParamMapping —— Step A 分析：遍历函数体，将每个 ReturnStmt 值反向映射到 Out/InOut 参数下标。
CollectOutIndices —— 收集 ParamDirection 为 Out / InOut 的参数位置。
ComputeReturnPermutation —— 综合上述两个分析，得到最终的 permutation[old_index] = new_index；不需重写或分析不完整时返回空。
ReorderReturns —— 基于 MutableCopy(func) 构造新的 Function，置换 ReturnStmt::value_ 与 Function::return_types_。
TupleIndexPermutationMutator —— Step B 改写器：跟踪调用结果变量，并重写 TupleGetItemExpr 索引。

属性: include/pypto/ir/transforms/pass_properties.h

inline const PassProperties kNormalizeReturnOrderProperties{
    .required = {IRProperty::SplitIncoreOrch, IRProperty::IncoreTileOps},
    .produced = {IRProperty::ReturnParamsExplicit}};

Python 绑定: python/bindings/modules/passes.cpp

passes.def("normalize_return_order", &pass::NormalizeReturnOrder,
           "Create a return order normalization pass\n\n"
           "Reorders return tuple values in InCore functions so that return[i]\n"
           "corresponds to the i-th Out/InOut parameter in declaration order,\n"
           "and updates TupleGetItemExpr indices at call sites accordingly.");

类型存根: python/pypto/pypto_core/passes.pyi

def normalize_return_order() -> Pass:
    """Create a return order normalization pass."""

测试: tests/ut/ir/transforms/test_normalize_return_order.py