跳转至

MemoryReuse Pass

利用依赖分析识别内存复用机会，并移除冗余的 alloc 操作。

概述

该 Pass 通过分析变量生命周期和依赖关系来实现内存共享。在同一内存空间中，生命周期不重叠的变量可以共享内存引用 (MemRef) 对象，从而减少内存占用。

应用 MemRef 共享后，该 Pass 还会移除冗余的 tile.alloc 语句 (Statement)——即那些不再被任何 TileType 变量引用的 MemRef 对应的 alloc 语句。

核心要点：

  • 生命周期不重叠的变量可以复用内存
  • 只有在同一内存空间中的变量才能共享 MemRef
  • 生命周期通过 def-use 分析确定
  • 共享完成后，已无引用的 MemRef 及其 alloc 语句会被清理

使用时机：这是 MemoryPlanner.PYPTO 的机会性复用阶段，在 MaterializeSemanticAliases 之后、 AllocateMemoryAddr 之前运行。 MemoryPlanner.DSA_RP 会跳过它，以便 DSA-RP 求解器仍能看到独立 buffer； MemoryPlanner.PTOAS 也会跳过它，因为生命周期复用由 ptoas 负责。三种模式中的 循环 carry 与原地强制别名都已由 MaterializeSemanticAliases 建立。

规划器边界

MemoryReuse 在地址分配前先选择共享的 MemRef 身份。DSA_RP 则保留这些独立 身份，并在 AllocateMemoryAddr 中结合容量与复用惩罚联合选择地址。二者同时运行 会在 DSA-RP 评估之前删除候选方案。生命周期干涉、语义 no-alias 规则与目标 hazard 等正确性事实，在 DSA-RP 问题中仍然是硬约束。

API

C++ Python 级别
pass::MemoryReuse() passes.memory_reuse() 函数级

工厂函数：

Pass MemoryReuse();

Python 用法：

from pypto.pypto_core import passes

reuse_pass = passes.memory_reuse()
program_optimized = reuse_pass(program)

算法

  1. 生命周期分析：遍历完整 IR 树（包括嵌套控制流体内的语句）通过 def-use 分析计算变量生命周期。在循环外定义但在循环内使用的变量，其生命周期会延展到循环结束（循环感知延展）
  2. 干涉检查：识别生命周期重叠的变量
  3. MemRef 共享（全局「最大优先 + first-fit」装箱，IdentifyReuseOpportunities）：在每个内存空间内，按 大小从大到小 装箱；后续每个区间加入第一个其全部成员都能与之共享的缓冲区（生命周期不重叠 + hazard / no-alias 安全，见 can_share）。缓冲区的分配大小由其首个（最大）成员固定，因此之后纳入更小的成员是「免费」的 —— 且 后定义的较大区间 现在可以承载 先定义的较小区间。（此前的定义序贪心带有单向的大小门槛 source.size >= target.size，因此两个生命周期不相交、但较小者先定义的 tile 永远无法合并。）每个成员被重定位到的「代表」是该缓冲区的最大成员；由于 InitMemRef 会把所有 tile.alloc 提升到函数体头部，代表的 alloc 支配整个函数，因此代表即使定义在其部分成员之后也是安全的。由于装箱器不再按程序序处理，每个成对门槛（hazard、no-alias）都会在两个方向上检查。

容量溢出时的子区间回退（仅 PYPTO）。 本 pass 先计算上述旧式整块装箱；只要其精确 allocator footprint 不超过 backend 容量，就逐字节保留原布局。只有确认溢出后，才重试把较小 root 放到现有最大成员 arena 的对齐子区间。例如，一块已经失活的 64 KiB root 预留 [0, 65536) 后，两个随后且同时存活的 32 KiB root 可以分别占 [0, 32768) 与 [32768, 65536)。64 KiB root 的生命周期必须与两个小 root 都不重叠；两个小 root 互相干涉，所以它们的字节区间必须分离。重试为候选选择不被「不能与它共享字节」的成员阻塞的最低对齐空洞。arena 永不扩容，所以两个 32 KiB arena 不会拼成新的 64 KiB arena；64 KiB 代表必须原本就存在。若子区间装箱后仍放不下，原有的流水线深度削减与 legacy 回退照常执行。

