pnnx for torch exported program

[lvy010]

PR：#6945

本项目为 pnnx 补齐了 torch.export exported program .pt2 的直接转换支持。实现思路是契约优先、失败即拒：前端只转换它能证明的东西，凡是证明不了的一律带着明确诊断信息拒绝，而不是近似转换出一个看起来能跑、其实算错的模型。

实现

项目要求	实现
pnnx 支持 torch dynamo exported program 直接转换 pnnx model.pt2	新增 PT2 前端：model_format + json + exported_program{,_reader} + exported_program_defaults + load_exported_program；共 11 个新文件约 5370 行
为 pnnx 单元测试添加 pt2 格式支持	CTest 注册 405 项 pt2，与 390 项 torchscript 并行成对；顶层算子测试双路径与 eager 对比
解析 pt2 中的模型图和权重	JSON 图 + constants/weights 原始张量载荷；sizes / strides / storage_offset / 共享存储视图 / 零 stride 展开全部在越界与溢出检查下密集物化
不引入第三方库	ZIP 复用自带 StoreZipReader，JSON 用新写的 pnnx_json，零第三方依赖；算子下降完全复用既有 pass_level2~pass_level5 与 pass_ncnn
通过多平台编译和测试	pnnx-pt2.yml：ubuntu / windows / macos × producer torch 2.12.1 / 2.13.0，converter 统一链接 LibTorch 2.12.1
本地全链路验证	前端 14/14、后端 ncnn 数值 1/1、torchscript 390/390、pt2 算子 390/390

方案：

torch.export(model, args) ──torch.export.save──> model.pt2 (zip)
                                                     │
                                    detect_model_format 判定格式
                                                     │
              ┌──────────────────────────────────────┴───────────────────────┐
              │ exported_program_reader：STORE 条目 + 逐条 CRC32 + ZIP64 边界 │
              │ pnnx_json → 类型化 schema 8.20 → 默认值/效应归一化            │
              └──────────────────────────────────────┬───────────────────────┘
                                                     │
                                        load_exported_program
                                                     │
                                    PNNX IR
                                                     │
                                   pass_level2~5 + pass_ncnn
                                                     │
                            .pnnx.param/.bin + _pnnx.py ／ .ncnn.param/.bin

前端按职责分模块：model_format 只做格式判定；exported_program_reader 负责归档与原始载荷；exported_program{.h,.cpp} 把图、签名、range constraints、tensor metadata 解成类型化结构，字段畸形直接变成解析错误而不是留到下游；exported_program_defaults 集中做 schema 归一化、默认参数补全与局部效应归一化；load_exported_program 只负责建 IR。由于与 TorchScript 路径共用 pass 管线，绝大多数算子零成本复用。

示例

import torch

net = Model().eval()
x = torch.rand(1, 3, 224, 224)

program = torch.export.export(net, (x,))
torch.export.save(program, "model.pt2")

# 与 torchscript 路径完全相同的命令行
pnnx model.pt2
pnnx model.pt2 inputshape=[1,3,224,224]

生成 model_pnnx.py、model.pnnx.param/.bin，ncnn 后端另有 model_ncnn.py 与 .ncnn.param/.bin。单元测试走共用 helper，一次同时验证两种格式：

from pnnx_test_utils import test_model_formats
ok = test_model_formats(net, (x,), expected, "test_F_softmax")

设计

1. 版本契约是三个独立条件，不是一个 PyTorch 版本号

archive_version="0"、ExportedProgram schema 恰好 8.20、opset_version.aten 等于所链接 LibTorch 的 torch::jit::getMaxOperatorVersion()——三者分别校验。这个 operator version 既不是 PyTorch 发行号也不是 ONNX opset，且没有实现任何升降级适配。所以 producer 与 converter 是两个独立选择项，CI 里用两个 venv 明确分开，避免版本号相近就假定兼容。

2. 归档读取按不可信输入处理

被消费的条目只接受 STORE，且每一条都验 CRC32；校验 ZIP/ZIP64 物理边界、目录计数、local/central header 一致性；拒绝加密与重名条目。新增 get_file_compression()/get_file_flags()，让 reader 在分配载荷之前先检查记录声明。没被消费的压缩附件可以留在索引里，但绝不读取。

