已合并
fix: support custom torch extension package #193
RuiWang_创建于 8月27日
fix: support custom torch extension package #193
已合并
RuiWang_创建于 8月27日
RuiWang_成员
8月27日

当前PR是否有AI参与:

[x] 否
[ ] 是

PR功能描述 / 为什么需要这个合入:

  1. torch_ops_package_loader 去硬编码,改约定探测:删除 NAMESPACE_PACKAGES 映射表,从 torch.ops.<ns>.<op> 的 api_name 按约定取 ns(namespace == 包名)自动 import 扩展包。CANN site-packages 定向 find_spec + 全局兜底,builtin(aten/prim)自动跳过。新增 cann_ops_transformer/cann_ops_nn 等包无需改源码。

  2. 删除 _preload_plugin_modules:bridge.py 中 blanket exec 所有 --plugin/*.py 的预加载逻辑已删除。它唯一服务于 golden 文件顶层的 register() 调用,签名解析改由算子自带 schema(_schemas)自动获取后不再需要。

  3. 删除 3 处 worker 冗余 ensure_registered:input_generation.py、api_resolver.py、simple_param_extractor.py 中的 ensure_registered 调用已删除。父进程校验期(FrameworkApiInfoKeeper.get cache-miss)已 import,fork 的 worker 继承 sys.modules。

  4. issue #122 修复:extract_api_params 按对象判 schema,不按名字判:去掉 startswith("torch.ops.") 前缀门控,由 hasattr(obj, "_schemas") 对象级短路。3 段 cann_ops_transformer.<op>(Python wrapper function)也能走 inspect.signature 解析 kw-only。_extract_params_from_aten_schemas 的 namespace 改从 schema.name 取(ns::op → ns),不再写死四段拆名。新增 _try_upgrade_to_inspect_if_kw_mismatch:docstring 解析 kw 全 False 但 inspect 有 KEYWORD_ONLY 时升级到 inspect。

  5. 补充 validate_testcase_params 的 scalar 校验:原代码读取 api_scalar_count 但未使用(F841),补上 scalar 数量比较。

  6. 文档 + 示例用例:E2E_Case_Writing.md / e2e-case.md / write-plugin SKILL.md 补充 torch.ops 自定义算子测试方法(4 段 api_name + 自动装包 + schema 签名)。examples/case_store/e2e/torch_ops.csv 合并 torch_add + torch_npu_conv2d + causal_conv1d 示例,配套 torch_ops_causal_conv1d_golden.py。

该PR关联的issue

fixes #122

希望检视人员了解:

  • FrameworkApiInfoKeeper.register() 保留(可手动调),但 _preload_plugin_modules 已删除,golden 顶层 register() 不再自动 exec。
  • 4 段 torch.ops.<ns>.<op> 推荐(走 _schemas 权威签名),3 段 cann_ops_transformer.<op> 也支持(走 inspect.signature)。
  • ruff 全过,1305 测试通过。

测试

  • causal_conv1d_fn E2E 用例(4 段 + 3 段)均 PASS。
  • 1305 passed, 20 skipped。
  • ruff check 全过。

文档更新

  • docs/Operator_Test_Guides/E2E_Case_Writing.md:torch.ops 自定义算子测试方法
  • .claude/skills/ttk-how-write-case/references/e2e-case.md:同步
  • .claude/skills/ttk-how-write-plugin/SKILL.md:禁 register 规则

类型标签

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 RuiWang_ 的贡献)
RRuiWang_成员
8月27日 创建了 pull request,commit f2ae6a85
atomgit-bot
atomgit-bot
8月27日 评论:

变更摘要

该 PR 旨在支持自定义 torch.ops 扩展包:将 TorchOpsPackageLoader 从"硬编码已知命名空间(cann_ops_transformer、cann_ops_nn)并提前注册"改为"由 api_name 本身推导命名空间并按需惰性导入"。扩展包不再通过 ttk/cli/bridge.py 的 _preload_plugin_modules 预加载,也不再在 api_resolver.py 的 resolve_api 与 input_generation.py 的 generate_inputs 中重复调用 ensure_registered,而是在父进程校验阶段由 FrameworkApiInfoKeeper.get 触发一次导入,fork 出的 worker 通过继承 sys.modules 获得该包。同时 FrameworkApiInfoKeeper 新增了标量参数数量的一致性校验。ttk/utilities/simple_param_extractor.py 无内嵌补丁,仅从统计可见改动量较大(+731 -574)。

主要改动

  • TorchOpsPackageLoader 改为按命名空间自动发现扩展包:删除硬编码的 NAMESPACE_PACKAGES 注册表,新增 _extension_namespace 从 torch.ops.<ns>.<op> 推导包名(即命名空间本身),ensure_registered 先通过 PathFinder.find_spec 探测 CANN 环境 cann_site_packages 目录、失败后回退到全局 importlib.util.find_spec,再由 _insert_paths(带 _paths_inserted 幂等标记)插入 sys.path 并惰性导入;内置命名空间(如 aten、prim)直接跳过。
  • 移除预加载与工作进程内的重复注册:ttk/cli/bridge.py 删除 _preload_plugin_modules 函数及其在 run_with_switches 中的调用;api_resolver.py 的 resolve_api 与 input_generation.py 的 generate_inputs 删除 TorchOpsPackageLoader.ensure_registered 调用,理由为扩展包已由父进程校验时经 FrameworkApiInfoKeeper.get 导入并被 fork 的 worker 继承。
  • FrameworkApiInfoKeeper 新增标量参数数量校验:在原有张量数量校验基础上增加 scalar_count != api_scalar_count 的检查,不匹配时返回包含 API 名、标量参数数量与 testcase 配置数量的错误信息。
  • input_generation.py 异常链与导入整理:assign_tensor_value 中 raise ValueError(...) 改为 from e 保留原始异常链,并清理函数内 import 空行格式。
likedislike
不准确?
atomgit-bot
atomgit-bot
8月27日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
CANN-robotCANN-robot成员
8月27日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
8月27日 评论:

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.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.


PR Approval Progress

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

Module Approval Details

module lgtm status approve status
repo-cann/ops-test-kit ✅ 王林木, 周奇龙 (2/2) ✅ 王林木, 周奇龙 (2/1)

💡 Tip:

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

CLA Signature Pass

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

likedislike
此处折叠了41条消息 查看更多
CANN-robotCANN-robot成员
8月28日 添加了label:approved
sunday成员
8月28日 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
8月28日 添加了label:lgtm
CANN-robotCANN-robot成员
8月28日 关闭了关联的issue
CANN-robotCANN-robot成员
8月28日 合入了pull request