子区间复用移动的是完整 root allocation；它不会利用 alias/view 未访问的内部空洞。候选 root 与 arena root 的 allocation offset 因此仍须是静态非负值。只有直接内部 tile.slice 才支持动态 offset，因为 PTO 能使用已经重定位的 source tile 与 slice 操作数重建其运行时地址。如果某个 memory space 中存在其他动态 offset 成员，例如动态 slice 之后的 reshape、reinterpret 或 transpose view，则整个 space 都禁用子区间回退，避免压缩无关 root 后把原本的 legacy overflow 变成错误代码。常量成员 offset 继续检查非负且不越过 root；受支持的动态 slice 原样保留相对表达式 arena_offset + placement + (member_offset - root_offset)。动态 root offset 仍不具备资格，声明式/pinned allocation 也保持封闭。生命周期干涉、pipeline 分离、目标 hazard 与算子 no-alias 规则仍由 can_share 判断；任何规则只要禁止字节重叠，两个 root 地址区间就必须分离。AllocateMemoryAddr 使用 arena root 的完整 allocation capacity 计算物理槽大小；所有常量子区间在 placement 阶段已经验证不会越过该容量。纯动态 tile.slice 的 MemRef 随后沿用旧行为记录 bare-base 地址，实际窗口由 PTO 使用 source tile 和 slice 操作数生成 pto.subview。后续 reshape fold 比较 base、offset 与 tile signature，PTO DPS alias 还会比较 extent。因此不会只因 base pointer 相同就把同一 arena 的兄弟 root 子区间误认为同一个 tile。 4. 循环携带变量重对齐（AlignLoopCarriesToInitMutator）：共享（步骤 3）只会重写由 AssignStmt 定义的变量（producer/init），而循环携带的 iter_arg/return_var 节点被排除在生命周期/共享映射之外、仍保留原始 MemRef。本步骤自外向内遍历 ForStmt，将每个循环的 iter_arg/return_var 重对齐到其（已复用的）initValue 的 MemRef，并在递归前写入 var_remap_，使嵌套循环能观察到已修正的外层 iter_arg 作为其 init。若缺少本步骤，被复用的嵌套流水化 matmul_acc 累加器会分裂到两个 Acc 缓冲区，导致步骤 7 插入非法的 acc→acc tile.move，被 Ascend 910B 的 ptoas 拒绝（#1352） 5. 累加器 if-phi 合并（TopDownRetargeter::CoalesceAccumulatorIfPhis）：LowerPipelineLoops 会把 stage-2 的 K 循环剥离成 if-phi，其活跃分支是就地累加的 matmul_acc（位于累加器缓冲区），而失效的 if k==0 分支是位于不同 Acc 缓冲区上的全新 matmul seed。若不处理，步骤 7 会尝试用 acc→acc tile.move 协调二者 —— 产生第二个同时存活的 L0C 缓冲区（溢出），且 ptoas 也会拒绝（不存在合法的 Acc→Acc tmov）。本步骤通过 reuses_input 的 producer 识别就地累加分支，并把另一分支的 seed 重定向到累加器缓冲区，使两个分支共享同一缓冲区、不再产生 move（符合 mad_acc 共享 %dst 的语义）。仅作用于 Acc；重定向是强制的（被拒绝的重定向会触发 INTERNAL_CHECK，绝不退化为 move —— 因为不存在合法的 Acc→Acc move）。它会跳过全局 dead-at-assign 活跃性检查（否则会因 if 之后合法的 phi 消费者而误判拒绝），但仅在验证分支互斥真正所需的两个前提之后：(a) seed 的 producer 是词法上位于该分支内部的 Call（经由分支透传的 if 前值会无条件执行，从而破坏 sibling 就地分支所读取的累加器），以及 (b) 限定分支范围 的 IsTargetDeadAtAssign（在所属 if 处停止）确认分支内 seed 之后没有对累加器缓冲区的尾部读取。任一前提不满足时，该 phi 保持未合并，步骤 7 会明确失败，而不会生成不受支持的 Acc→Acc IR

