已合并
test(fx): add NPU tests for torch.fx graph APIs #35514
test(fx): add NPU tests for torch.fx graph APIs #35514
已合并
zhouzirui1234创建于 5月13日
zhouzirui1234
5月13日

【合入来源】

关联 issue:

【修改方案】

本 PR 新增 test/fx/test_fx_graph_api.py,补充 torch.fx Graph 相关 API / 内部 codegen 链路在 NPU 环境下可直接运行的轻量验证用例。测试中涉及 Tensor 的输入均通过 .to(device_type) 放到当前 accelerator 上执行。

任务 API 功能及上游社区用例情况如下:

API / 内部链路 功能说明 上游社区用例情况 本 PR 处理
torch.fx.Graph.inserting_after 设置 FX Graph 插入点,使后续新节点插入到指定节点之后。 PyTorch 上游 test/quantization/fx/test_model_report_fx.py::TestFxModelReportDetectDynamicStatic.test_nested_detection_case 中有使用,但该用例在当前 Kunpeng/ARM 环境会因 FBGEMM 条件被 skip,无法直接进入目标主体。 参考上游构图方式,新增轻量 NPU 测试直接验证插入顺序、节点参数传递和 NPU 执行结果。
torch.fx.Graph.inserting_before 设置 FX Graph 插入点,使后续新节点插入到指定节点之前。 PyTorch 上游 test/test_fx.py::TestFX.test_insertion_point 已有直接用例,但当前环境运行 test/test_fx.py 会受 libtorchbind_test.so 缺失影响。 新增独立轻量 NPU 测试,避免大文件环境依赖,同时验证插入顺序、参数重写和 NPU 执行结果。
torch.fx.Graph.lint 校验 FX Graph 结构合法性,包括节点拓扑顺序、引用关系、owner、target 等约束。 PyTorch 上游 test/test_fx.py 已有直接专项测试,例如 TestFX.test_wrong_topoTestFX.test_copy_no_remap 不重复新增 lint 专项测试;本 PR 仅在新增测试中调用 graph.lint() 校验所构造图的合法性。
torch.fx.graph.magic_methods.format torch.fx.graph 内部 codegen 模板链路,用于生成 magic method 对应的 Python 表达式。 PyTorch 上游 test/test_fx.py::TestFX.test_fx_shifts 通过 shift 表达式覆盖相关 codegen 链路。 新增 NPU 测试,通过 symbolic_trace 验证 Proxy 节点生成、节点参数、codegen 文本和 NPU 执行结果。
torch.fx.graph.inplace_methods.format torch.fx.graph 内部 inplace codegen 模板链路,用于生成 inplace 操作对应的 Python 代码。 PyTorch 上游 test/test_fx.py::TestFX.test_imul_code_print 通过 operator.imul 覆盖 inplace codegen。 新增 NPU 测试,手动构造 operator.imul 节点触发 a *= b codegen,并验证节点参数、输出结果和输入 Tensor 原地更新。

主要开发思路:

  1. 先复查 PyTorch 上游社区用例,确认 5 个任务 API / 内部链路的已有覆盖情况。
  2. torch.fx.Graph.lint,上游已有直接专项测试,因此本 PR 不重复新增 lint 测试。
  3. torch.fx.Graph.inserting_after,上游 quantization 用例会因当前 Kunpeng/ARM 环境不支持 FBGEMM 被 skip,因此参考其插入点使用方式新增轻量测试。
  4. torch.fx.Graph.inserting_beforetorch.fx.graph.magic_methods.formattorch.fx.graph.inplace_methods.format,上游相关用例集中在 test/test_fx.py,该文件在当前环境存在 libtorchbind_test.so 依赖问题,因此新增独立小文件验证核心行为。
  5. 新增测试聚焦 FX 图构建、Proxy 生成、节点 target、节点参数传递、codegen 文本以及 NPU Tensor 执行结果,不修改 API 实现。

本 PR 覆盖以下 4 个新增测试目标项:

API / 内部链路 验证内容
torch.fx.Graph.inserting_after 验证 FX Graph 插入点设置后,新节点可插入到指定节点之后,并校验节点顺序、节点参数传递、节点 target 以及 NPU 执行结果。
torch.fx.Graph.inserting_before 验证 FX Graph 插入点设置后,新节点可插入到指定节点之前,并校验节点顺序、节点参数传递、节点 target 以及 NPU 执行结果。
torch.fx.graph.magic_methods.format 通过 symbolic_trace 触发 Proxy 生成链路,验证 operator.lshift / operator.rshift 节点生成、参数传递、生成代码格式以及 NPU 执行结果。
torch.fx.graph.inplace_methods.format 手动构造 operator.imul FX 节点,验证 inplace codegen 生成 a *= b、节点参数传递、NPU 执行结果以及输入 Tensor 被原地更新。