同时设了明确的信任上限：原始张量合计 512 MiB、单条记录 512 MiB、单个密集 attribute 512 MiB、ZIP 中央目录 256 MiB、条目数 1Mi、条目名累计 64 MiB。这些是信任边界，不是进程总内存上限。

3. in-place 算子的所有权证明

torch.export 会把 relu_、add_、masked_fill_ 这类原地写保留在图里。要安全地把它们改写成函数式版本，必须先证明被写的张量是完整、单次使用、无别名的本地临时量，否则就会悄悄改掉别人还在用的数据。

实现上用的是显式 dispatcher overload 配对表，并做全程序使用计数。外部输入、参数、buffer、常量、视图、已存在的别名——全部 fail-closed。

这里有一个必须显式白名单的原因：schema 上没有 alias 标注并不足以证明是新分配。aten::_unsafe_view 就是单返回、无 alias 标注、无参数写入，但它返回的是视图。所以白名单必须靠人工判断这个算子是否真的写进新缓冲区，测试里专门钉了一条 pure does not mean whitelisted allocator 的负向用例。

按这个标准，卷积 / matmul / bmm / linear / batch_norm / layer_norm / group_norm 都是单返回、无 alias 标注、无参数写入，且语义上确实物化新输出缓冲，因此可以作为本地根。补上它们之后，nn.ReLU(inplace=True) 接在 conv/BN 之后这个最常见的写法就能正常转换了。

4. 输入契约保留到生成代码里

导入时把 dtype、rank、静态维度、裸符号、range、共享符号相等关系保存在 pnnx.Input 的 __pt2_input_* 保留参数上，ir.cpp 据此在生成的 Python 里、计算之前发出显式 ValueError 检查，并且在 python -O、-OO 下依然生效。

需要说明清楚：裸 ncnn 模型不执行这些检查，且这些 PT2 参数不会出现在 ncnn 产物里。

5. guard 与工厂算子 dtype

只接受已证明为真的具体布尔 aten._assert_scalar，或静态可验证的 aten._assert_tensor_metadata，并且在移除 guard 前再交叉核对操作数的 dtype/shape。无法解析或依赖运行期的 guard 一律拒绝。

工厂算子只为 aten.full.default 和 aten.full_like.default 从输出 metadata 补回被省略的 dtype——不推断、不强转 fill_value、不覆盖已显式指定的 dtype，目的是保住 producer 记录的 dtype，而不是依赖生成代码运行时的默认 dtype。

6. CLI 提前拒绝不支持的意图

PT2 在创建任何输出文件之前就拒绝非空 customop/moduleop 以及 device != cpu。inputshape 的 dtype 后缀必须与导出 dtype 一致，省略后缀则沿用导出 dtype；input/input2 的 npy 会被完整读取校验，但不执行、也不用于常量特化。另外把 save/python/save_ncnn 的失败改为向上传播，不再静默继续。

多平台 CI

新增 .github/workflows/pnnx-pt2.yml：

Job 维度	取值
平台	ubuntu-latest / windows-latest / macos-latest
producer torch	2.12.1 / 2.13.0
converter	统一链接 LibTorch 2.12.1
每个 job 运行	pt2_frontend 全量 + torchscript/pt2 的 ReLU/Linear/Conv2d 冒烟
仅 Linux 追加	pt2_backend

日志与 CMakeCache 作为 artifact 上传。需要如实说明：该 workflow 不跑完整 pt2_operator 套件，也不跑 Transformers 模型矩阵；workflow 配置本身不等于这些 job 已经成功执行过。

验证

本地 Windows x64 实跑，工作树干净：

环境项	取值
平台 / 编译器	Windows x64，VS2019 Build Tools，Ninja，Release
producer	Python 3.13.3，Torch 2.12.1+cpu，ExportedProgram schema (8, 20)
converter	LibTorch 2.12.1
依赖	TorchVision 0.27.1+cpu、Transformers 4.51.3、NumPy 2.5.1、ncnn 1.0.20260526
pnnx.exe SHA-256	41450B2ABFF0FB8B0E47EA1C5961347939119ABDEED31438BA5C3F50B238264C

