pnnx for torch exported program - #6953
magician336 wants to merge 69 commits into
Conversation
- pt2_sweep.py: 全量测试扫描, pt2/ts 双路径 .ncnn.param 结构对拍, 三分类 (PASS/DIFF/UNSUPPORTED_OP), 当前基线 PASS 117 / DIFF 8 / UNSUPPORTED_OP 90 - pt2_crosscheck.py: 算子级对拍矩阵 - testutil_pt2.py + test_pt2_smoke.py: pyncnn 数值对拍 (ctest: pt2_smoke) - test_pt2_weights.py: 权重链路验证 注: pt2_crosscheck.py 位于 tests/ncnn/ 下, 非 tools/pnnx/ 根 (CLAUDE.md 旧引用有误)
json.hpp / load_pt2.cpp(h) / load_pt2_parse.cpp(h) + main.cpp/storezip 探测分发。 作废原因 (2026-08-30 实测): torch 2.13 的 .pt2 为 JSON 图 (schema 8.20) + ZIP_STORED 原始权重, 纯 C++ 可解析, "pickle + deflate 需 Python 预处理" 的 架构前提不成立; 竞品 PR Tencent#6933 已以纯 C++ 方案占位并系统性覆盖本方案的 参数兜底/形态判别等全部设计点。 保留价值: kPt2PreprocessPy 内嵌的 aten->pnnx 语义映射与 11 条踩坑记录, 可转为对 PR Tencent#6933 的评审材料; 验证基建见上一 commit。
- 删除 load_pt2.cpp/h、load_pt2_parse.cpp/h(桥接前端,格式前提已证伪) - main.cpp / CMakeLists.txt 恢复 master(c189d88)基线 - 保留 json.hpp(自写 JSON 解析器,独立重建将复用) - 保留 storezip central directory 读取增强(通用基础设施) - 验证基建(dcedd54)与格式逆向(docs/11)原样保留 新方向:独立重建 pt2 前端,PR Tencent#6933 仅作思路参考(不抄代码)。
复用 pt2_crosscheck 的 normalize_param,直接比对已落盘的 sw_*.ncnn.param 与 sw_*_ts.ncnn.param,无需重跑 pnnx。
- pt2_schema.h/cpp:L4 schema 层,解析 model.json(header/nodes/signature) 与 weights/constants config;节点参数按 as_* 变体建模(12 种,实测 237 个 .pt2 全量普查),kind=1/2 区分 positional/keyword 实参 - load_pt2.h/cpp + main.cpp:model_file_maybe_pt2 探测(zip 签名 + <root>/models/model.json 特征),插在 torchscript 分发之前; 忠实转写 builder 为 N2 留桩 - 纯 C++ 零 libtorch,无条件编入 pnnx(pnnx_SRCS) - json.hpp 边界加固:自包含 <cstdio>、未闭合字符串报错、数字扫描收紧 (前导零/小数点/指数校验)、\u 代理对展开为 4 字节 UTF-8 - tests/test_json.cpp:70 项边界单测(独立 harness); tests/test_pt2_schema.cpp:schema dump harness(-c 规范化输出) - 验收:239 个真实 .pt2 双侧规范化对拍全一致(C++ harness vs pt2-dump/dump_canonical.py);WSL 构建通过,pnnx mini.pt2 正确命中 pt2 分发
- load_pt2 builder:忠实转写,与 ts level1 通用转写形态对齐—— 权重/buffer/tensor_constant → pnnx.Attribute(op 名 = state_dict 名, 裸字节从 zip 读入);标量参数 → prim::Constant operand;张量列表 → prim::ListConstruct;Input/Output 命名与 ts 相同;op 保持 aten 原名 (target 去 torch.ops. 前缀与 overload 后缀);零归一化 - pass_level2/F_pt2.cpp:PT2 形态分支,把 torch.export 的缺省实参形态 归一成 ts 等价形态——flatten 2-op 补 end_dim=-1 汇合 torch_flatten; conv2d 四种缺省/全参变体 emit params 化 F.conv2d(缺省 dilation/groups 内联,N3 换静态表)。priority 55:先于 torch_flatten(60)归一形态, 先于 F_conv2d_1(140)避免其 7-input pattern 误吃 - schema:补 constants 路径(tensor_N 位于 data/constants/)与 zippath - 验收(WSL,同权重模型):conv2d 模型 ts/pt2 双路径 .ncnn.param 与 .ncnn.bin 全部逐字节一致(权重字节链路端到端打通);smoke 模型 (flatten+cat+relu)param/bin 逐字节一致;aten::add 2-op 与 cat 列表 形态由现有 fuse_expression/torch_cat 直接消化,零新 pass
- scripts/dump_aten_defaults.py:离线 dump torch._C._jit_get_all_schemas, --scan 语料驱动按需收录,180 算子入库(可审计、可再生成、零运行时依赖) - src/aten_defaults_table.h:生成的静态表(header-only,按 overload 全名查表) - load_pt2:节点缺参查表补全为完整 schema 形态(带 fill default 标记), provided 形参名与表不符时回退原样转写;F_pt2.cpp 的 5 条内联缺省分支删除, pt2 图与 ts 图同构后由既有 ts 形态分支零改动消费 - pt2_schema:解析 graph.tensor_values 张量元数据表(含中间张量形状/dtype, 修正"JSON 不携带中间张量形状"的误记);builder 填 operand.shape, 文件事实优先于 CLI inputshape - 空列表默认值编为 type 0(None)对齐 ts 形态(max_pool2d stride=() 段错误根因, 修复后 sweep PNNX_PT2_FAIL 120 -> 6)
- pt2_schema: 新增 json_as_int 助手替换全部裸 asInt(),整数字段被未来 torch 版本写成浮点时不再经 JSON_INT/DOUBLE union 静默错值; schema_version major/minor 带 isNumber 守卫,保持缺失=-1 语义 - load_pt2: 删除与 3.5 节 tensor_values 回填矛盾的过期注释(N3 漏改) - dump_aten_defaults: 空列表编码注释改为实际机制说明(编 INTS 空串, builder 转 None 对齐 ts 形态) - 验证:WSL 重建后 test_pt2_schema mini/smoke_cat/mega 解析 OK, 全量 sweep 复跑 PASS 115 / DIFF 95 / PNNX_PT2_FAIL 6 / 219 与基线 逐项一致(零回归)
…S 115→178 builder(load_pt2): - hoist_constants:标量 prim::Constant 前移到消费者之前。builder 惰性创建 的常量滞后于消费者,fuse_expression(level3)反向扫描会先把常量包成 pnnx.Expression,算术链融合无法内联字面量,产出与 ts 不同的常量 blob 形态(激活族等 36 个 DIFF 的根因) - 切分族(unbind/split/split_with_sizes/chunk/tensor_split)转写为 1 输出 + prim::ListUnpack,对齐 ts level1 形态,复用现有 torch_* 形态 分支与 fuse_op1ton_unpack(level3)折叠(getitem 已被 exporter 折叠成 as_tensors 多输出,1 输出 pattern 无法匹配) pass_level2 PT2 形态分支: - F_conv1d_1/F_conv3d_1/F_conv_transpose1d_1/2d_1/3d_1:torch.export 产出 公开算子 aten::conv1d/conv3d/conv_transpose*(与 ts 内部算子 aten::_convolution 不同),参数实参经 fuse_constant_expression 折入 params,后续 fuse_static_conv* 正常折叠权重(F_conv2d_1 同款先例) - F_pt2_weight_norm:aten::_weight_norm(v,g,dim=0) 折成 pnnx.Attribute, 复用 utils.cpp apply_weight_norm(与 ts level1 同一 float 实现,权重 字节一致);v/g 非常量等形态不匹配显式失败 - F_pt2_adaptive_pool_*:adaptive pool output_size 中被 torch.export 实例 化的 None(=输入空间维)还原为 0(pass_ncnn 对 0 写 -233 哨兵,恒等池化 语义不变);仅匹配 prim::Constant 形态,替换图同构靠'确会改写才匹配' 终止重写循环 验证:全量 sweep PASS 178 / DIFF 32 / PNNX_PT2_FAIL 6 / 219(基线 115/95/6),0 回归 63 改善;F_relu/stack 等代表性场景 param+bin 逐字节 一致;clang-format 10.0.1 已跑(astyle 同前,提 PR 前补验)
builder 解析 node.metadata.nn_module_stack(export 图一等事实),对白名单 (模块类, aten 算子)对转写为 nn.<类> + params 化形态——折参规则逐模块对齐 ts 侧 level1 模块转换(FuseModulePass)的产出,与 F. 形态对齐 level1 通用 转写同理,loader 侧转写与 ts 层次对等。 - 白名单:ReLU6/Softmax2d/ChannelShuffle/PixelShuffle/MaxPool1d-3d(含 with_indices)/AdaptiveAvgPool1d-3d/pad 族 9 类/Upsample 三类 - 折参改名对齐 level1:pad→padding、output_size→size、scale_factors→ scale_factor;折参集合精确对齐(pnnx 模式匹配对参数集合敏感) - 单值 int/float 实参按算子空间维广播成列表(对齐 JIT 的 schema 泛化) - export 省略的默认实参按 L7 表补进 params(MaxPool 的 dilation/ceil_mode) - F_pt2 新增 nn.AdaptiveAvgPool* params 形态的 None 实例化还原分支 - MaxPool with_indices → nn.MaxPool* + return_indices=True,indices 由 eliminate_maxpool_indices 消除后与 ts 形态汇合 sweep:PASS 197 / DIFF 13 / PNNX_PT2_FAIL 6 / 219(基线 178/32/6)
…7→204,PNNX_PT2_FAIL 清零 - F_pt2_fold_ones_like:torch.ones_like+add(标量) 静态折成常量 Attribute。 ts 侧靠 pass_level0 跑 libtorch 折常量子图;pt2 零 libtorch,利用 ones_like 值语义恒为全 1 静态求值,匹配 torch_ones_like(priority 20)归一后的形态, 修复 maximum/minimum/atan2/pow 的 ones_like(z)+0.5 常量族(4 场景) - LayerNorm/RMSNorm 模块形态:折 normalized_shape/eps,elementwise_affine 由 weight 实参存在性判定;γ/β 折 op attrs(ts level1 模块转换同构, pass_ncnn 按 @weight/@bias 捕获)(2 场景) - argument_to_constant 补 MEMORY_FORMAT/DEVICE 枚举实参转写(clone/stft/ istft 的 memory_format/device 实参,PNNX_PT2_FAIL -1) - testutil_pt2 数值对拍语义修复:batch_index=233 的 3D size-1 输入先剥 batch 维喂 ncnn;batch_index=0 输出还原形状再比。权重数值对拍 7/7 全绿 (fp16 存储下 max|d|≈3e-4),权重字节链路 state_dict→zip→Attribute→ .bin→推理端到端验证 sweep:PASS 204 / DIFF 12 / PNNX_PT2_FAIL 0 / EXPORT_FAIL 1 / SKIP 2 / 219
…0 不变 - F_pt2_fold_ones_like:输出 dtype/静态 shape/溢出校验前移到 match(), 非 f32、缺 shape、非正维、非标量 other 保持原图,write() 不再静默 留下无 data 的空 pnnx.Attribute - load_pt2 argument_to_constant:DEVICE 保留 index 编码为 type:index, cuda:1 不再降级为 cuda;无 index 为裸 type;空 device 编 None - testutil_pt2._restore_ncnn_output:仅允许剥 size-1 batch 轴这一种 还原关系(batch_index=0/233 统一语义),错序但 numel 相同的形状拒绝; 修复上一轮 233 分支过度收紧误杀 weights 场景合法剥离 - 新增 test_pt2_regress.cpp:ones_like 5 形态 + DEVICE 5 编码白盒单测 (零 libtorch,include 产品 cpp,链接行需置于最后避开注册表静态 初始化顺序问题) - 新增 test_pt2_testutil.py:helper 剥离/拒绝语义 9 断言 验证:全量 sweep PASS 204 / DIFF 12 / PNNX_PT2_FAIL 0 / TOTAL 219 与 N4 基线一致;ctest pt2 smoke 1/1;数值对拍 7/7;clang-format 10.0.1 dry-run 零 violation
新增 DIFF 基线夹具,sweep PASS 204→208 - 新增 tests/ncnn/pt2_diff_fixture.py:12 个 DIFF 场景基线夹具 (capture/verify 双模式),落盘 pnnx IR/ncnn param/bin/源模型供 M3-M5 改造后快速回归 - fuse_static_conv 新增 6 个 *_pad 融合 pass(1/2/3D × bias 有无, 先于静态折叠):F.pad(reflect/replicate) + F.conv*d(zeros) → nn.Conv*d(padding_mode=...,padding 折算),与 ts level1 nn_Conv* 同构,收敛 Conv3d reflect/replicate + groups 场景 - weight_norm 改判命令式 fold_pt2_weight_norm(pass_level2 主函数 fuse_constantlist 后调用):v/g 均为 pnnx.Attribute 且 dim=0 时就地 转 Attribute,绕开 pattern 引擎"消费者数==anchor"约束,收敛 weight_norm→Linear 多消费者场景 - F_pt2 新增 F_pt2_local_response_norm_base + nn.LocalResponseNorm / F.local_response_norm 两子类:pt2 静态 shape(含 -1)对齐 ts 动态 size+ListConstruct,收敛 2 个 LRN 场景 - Upsample/interpolate 含 .5 scale 场景定性「格式级不可达」:export 将 floor(size×scale) 子图内联为常量且 floor 有损,scale 不可恢复 验证:全量 sweep PASS 208 / DIFF 8 / EXPORT_FAIL 1 / SKIP 2 / PNNX_PT2_FAIL 0 / TOTAL 219;pyncnn 数值对拍 7/7;ctest pt2 smoke 1/1;clang-format 10.0.1 dry-run 对 4 个改动文件零 violation
|
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 98f5499f1c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (fread((char*)&z64_cd_records, sizeof(z64_cd_records), 1, fp) != 1) | ||
| continue; | ||
| if (fread((char*)&z64_cd_size, sizeof(z64_cd_size), 1, fp) != 1) | ||
| continue; | ||
| if (fread((char*)&z64_cd_offset, sizeof(z64_cd_offset), 1, fp) != 1) |
There was a problem hiding this comment.
After the fixed fields, a Zip64 EOCD contains four consecutive 64-bit values: records-on-disk, total-records, central-directory size, and central-directory offset. These reads consume only three, so z64_cd_size receives the total-record count and z64_cd_offset receives the directory size; cd_offset_valid() consequently fails. This breaks every archive emitted by this file's StoreZipWriter, which always writes a Zip64 EOCD, preventing pnnx from reopening its own parameter and foldable-constant archives.
Useful? React with 👍 / 👎.
| if (dim_index >= 0 && dim_index < (int)ishape.size() && osz.ai[i] == ishape[dim_index]) | ||
| { | ||
| osz.ai[i] = 0; |
There was a problem hiding this comment.
When an explicit adaptive-pool size happens to equal the traced input dimension, this treats the equality as proof that the source value was None and rewrites the literal to 0. The ncnn conversion maps 0 to the preserve-input-dimension sentinel, so a model traced with an 8×8 input and explicitly requesting (8, 8) incorrectly returns 10×10 when later run on 10×10 input instead of pooling to 8×8. Because this pass is globally registered, the regression also affects existing TorchScript graphs, not only PT2 inputs.
Useful? React with 👍 / 👎.
| if (spec.kind == Pt2InputSpec::TENSOR_CONSTANT) | ||
| entry = program.find_constant(spec.state_dict_name); | ||
| else | ||
| entry = program.find_weight(spec.state_dict_name); |
There was a problem hiding this comment.
For an exported module using register_buffer(..., persistent=False), the graph signature still identifies the input as BUFFER, but its bytes reside in the exported program's constants rather than its state-dict weights. This branch sends every buffer to find_weight, and the later call also selects the weights path, so conversion aborts with "weight entry not found" for such models. Treat non-persistent buffers as constants for both lookup and archive-path selection.
Useful? React with 👍 / 👎.
| Pt2OutputSpec s; | ||
| s.graph_name = parse_spec_graph_name(it->second); | ||
| out.push_back(s); |
There was a problem hiding this comment.
When the exported signature contains buffer_mutation or user_input_mutation output specs, this loop appends them exactly like user_output without checking the spec kind. load_pt2 subsequently creates a public pnnx.Output for each mutation value, changing the model's output count and ordering while not implementing the mutation semantics. Restrict this list to user_output entries, or reject unsupported mutation-bearing exports explicitly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f15074dffd
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| pnnx_ncnn_add_test(pt2_smoke) | ||
| pnnx_ncnn_add_test(pt2_weights) |
There was a problem hiding this comment.
These tests are registered unconditionally, so the existing workflow's bare ctest --output-on-failure -j 8 also runs them in every legacy matrix entry, including Torch 1.8–1.13 where torch.export does not exist. CTest documents -R as “Run tests matching regular expression”; because that job has no selector or exclusion, run_pt2_test() catches the resulting export error and returns false, causing at least test_ncnn_pt2_smoke and test_ncnn_pt2_weights to fail every affected matrix job. Gate these registrations/tests by Torch capability or exclude them from the legacy run while retaining the dedicated pt2-test job.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 39dc6a04a0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (aten_type == "aten::adaptive_avg_pool1d" || aten_type == "aten::adaptive_avg_pool2d" | ||
| || aten_type == "aten::adaptive_avg_pool3d" || aten_type == "aten::adaptive_max_pool1d" | ||
| || aten_type == "aten::adaptive_max_pool2d" || aten_type == "aten::adaptive_max_pool3d") | ||
| { | ||
| op->name = "pt2_" + op->name; |
There was a problem hiding this comment.
The new source guard does not resolve the prior issue because the loader assigns the pt2_ marker to every exported adaptive-pool node, regardless of whether output_size originally contained None. Consequently, an explicitly requested size equal to the export input dimensions still passes the guard in F_pt2_adaptive_pool_base and is rewritten to the preserve-dimension sentinel; for example, exporting explicit (8, 8) at 8×8 and later running at 10×10 produces 10×10 instead of 8×8.
Useful? React with 👍 / 👎.
| fseek(fp, extra_size - 4, SEEK_CUR); | ||
| extra_offset += extra_size; |
There was a problem hiding this comment.
For a valid stored Zip64 archive where another extra field precedes the 0x0001 Zip64 field (for example an extended-timestamp field), extra_size is the payload length after the four-byte ID/length header. Seeking by extra_size - 4 and advancing extra_offset by only extra_size leaves the cursor inside that payload, so subsequent bytes are interpreted as another header; depending on those bytes, opening the archive can fail or loop indefinitely before pnnx even detects the PT2 model. Skip all extra_size payload bytes and account for the four-byte header in the offset.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ebd45082c5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ) | ||
| output_count = len(re.findall(r"^pnnx.Output\s", param, re.MULTILINE)) | ||
| check( | ||
| result.returncode == 0 and output_count == 1, |
There was a problem hiding this comment.
For the BufferMutation model, the exported signature contains a buffer_mutation output, and the current parse_output_specs() now deliberately returns -1 for every non-user_output spec. Therefore this assertion can never pass: _export_and_convert() receives a nonzero return code and no valid parameter output. The dedicated pt2-tests workflow runs test_ncnn_pt2_review, so update the regression to expect rejection or change the parser to filter mutations while succeeding.
Useful? React with 👍 / 👎.
| if (extra_size < sizeof(zip64_extended_extra_field)) | ||
| { | ||
| // skip this extra field block | ||
| fseek(fp, extra_size - 4, SEEK_CUR); | ||
| fseek(fp, extra_size, SEEK_CUR); | ||
| extra_offset += extra_size; | ||
| continue; |
There was a problem hiding this comment.
When a standard Zip64 central-directory entry saturates only its sizes or its sizes plus local-header offset, the 0x0001 payload is respectively 16 or 24 bytes; fields are included conditionally in the order of the saturated 32-bit fields. Requiring the full 28-byte struct skips these valid payloads, leaving the sizes at 0xffffffff, so archives containing an entry over 4 GiB are recorded with the wrong size and subsequently fail or attempt a multi-gigabyte allocation. Read each required field individually based on which legacy field is saturated.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8d8f4ac646
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (dim_index >= 0 && dim_index < (int)ishape.size() && osz.ai[i] == ishape[dim_index]) | ||
| { | ||
| osz.ai[i] = 0; |
There was a problem hiding this comment.
For an export at 8×8 with output_size=(None, 8), both serialized values are 8, and this loop rewrites both to the preserve-input sentinel, so inference at 10×10 produces 10×10 instead of 10×8. Fresh evidence after the earlier guard change is that adaptive_pool_has_none records only whether any axis was None, losing which axis it was; preserve the per-axis mask rather than rewriting every value equal to the traced input size.
Useful? React with 👍 / 👎.
| return -1; | ||
| } | ||
|
|
||
| attr.data = raw; |
There was a problem hiding this comment.
When a parameter or buffer is non-contiguous, such as a transposed 2-D weight, the raw entry contains storage-order bytes while entry.sizes, entry.strides, and entry.storage_offset describe its logical view. Assigning those bytes directly to a contiguous pnnx attribute silently permutes values when the storage size equals the tensor's element count, and offset views instead fail the preceding size check; materialize the logical contiguous tensor using the serialized strides and offset.
Useful? React with 👍 / 👎.
| if (z64_total_cd_records != z64_cd_records) | ||
| continue; | ||
|
|
||
| if (cd_offset_valid(z64_cd_offset) && z64_cd_offset + z64_cd_size <= (uint64_t)file_size) |
There was a problem hiding this comment.
When Graph::save serializes a weightless graph, StoreZipWriter emits a valid Zip64 archive with zero records, a zero-size central directory, and z64_cd_offset == 0. This condition nevertheless requires a central-directory file-header signature at offset zero, so StoreZipReader::open rejects the generated .pnnx.bin; handle the zero-record/zero-size directory as a valid special case.
Useful? React with 👍 / 👎.
| if (eocd_cd_offset != 0xffffffff) | ||
| { | ||
| zip64_end_of_central_directory_locator eocdl64; | ||
| fread((char*)&eocdl64, sizeof(eocdl64), 1, fp); | ||
| if (cd_offset_valid(eocd_cd_offset) && (uint64_t)eocd_cd_offset + eocd_cd_size <= (uint64_t)file_size) | ||
| { | ||
| cd_offset = eocd_cd_offset; | ||
| cd_size = eocd_cd_size; | ||
| cd_records = eocd_records; |
There was a problem hiding this comment.
For a valid Zip64 archive with 65,536 or more entries but a central-directory offset below 4 GiB, the legacy EOCD offset remains finite while its record-count fields are 0xffff. This branch therefore skips the Zip64 locator and scans only 65,535 entries, silently omitting the rest from filemetas; select the Zip64 record when any legacy EOCD field is saturated, not only the offset.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2da8323f58
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| fprintf(stderr, " %s", it->first.c_str()); | ||
| } | ||
| fprintf(stderr, "\n"); | ||
| return Pt2Argument::NONE; |
There was a problem hiding this comment.
When a PT2 node explicitly supplies a tensor layout, such as layout=torch.sparse_coo, model.json represents it with the as_layout argument variant. This detector has no as_layout branch, logs the unknown variant, and then returns NONE; argument_to_constant() accepts that value, so conversion continues as though the layout were omitted and can silently replace sparse semantics with the default dense layout. Fresh evidence after symbolic argument handling was added is that the standard as_layout variant still reaches this fallback; either translate supported layouts or fail conversion explicitly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a81ce57381
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| const int elemsize = (attr.type == 1) ? 4 : ((attr.type == 5) ? 8 : 0); | ||
| if (elemsize == 0) | ||
| { | ||
| fprintf(stderr, "load_pt2: unsupported attribute type %d\n", attr.type); | ||
| return -1; |
There was a problem hiding this comment.
Fresh evidence after the dtype-mapping expansion is that this gate still assigns a size only to f32 and i64, so every newly mapped type is immediately rejected. For example, a valid float16 parameter maps to pnnx type 3 and then fails here with unsupported attribute type 3, preventing conversion of common half-precision PT2 models; derive the size from the mapped attribute type, as Attribute::elemsize() already does.
Useful? React with 👍 / 👎.
| if (index == 2) | ||
| return value.type == 4 && value.s == "cpu"; | ||
| if (index == 3) | ||
| return value.type == 1 && !value.b; |
There was a problem hiding this comment.
When a PT2 window node omits the schema-default pin_memory argument, the defaults table materializes it as a type-0 None, but this predicate accepts only an explicit boolean False; the analogous device default is also materialized as None while this requires the string cpu. Consequently the fold is skipped for such valid default-form hann_window/hamming_window nodes, and because this is the repository's only handler for those ATen operators, models such as STFTs retain an unsupported window operator.
Useful? React with 👍 / 👎.
| const int window_length = length.i; | ||
| Attribute attr; | ||
| attr.type = 1; | ||
| attr.shape = std::vector<int>(1, window_length); | ||
| attr.data.resize((size_t)window_length * sizeof(float)); | ||
| float* data = (float*)attr.data.data(); |
There was a problem hiding this comment.
When the default floating dtype at export is float64, a window call may still serialize dtype=None and therefore pass this fold, while its output metadata records f64. Hard-coding the replacement attribute to f32 silently changes the constant and downstream STFT numerics; use the output operand's recorded type and matching storage representation, or decline the fold when it is not f32.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 51a42bb495
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| {"LayerNorm", "normalized_shape", "normalized_shape"}, | ||
| {"LayerNorm", "eps", "eps"}, | ||
| {"RMSNorm", "normalized_shape", "normalized_shape"}, | ||
| {"RMSNorm", "eps", "eps"}, |
There was a problem hiding this comment.
When nn.RMSNorm is exported without an explicit eps, the defaults table supplies None, and this rule preserves that value in the normalized module. The downstream lowering in pass_ncnn/nn_RMSNorm.cpp translates a type-0 epsilon to 0.f, whereas PyTorch defines the default as torch.finfo(input.dtype).eps; for small-magnitude inputs this can change outputs dramatically (for example, float32 values near 1e-8 normalize near 1 with zero epsilon). Resolve None using the input dtype before lowering, or reject this default form.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ff3aec8eec
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| StoreZipReader zip; | ||
| if (zip.open(program.zippath) != 0) |
There was a problem hiding this comment.
When converting a model with many parameters or buffers, load_weight_attribute() reopens the same archive for every entry. Each StoreZipReader::open() rescans the complete central directory and seeks through every local header, while this function is itself called once per input-spec weight, making loading quadratic in the number of archive entries and potentially rendering large transformer conversions impractically slow. Open the archive once in load_pt2() and reuse that reader for all attributes.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ade9f3454b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| std::vector<float> weight = attr_v.get_float32_data(); | ||
| const std::vector<float>& weight_g = attr_g.get_float32_data(); | ||
|
|
||
| const int size = (int)(weight.size() / dim0); |
There was a problem hiding this comment.
When a valid weight-normalized module has a zero-sized leading dimension, such as weight_norm(nn.Linear(8, 0)), attr_v.shape[0] is zero and both weight attributes are empty, so the preceding size check succeeds and this integer division by zero can crash pnnx during conversion. Handle dim0 == 0 by producing the corresponding empty attribute or declining the fold safely.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8c19849730
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| uint64_t lfh_offset; | ||
| }; | ||
| std::vector<CDEntry> cdentries; | ||
| cdentries.reserve((size_t)cd_records); |
There was a problem hiding this comment.
When a malformed Zip64 archive advertises a huge cd_records value while pointing to a small in-bounds central directory, find_central_directory() accepts the metadata and this reserve attempts an enormous allocation before any header is read. On 64-bit builds this can throw std::length_error/std::bad_alloc, which is uncaught when model_file_maybe_pt2() calls StoreZipReader::open(), terminating pnnx instead of rejecting the input; verify that the record count is representable and no greater than cd_size / 46 before allocating.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dd3bf3c111
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for (size_t j = 0; j < node.inputs.size(); j++) | ||
| { | ||
| ordered_inputs.push_back(&node.inputs[j]); |
There was a problem hiding this comment.
When a PT2 model contains torch.argmax(x) or torch.argmin(x) without explicit dim and keepdim, export omits those default-valued inputs, but the generated defaults table contains neither operator, so this fallback emits only self. The existing pass_level2/torch_argmax.cpp and torch_argmin.cpp rewrites require all three inputs and therefore never match, leaving unsupported ATen operators in the converted model. Fresh evidence after the earlier zeros_like fix is that the final table still omits both aten::argmax.default and aten::argmin.default.
Useful? React with 👍 / 👎.
| attr.shape.push_back((int)entry.sizes[i]); | ||
| } | ||
|
|
||
| size_t elem_count = 1; |
There was a problem hiding this comment.
For a zero-dimensional f16 or bf16 parameter/constant, entry.sizes is empty, so this loads one element while leaving attr.shape empty. The ncnn convert_half_to_float() pass later computes the replacement size with Attribute::elemcount(), which returns zero for an empty shape, and replaces the attribute with empty data; the resulting scalar MemoryData expects one value but has no weight bytes. Handle scalar attributes in that conversion or represent their element count consistently.
Useful? React with 👍 / 👎.
The torch matrix jobs never compile; they run ctest against the prebuilt pnnx binary, so the four native harnesses added by this branch could not be found and failed all 19 jobs. Label them pt2_native, exclude the label from the matrix job and select it in the pt2-tests job, which builds them.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e30df77a28
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (attr_v.type != 1 || attr_g.type != 1) | ||
| continue; |
There was a problem hiding this comment.
When a PT2 model applies weight normalization to an f16 module (for example, weight_norm(nn.Linear(...).half())), both v and g are loaded as type 3, so this guard skips the only _weight_norm fold in the repository. The half-to-float attribute conversion does not run until pass_ncnn, after level 2, and there is no ncnn lowering for the retained aten::_weight_norm; consequently conversion leaves an unsupported operator. Convert these attributes to float for the fold (and emit a consistent output attribute) or perform half conversion before this pass.
Useful? React with 👍 / 👎.
2 participants
Motivation
torch.exportis PyTorch's modern export path and produces.pt2exported programs.pnnx currently supports TorchScript archives (
.pt) but cannot consume.pt2filesdirectly.
This PR adds native
.pt2loading and conversion support.Design
The implementation follows a clear separation:
.pt2loader faithfully transcribes the exported graph.pass_level2.pnnx.Attributeoperands.graph.tensor_valuesprovides authoritative tensor shape and dtype metadata.The implementation includes:
.pt2file detection and dispatchtorch.exportand TorchScript dual-path comparison helpersVerification
Structural parity sweep:
Numeric parity:
Regression tests:
test_ncnn_pt2_smoketest_ncnn_pt2_weightstest_ncnn_pt2_testutilA dedicated PyTorch 2.13 PT2 CI job has been added. The existing TorchScript test
matrix remains unchanged.
Known limitations
The remaining 12 structural differences are explicitly classified as:
Additional current boundaries:
dim=0ones_likefolding is restricted to validated static f32 forms