本 PR 不新增 torch.fx.Graph.lint 专项测试,原因如下:

  1. 复查 PyTorch 官方 test/test_fx.pytorch.fx.Graph.lint 已有直接专项测试覆盖,例如 TestFX.test_wrong_topoTestFX.test_copy_no_remap 等。
  2. 本 PR 新增的 4 个测试中也在构造或 trace 完 FX Graph 后调用 graph.lint(),用于校验测试图结构合法性,但不把它作为本 PR 的新增目标 API。

本 PR 不修改 API 实现,不修改 test/test_fx.py。新增独立测试文件 test/fx/test_fx_graph_api.py,避免受 test/test_fx.py 中历史环境依赖 libtorchbind_test.so 缺失问题影响。

新增测试文件:

test/fx/test_fx_graph_api.py

新增测试方法:

test_graph_inserting_after
test_graph_inserting_before
test_magic_methods_format_codegen
test_inplace_methods_format_codegen

各测试核心验证点如下:

测试方法 核心验证点
test_graph_inserting_after 构造 placeholder -> neg -> relu -> output;使用 graph.inserting_after(neg);断言 relu 紧跟 neg;断言 neg.args == (x,)relu.args == (neg,)relu.target == torch.relu;在 NPU Tensor 上执行 GraphModule 并校验结果。
test_graph_inserting_before 先构造 placeholder -> relu -> output;使用 graph.inserting_before(relu) 插入 neg;断言 neg 位于 relu 前且相邻;断言 neg.args == (x,)relu.args == (neg,)neg.target == torch.neg;在 NPU Tensor 上执行 GraphModule 并校验结果。
test_magic_methods_format_codegen 使用 symbolic_trace trace x << 3, x >> 3;断言 FX Graph 中生成 operator.lshift / operator.rshift 节点;断言节点参数为 (x_node, 3);断言 codegen 包含 x << 3x >> 3;在 NPU LongTensor 上执行并校验结果。
test_inplace_methods_format_codegen 手动构造 operator.imul(a, b) 节点;断言节点参数和 target;断言 codegen 包含 a *= b;在 NPU Tensor 上执行并校验输出结果,同时校验输入 Tensor 被原地更新。

【资料变更】

当前不涉及资料补齐 PR。三个公开 torch.fx.Graph.* API 已在 docs/zh/native_apis 中记录;两个 torch.fx.graph.*.format 项为内部 codegen 字典模板链路,不宜按普通公开 API 强行补入 native API 表。

【接口变更】

不涉及。

【功能验证】

执行命令:

cd /root
python /workspace/pytorch/test/fx/test_fx_graph_api.py

执行结果:

....
----------------------------------------------------------------------
Ran 4 tests in 7.673s

OK

最终验证结论:

  1. 本 PR 新增 4 个轻量 NPU 测试,覆盖 inserting_afterinserting_beforemagic_methods.formatinplace_methods.format
  2. 新增测试不仅验证 API 调用不崩溃,也验证 FX 图构建、Proxy 生成、节点 target、节点参数传递、codegen 文本和 NPU Tensor 执行结果。
  3. torch.fx.Graph.lint 已有 PyTorch 官方直接专项测试覆盖,因此本 PR 不重复新增 lint 专项测试。
  4. 本 PR 不修改 API 实现,不涉及资料变更和接口变更。

【CheckList】

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 zhouzirui1234 的贡献)
Zzhouzirui1234
5月13日 创建了 pull request,commit cbd7144d
ascend-robot
ascend-robot成员
5月13日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here
You can get sig-info at here


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
test 李伟, sunyu-xuan (2/2) 李伟 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

zhouzirui-2026, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
ascend-robotascend-robot成员
5月13日 添加了label:ascend-cla/yes
ascend-robot
ascend-robot成员
5月13日 评论:

当前仓库存在以下 保护分支

Protected Branch Version Release
master
v2.7.1
v2.9.0
v2.10.0
v2.11.0
v2.12.0

评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支(创建同步 PR):
a) 如果当前 PR 是 Open 状态,同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged,将立即执行同步操作

注意:

  1. /sync 命令可以指定同步到多个分支,仅最后一个 /sync 命令生效
  2. 如果创建的同步 PR 不正确,可通过向同步 PR 的源分支提交轻量级 PR 完善,或使用 /close 命令关闭
likedislike
此处折叠了111条消息 查看更多
sunyu-xuan成员
5月21日 评论:

/lgtm

likedislike
liwei386成员
5月21日 评论:

/approve

likedislike
ascend-robotascend-robot成员
5月21日 添加了label:approvedlgtm
ascend-robot
ascend-robot成员
5月21日 评论:

Review Guide

This pull-request passes review.
Committers who wrote a comment of /approve are: 李伟.
Reviewers who wrote a comment of /lgtm are: 李伟, sunyu-xuan.

likedislike
ascend-robotascend-robot成员
5月21日 合入了pull request