使用 --output-on-failure --no-tests=error：

选择	注册数
-L '^pt2_frontend$' -j 4	14
-L '^pt2_backend$' -j 1	1
-L '^torchscript$' -j 8	390
-L '^pt2_operator$' -j 8	390

torchaudio 没有与 torch 2.12.1 匹配的发行版，所以 4 个 test_torchaudio_* 在两条链路上都因 import 失败，这正是 torchscript 那 4 个失败的全部来源。以本机实际可运行的 386 个用例计：torchscript 386/386，pt2 377/386。

端到端

模型 Conv2d + BatchNorm2d + ReLU + AdaptiveAvgPool2d + Linear，torch.export.save 导出后执行 pnnx model.pt2 inputshape=[1,3,32,32]：

• .pt2 归档条目全为 STORE；
• ncnn 图降到 4 层、0 个 ATen 残留——BatchNorm 折进 Convolution 且 ReLU 融合，其后 Pooling、InnerProduct；
• 与 eager 的最大绝对误差：生成的 PNNX Python 1.49e-08，真实 ncnn.Net 7.55e-05。

两个编译期探测为负，限定了本次结论范围

PNNX_TORCH_HAS_PICKLE_LOAD 探测失败 → legacy 张量字典路径未参与编译，test_pnnx_exported_program_legacy 未注册；protobuf 缺失 → 本次构建无 ONNX 输出可比。这两点在结果文档里单独记录，不混进通过数。

尝试与限制

有 9 个 pt2_operator 用例失败在转换器内部。每一个都单独归因到自己的诊断信息，且都落在本契约已声明的限制范围内：

用例	原因
test_mobilenet_v3_small_pt2	dropout_ 未登记 in-place 配对：aten::dropout 把首参 self 改名为 input，当前按 schema 归位的机制表达不了这种重命名
test_swin_t_pt2	fill_ 写入的目标不是单次使用的无别名本地量
test_vit_b_32_pt2	载荷超过 512 MiB 合计预算
test_transformers_flaubert_attention_pt2、..._xlm_...	masked_fill_ 目标非单次使用本地量
test_transformers_mt5_attention_pt2、..._t5_...	add_ 目标非单次使用本地量
test_transformers_funnel_attention_pt2	aten._assert_scalar 不是具体布尔值，运行期 guard 未实现
test_transformers_longformer_attention_pt2	producer 把整数标量序列化进了 aten.div.Tensor_mode，而该 overload 的 schema 声明 self 为 Tensor

其中 4 个 in-place 用例需要的是通用 functionalization / alias-update 下降，而不是把白名单再放宽——当前的所有权证明选择拒绝，而不是去改写一个仍被共享的张量。

其余按设计尚未实现的部分：

• 被消费条目的 DEFLATE 解码；
• 受限 pickle 解码：legacy 路径使用 LibTorch torch::pickle_load，仅接受可信输入，既不是受限 pickle 实现也不是沙箱；
• 完整 PyTree 输入/输出重建：目前边界是扁平位置张量入参 + 空 kwargs，输出为张量/数值标量或其非空扁平 tuple；
• 通用动态 shape：symbolic 维度由显式 inputshape 固定化，不向 ncnn 发出运行期动态 shape 断言；
• 有 4 个用例故意期望转换器拒绝并钉死诊断字符串：Tensor_fill、Tensor_slice_copy、Tensor_index_put、torch_masked_select；它们的 TorchScript 变体保持正常的转换与数值比对。这类负向测试通过，只代表观察到了预期诊断且没有崩溃/超时，不代表该算子转换成功。

[lvy010]

flowchart LR
    A[model.pt / model.pt2] --> B{检查 ZIP 成员}
    B -->|TorchScript| C[load_torchscript]
    B -->|PT2 Archive| D[load_exported_program]
    D --> E[读取 archive/schema]
    E --> F[解析 JSON graph/signature]
    E --> G[读取 weights/constants]
    F --> H[构造 pnnx::Graph]
    G --> H
    C --> H
    H --> I[level2 到 level5]
    I --> J[pnnx / ncnn 输出]

Loading