长期演进方向。 本步骤修复的是手写剥离式 split-K 写法。tile.matmul_acc(acc, lhs, rhs, init_cond=(k == 0)) 用同一块 buffer 表达同一个归约，无需任何修复；AutoTileMatmulL0 现在直接生成该形式（对没有 init_cond 操作数的 bias K-loop 则采用 head-peel），因此已不再有编译器生成的 IR 走到本步骤。保留它是为了作者源码中的剥离写法 —— 仅 pypto-lib 就有 95 处 —— 要移除它必须先完成这些源码的迁移。新代码请优先使用 init_cond；但不要期待本步骤会在近期消失 6. 恒等拷贝缓冲区归一化（NormalizeIdentityCopyBuffersMutator，在步骤 7 前后各运行一次）：在步骤 5 重定向累加器 if-phi 后，对（已被移动的）return_var 的下游裸 Var SSA 恒等拷贝可能仍携带合并前的缓冲区（例如 c_phi 移到 mem_acc_5 后，c: …mem_acc_17 = c_phi）。x = y 拷贝（值为裸 Var 而非 Call）是纯重命名、必须与 y 共用缓冲区，因此本次单向前向遍历把这类拷贝的 LHS 重定型到 RHS 的 MemRef，并替换 LHS 的下游使用。无不一致时为空操作。它在 yield 修复前后各跑一次，因为两侧都会从相反方向让这类重命名滞留。之前：lag2 = lag1 这类携带值重命名要到这里才落到 lag1 的缓冲区上，而步骤 7 必须看到这一点才能让携带值回写之间正确排序——若按归一化之前的缓冲区来看，每个重命名都还在各自的缓冲区上，冲突根本不可见。之后：步骤 7 自身的 IfStmt 修复会把 phi 的 return_var 重指到规范分支缓冲区上，从而以同样的方式让该 phi 的下游重命名滞留。该 mutator 是幂等的，因此没有发现不一致的那一次就是空操作 7. Yield 修复：修复控制流返回变量的 MemRef 不一致。Acc→Acc 拷贝不合法，因此残留的不匹配 Acc 循环携带值会触发内部错误，而不是生成 tile.move： - ForStmt：确保 4 个循环携带变量（initValue、iter_arg、yield value、return_var）共享同一个 MemRef。若 MemRef 不同则在 yield 前插入 tile.move。pl.yield_ 是同时重新绑定所有携带值，但实现它的拷贝是顺序执行的，因此这些拷贝之间要互相排序，而不是按 iter_arg 顺序发射：读取了另一条拷贝所写字节范围的拷贝先执行。缺少这一排序会让移位寄存器塌缩——对 lag2 = lag1; lag1 = v，先写 lag1 的缓冲区会让 lag2 读到新的 lag1，循环携带的变成 lag2 == lag1（#2481）。交换（cur, prev = prev, cur）则不存在任何合法顺序，因此把其中一方先拷进一块新的暂存缓冲区、其回写再从暂存区读取；该分配像其他分配一样被提升到函数体首部。若两个携带值的缓冲区重叠——最常见的是两个携带值由同一个 tile 初始化，从而变成同一块缓冲区——任何顺序都无法挽救，在这里用 CHECK 直接报错，此时携带值还带着名字 - IfStmt：把 return_vars 修补到规范缓冲区（then 分支），并为 yield 值位于其他缓冲区的分支插入一条写入该缓冲区的 tile.move，使经过该 if 的每条路径都写入 phi。分支的 yield 值可能是外层循环的 IterArg——if cond: acc = f(acc) 会让另一分支只剩裸的 yield acc_iter——因此分支用 AsVarLike 而非 As<Var> 匹配：IterArg 有自己的 ObjectKind，漏掉它会让该路径上的 phi 缓冲区从未被写入，而 return_var 仍被修补到兄弟分支的缓冲区上，于是循环携带值被喂入陈旧数据（#2481） 8. 移除冗余 alloc：收集仍被 TileType 变量引用的所有 MemRef，然后移除不再使用的 tile.alloc 语句

复用条件：

  • 生命周期不重叠（无干涉）。当 prev.last_use <= curr.def 时，两个变量不重叠（即源的最后使用可以和目标的定义在同一语句，因为在同一语句内输入先于输出被消费）
  • 相同内存空间
  • 缓冲区大小取其最大成员；由于按最大优先装箱，后纳入的成员都不大于代表，故无需显式字节大小检查（复用方向也不再被限制为「先定义且更大」）
  • 子区间放置只在已知容量的精确 footprint 确认溢出后尝试。候选必须按 backend 对齐策略完整落在已有代表的容量内；每个同时存活或因其它原因不可共享的成员都会阻塞其完整字节范围。
  • No-alias 守护（算子语义）：定义复用变量的算子可以禁止其输出与某些输入操作数共享缓冲区——因为硬件在写输出的同时读取这些输入,原地写会中途破坏该算子。三个来源汇入同一个"每个输出禁止 alias 的输入集合"（ForbidAliasCollector）：
  • not_inplace_safe() —— 该算子无法以 src == dst 运行，因此其输出不得 alias 任何输入操作数。
  • forbid_output_alias(i) —— 该算子对其值操作数 in-place-safe，但在写输出时读取某个特定操作数，因此输出不得 alias 该操作数的缓冲区。
  • 升精度 tile.cast（直接在 ForbidAliasCollector 处理）—— 输出 dtype 比输入更宽时,cast 无法原地：元素 i 在 i*in_bytes 处读、i*out_bytes 处写,写指针超前于读指针,冲掉尚未转换的输入。降精度 / 同宽 cast 仍 in-place-safe（保留下方的跨 dtype 复用）。

MemoryReuse 拒绝将输出放到任一禁止操作数的物理缓冲区上,并通过 reuse-map 合并与 VIEW 继承（reshape/slice 共享其源的 MemRef base）解析每个操作数——因此间接到达的禁止操作数（其 owner tile 被复用到别的缓冲区,或经 view 占用）也能被捕获。对于返回 tuple 的算子，约束会从 tuple 临时变量传播到每个 tile 类型的 TupleGetItem 结果。

当前声明 no-alias 约束的算子：

算子 约束 为何输出不能 alias 输入
tile.recip、tile.rsqrt not_inplace_safe 高精度路径在写输出时读取输入和 tmp scratch
tile.row_sum / row_max / row_min not_inplace_safe TROW* 在写规约输出 [M, 1] 时读取整行输入 + tmp scratch
tile.mrgsort_format1 not_inplace_safe 归并排序 intrinsic 要求 src != dst
tile.move not_inplace_safe TMOV 要求源和目标地址不同；baked-address PTO codegen 会拒绝残留的同地址 move，而不是静默把它当成 no-op
tile.rem、tile.rems、tile.fmod、tile.fmods not_inplace_safe remainder 内核在仍需读取原始被除数时，会先向 dst 写入中间商或余数；若 dst 与该输入复用，最终符号修正或减法会读取已被覆盖的数据。此外，A2/A3 的 tile.rem / tile.rems codegen 会拒绝 scratch tmp MemRef 与任一仍存活的 tile source（被除数，以及 tile.rem 的除数）重叠，包括重叠 subview；A5 不消费该 scratch 操作数
tile.transpose not_inplace_safe pto.ttrans 非 in-place 安全：a2a3 非对齐标量路径直接从 src 写 dst（不经 tmp 暂存），dst == src 会边写边读损坏数据。输出始终分配新 buffer（InitMemRef 也不会为其继承输入的 buffer）。
tile.sel forbid_output_alias(0)（mask）、(3)（tmp） TSEL 在写 dst 时读取 mask + tmp scratch
tile.sels 感知 target TSELS 始终要求 dst 与 predicate mask 分离，并允许复用 src 或 tmp；A2/A3 会先将 scalar 写入 tmp，再通过 set_cmpmask 读取它，之后才写 dst，因此 tmp 可以 alias dst，但不得与 mask/src 重叠；A5 保留但不读取 ABI 中的 tmp，允许其 alias 任一操作数
tile.prelu 感知 target A2/A3 的 TPRELU 在写 dst 时读取 src、slope 与 tmp，因此是 not_inplace_safe；A5 保留 ABI 要求的 tmp 操作数但不读取它，所以 dst 可复用 tmp，但不可复用仍参与运算的 src/slope
tile.{row,col}_expand{,_mul,_add,_sub,_div} forbid_output_alias(1)（广播向量） 行/列向量（arg 1）会被每个输出行/列重读,输出若 alias 它则在第一行/列后被覆盖
tile.cast（仅升精度） 输出 ≠ 输入缓冲区（条件式,在 ForbidAliasCollector） 更宽的输出写指针超前于读指针（见上）
  • 流水线 stage 守卫（容量门控，仅针对复制路径）：pl.pipeline(stage=F) 将循环体复制 F 份以实现 ping-pong，LowerPipelineLoops 给每个副本产生 tile 的 Call 打上 pipeline_membership (group, stage)（见 30-lower_pipeline_loops.md）。本守卫只作用于这条复制路径。在 memory_planner=PTOAS 下，合格循环会先被 LowerPipelineToSlots 接手：它给每个 load 一个以 iv % F 索引的 pl.MemRef(..., slots=F)，循环体只有一份——没有副本，也没有 pipeline_membership，而且该流水线本就跳过 MemoryReuse。以下描述仅适用于被那个 pass 拒绝的循环。F 份副本在调度器下并发执行，因此它们程序序不相交的生命周期不是安全的复用信号——把并发副本合并到同一块缓冲会注入一条虚假的写后读（write-after-read），使各 stage 串行化（即 #1475 的 cube matmul 操作数坍缩）。MemoryReuse 因此在每个内存空间（包括 L0 matmul 空间 Left/Right/Acc/Bias，且无论 tile 是 load 还是 tile.move 的结果）都把并发副本保持在不同的缓冲中，最多到可负担的双缓冲深度 F_g = min(depth_g, ⌊C_s / slot_g⌋)：stage k 的副本落在残数 ordinal(k) mod F_g（稠密的 stage 序号，因此稀疏 stage ID 如 {0, 2} 不会因 2 mod 2 == 0 mod 2 而错误合并），因此并发副本永不共享（放得下时是完整 ping-pong，空间紧张时尽量分散）。分离是否放得下由精确的按空间分配器足迹（SpaceFootprint，与 AllocateMemoryAddr 共享——按构造保证一致）决定，而非估算。当某空间在所有 group 的可负担深度下仍然溢出时，采用优雅的跨 group 削减：将某个 group 的深度降低一个残数并重新打包（按 MaxRelief 启发式选择 group——优先释放最多字节，平局取最小 group id）；若削减耗尽，则整体回退到 legacy 重新打包（force_legacy），从而绝不会在 legacy 打包本可放下的情况下溢出。容量未知的空间（未配置 backend）使用 legacy 判据，因此容量门控路径绝不比 legacy 更差。当门控把某个 group 的深度降到其请求的 stage= 之下（或某空间触发 legacy 回退）时，MemoryReuse 通过统一诊断通道发出诊断——一条 PH-MR-001 性能提示（回退情形则为 warning），指出请求深度与实际深度以及修复方式（把每 stage 的 tile 缩小到 ≤ C_s / stage，或把 stage= 降到能放下的值）——因此容量导致的串行化绝不会静默发生。若单个物理 buffer 自身已经超过整个内存空间，则不发出该 fallback warning：这不是复用或 stage 深度退化，而是不可能完成的分配，应由面向具体算子的检查或 AllocateMemoryAddr 作为硬错误报告。复用决策完成后，MemoryReuse 会剥离已消费的 pipeline_membership attr，使其不会带到下游 pass 或 codegen。
  • DSA_RP 跳过 MemoryReuse；它把请求的流水线深度表示为硬分离，并在 AllocateMemoryAddr 中执行仅放宽流水线意图的回退。

不再有 shape / dtype / TileView 兼容性门槛：共享同一物理 MemRef 的 tile 可以携带不同的 shape、dtype 或 TileView 属性。PTO codegen 为每个 tile 绑定一条 per-variable 的 alloc_tile，因此每个别名都以各自的静态 shape / dtype / layout / valid_shape 声明共享基址。这允许例如：

  • 跨 dtype 复用 —— BF16 tile 复用已死亡的 FP32 tile 的缓冲区（例如跨 tile.cast）；
  • tile.fillpad 输出复用其输入，以及两个 pad 不同的 fillpad 输出共享一个缓冲区；
  • N-D tile 在 valid_shape 不同的情况下共享缓冲区（各自在自己的 alloc_tile 上保留各自的 valid_shape）；
  • L0 cube 输入 Left / Right 中 shape 不同的子 tile 共享同一槽位（例如 fused-attention QK 的 Right [k, SEQ] 被 PV 的 Right [k', HEAD] 复用，将 L0B 峰值减半 —— issue #1595）。

早期版本以 AreTileTypesCompatible（shape / dtype / view 匹配，外加一个狭窄的 L0 字节复用例外）作为门槛；该门槛已移除。对读-写同体（read-while-write）算子的正确性现由上面的 no-alias 守护精确处理，而不再依赖粗粒度的整块匹配。

Alloc 清理：

MemRef 共享完成后，部分 MemRef 对象变为无引用状态（其变量现在指向不同的共享 MemRef）。该 Pass 遍历周围的 SeqStmts，移除所有左值 MemRef 指针不在仍使用集合中的 tile.alloc AssignStmt。

声明式分配

复用是机会主义的：任意两个生命周期不重叠的 tile 都是合并到同一块 buffer 的候选。就容量而言 这是正确的默认策略，但它并非没有代价——共用一块 buffer 的两个 tile 之间被强加了一条源码从未 要求的 WAR 依赖，硬件因此必须串行执行本可由调度器重叠的工作。

在 tile 注解中引用一个声明好的 pl.MemRef("name")，作者即可把某块分配从 packer 手里收回。 InitMemRef 将其物化为 tile.alloc(..., pinned=True)（见 InitMemRef），本 pass 随后视其为封闭的：pinned 区间在 first-fit 打包中自开一个槽位，之后每个候选在扫描槽位时都会跳过它。（隔离是打包循环内的 per-slot 标记，而非又一个 can_share 门——can_share 是 O(M²) 打包的最内层，且每次 shed 都会重跑，所以该判定按区间解析一次，而不是按配对解析。）具体而言：

  • 作者绑定到不同 buffer 的 tile 永远不会被合并，无论其生命周期多么不相交——保持彼此独立 正是绑定的目的。
  • 作者绑定到同一 buffer 的 tile 在 InitMemRef 阶段已共享同一 base，并保持如此。
  • 未绑定的 tile 照旧参与打包，且绝不会被塞进声明式分配。

代价由作者承担：pin 是用容量换并行度，pin 过头的 kernel 会在 AllocateMemoryAddr 处硬报错， 而不是被静默合并回去。

重叠检查。 独立绑定到同一槽位的两个 tile 不得同时存活——那不是复用，而是后写的数据 摧毁了先前 tile 仍需要的内容。该检查按槽位而非按分配进行：pl.MemRef(slots=N) 声明中位于 不同槽位的两个 tile 本来就应该同时存活（这正是 ping-pong 的用途），只有落在同一槽位的 tile 才会互相破坏。下标是运行期表达式时（l0c[i % 2]）没有静态的槽位归属可比，该检查被跳过——轮转 的正确性由作者负责——而与其他分配之间的隔离依然成立。该检查由本 pass 负责，因为这里才有生命周期信息 （ComputeLifetimes）；判据与 packer 自己的 var_overlap 一致，因此首尾相接是允许的 （一个 tile 的最后一次读，可以正是产出下一个成员的那条语句）。通过继承而非绑定落到该 buffer 上的 tile——view、原地结果、裸 SSA 别名——被排除在外：它们与源 tile 是同一份数据，重叠是预期 行为。

由于隔离保证由本 pass 提供，而 ptoas 会整体替换本 pass，因此 memory_planner=PTOAS 下的 单参数 pl.MemRef(...) 会在 InitMemRef 处直接报错，而不是"分配了但不隔离"地静默生效——只分配不 隔离，恰好会把作者写这条绑定所要避免的合并原样还回去。 memory_planner=DSA_RP 同样会跳过本 pass，但会在分配问题中保留该契约： 每个声明式分配都会与同一内存空间中的其他所有分配建立硬分离；若同一声明中独立绑定的 成员同时存活，则在求解前直接报错。

Ascend910B load + tpop_from_aic 危害

在 SplitMode 非 None 的 Ascend910B AIV 函数中，如果某个 writer 同时消费 tile.load 的结果（或其合法 view 派生）和 tile.tpop_from_aic 的值，则它的输出不能与该 load 结果落在同一块物理 buffer 上。让 writer 的输出原地复用 load buffer 会在该硬件上产生静默的错误结果。

MemoryReuse 掌管所有 buffer 合并决策，因此它从源头上阻止这种危害共享的形成，而不依赖后续的拆分。当 guard 生效时，复用决策在以下条件同时满足时被阻止：

  • writer 的定义 op 消费了 tile.tpop_from_aic 的值，且
  • 它将要原地复用的那个 buffer 成员（其 last use 正是该 writer 的定义语句）是 load 派生的。

两种分类都以 Var 标识（identity）为键。读取操作数（operand）时还需多一步：值可能经由循环携带（loop carry）到达 writer，而 IterArg 本身从来不是 AssignStmt 的定义，因此 Var 标识永远无法对其分类。由于 MaterializeSemanticAliases 已经把每条 carry 链（init value、IterArg、yield value）融合到同一个 MemRef base 上，IterArg 操作数改为按该 base 的污染（taint）状态来分类。而让这条 carry 首先变得可见的前提，是用 AsVarLike 读取操作数（绝不能用 As<Var>，它不匹配 IterArg 自己的 ObjectKind）；否则 down_next = tile.add(down_prev, pipe_carry) 中被携带的 tpop 值会静默丢失污染标记，从而形成危害性的原地复用。

carry 同时也打破了程序序，因此收集器会对函数体遍历两次。污染某条 carry buffer 的 producer 可能位于它所污染的那个 use 之后——例如 w = tile.add(l, carry); p = tile.tpop_from_aic(); yield p，从第 1 次迭代起 w 就读到了 tpop 值，但单次前向遍历在见到 p 之前就已经对 w 完成了分类。第三次遍历不会带来任何新信息：第一次遍历后 base 集合即已完备（tile.tpop_from_aic 的定义与顺序无关，而 view 与其 source 共享 base），因此第二次遍历中只有 Var 集合会增长，且它们在该次遍历内部按程序序传播。两次遍历均为 O(N)。

该 guard 由 BackendHandler::RequiresSplitLoadTpopWorkaround()（仅 Ascend910B 为 true）以及函数为 split-AIV 这两个条件门控；在其他任何 backend / 函数类型下输入集合为空，复用行为不变。writer 仍可自由复用任何非 load buffer —— 只有 load + tpop 的原地组合会被拒绝。（该 guard 此前由独立的 LegalizePTOBufferReuse pass 在事后拆分 buffer 来实现，现已并入 MemoryReuse。）

示例

MemRef 共享与 Alloc 清理

之前（InitMemRef 之后）：

# SeqStmts [
mem_vec_0: MemRefType = tile.alloc(Vec, -1, 16384, 0)
mem_vec_1: MemRefType = tile.alloc(Vec, -1, 16384, 1)
mem_vec_2: MemRefType = tile.alloc(Vec, -1, 16384, 2)
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.add(tile_a, ...)
# tile_a last use ↑
tile_c: Tile[[64, 64], FP32, memref=mem_vec_2] = tile.load(...)
# ]

之后（tile_c 复用了 tile_a 的 mem_vec_0，mem_vec_2 的 alloc 被移除）：

# SeqStmts [
mem_vec_0: MemRefType = tile.alloc(Vec, -1, 16384, 0)
mem_vec_1: MemRefType = tile.alloc(Vec, -1, 16384, 1)
# mem_vec_2 alloc removed — no longer referenced
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.add(tile_a, ...)
tile_c: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
# tile_c now shares mem_vec_0 with tile_a
# ]

生产者-消费者复用

当某个变量的最后一次使用恰好发生在定义另一个新变量的同一条语句上（生产者-消费者关系）时， 新变量可以复用旧变量的内存，因为输入在输出写入之前被读取：

# Before:
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.create(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.muls(tile_a, 0.0)
# tile_a.last_use == tile_b.def → reuse allowed

# After:
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.create(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.muls(tile_a, 0.0)
# tile_b reuses mem_vec_0

生命周期重叠（不可复用）

之前/之后（无变化——alloc 语句保留）：

# SeqStmts [
mem_vec_0: MemRefType = tile.alloc(Vec, -1, 16384, 0)
mem_vec_1: MemRefType = tile.alloc(Vec, -1, 16384, 1)
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.load(...)
tile_c: Tile[[64, 64], FP32, memref=...] = tile.add(tile_a, tile_b)
# tile_a and tile_b are both live here → cannot reuse
# ]

实现

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

Pass MemoryReuse();

实现文件：src/ir/transforms/memory_reuse_pass.cpp

  • LifetimeAnalyzer 遍历完整 IR 树计算变量生命周期（包括嵌套控制流）
  • ComputeLifetimes 构建 MemRef 共享组和生命周期区间
  • IdentifyReuseOpportunities 查找复用候选，在已知容量溢出时重试对齐子区间放置，并记录每个成员的 placement offset
  • ApplyMemRefSharing 通过 MemRefSharingMutator 更新 MemRef 指针
  • TopDownRetargeter::CoalesceAccumulatorIfPhis 通过把失效分支的 seed 重定向到就地累加器缓冲区，合并被剥离的循环携带累加器 if-phi，使 YieldFixupMutator 不再产生非法的 acc→acc tile.move（见算法步骤 5）
  • YieldFixupMutator 修复 ForStmt/IfStmt 在复用后的 yield/return_var MemRef 不一致（合法时插入 tile.move；拒绝残留的 Acc→Acc 不一致）
  • NormalizeIdentityCopyBuffersMutator 协调累加器 if-phi 合并后 LHS/RHS 缓冲区不一致的裸 Var SSA 恒等拷贝（见算法步骤 6）
  • UsedMemRefCollector 收集共享后仍被引用的 MemRef 指针
  • RemoveUnusedAllocStatements 从 SeqStmts 中过滤掉冗余的 tile.alloc 语句

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

passes.def("memory_reuse", &pass::MemoryReuse, "Memory reuse optimization");

测试：tests/ut/ir/transforms/test_memory_reuse.py

  • 测试非重叠生命周期的 MemRef 共享复用
  • 测试重叠生命周期不复用
  • 测试内存空间隔离
  • 测试字节大小兼容性
  • 测试跨 dtype / 跨 TileView 复用（现已允许：BF16↔FP32、fillpad 输出↔输入、valid_shape 不同）
  • 测试 no-alias 守护（TestForbidOutputAlias + TestInplaceOps），上表每条约束一个用例：
  • tile.recip / tile.rsqrt / tile.row_sum —— 输出不得 alias 输入（not_inplace_safe）
  • tile.sel —— 输出不得 alias mask / tmp（forbid_output_alias）
  • tile.sels —— 输出始终不得 alias mask；A2/A3 与 A5 均允许 tmp/输出 alias，但 A2/A3 backend 仍会拒绝 tmp 与 mask/src 重叠
  • tile.prelu —— A2/A3 输出不得 alias 任一输入；A5 输出仅可 alias 未使用的 tmp
  • tile.col_expand_mul —— 输出不得 alias 广播向量
  • 升精度 tile.cast —— 输出不得 alias（更窄的）输入
  • 经 VIEW 间接到达的禁止操作数也被遵守（物理缓冲区解析）
  • 测试切片操作的 MemRef 共享保持
  • 测试冗余 alloc 语句移除
  • 测试控制流生命周期分析（ForStmt 内嵌套 IfStmt、分支变量共享）