Pull Request已成功合入, 合并人@ascend-robot
(感谢 XianglongZeng 的贡献)变更摘要
该 PR 为 GRU 算子新增了基于 aclnn API 的 NPU 实现(此前仅存在 acl_op 路径),包含前向 _gru_npu 与反向 _gru_npu_backward 两个内核,同时修改了公共类型转换逻辑以适配 aclnn GRU 反向算子对 ACL_FORMAT_ND 格式的强制要求,并补充了更全面的前向/反向测试用例。
主要改动
- 新增 GRU 前向 aclnn 内核 (
GRUKernelNpuOpApi.cpp): 实现_gru_npu与gru函数,调用aclnnGRU执行前向计算;包含 BF16/混合精度回退到acl_op::gru的逻辑,训练模式下输出 gate 状态张量并在at::stack后显式转为ACL_FORMAT_ND以满足反向算子格式约束。 - 新增 GRU 反向 aclnn 内核 (
GRUBackwardKernelNpuOpApi.cpp): 实现_gru_npu_backward,调用aclnnGRUBackward;处理grad_y/grad_h为undefined时回退为zeros_like张量,并对hx/r/z/n/h_n/h等输入执行chunk+squeeze以适配 aclnn 所需的TensorList形参。 - 修改公共类型转换逻辑 (
op_api_common.cpp): 在ConvertType、ConvertType(重载) 和ConvertTypeV2三处,对已显式标记为ACL_FORMAT_ND的 3D 张量保留原格式,不再按维数推断为ACL_FORMAT_NCL,以避免aclnnGRUBackward因格式不匹配报错。 - 新增/注册算子声明与导数规则: 在
op_plugin_functions.yaml中将gru.input标记为op_api: all_version,新增_gru_npu和_gru_npu_backward自定义算子;在derivatives.yaml中添加_gru_npu的反向导数映射;在torch_npu_OpApi_schema_all.json中补充对应 schema。 - 新增 GRU 输出形状计算函数 (
KernelNpuOutputSize.cpp/.h): 添加gru_npu_output_size、gru_npu_hy_output_size、gru_npu_gate_output_size、gru_backward_npu_output_size、gru_backward_npu_h_prev_output_size五个工具函数,支持batch_first和双向模式。


代码审查
审查总结
| 优先级 | 数量 | 说明 |
|---|---|---|
| P1 | 1 | debug fprintf 残留在生产代码中 |
| P2 | 2 | grad_y fallback shape 错误;aclnnGRUBackward 缺少 dropout/train 参数 |
| P3 | 1 | ConvertType 注释引用不够明确 |
逐文件审查结果:
| 文件 | 结果 |
|---|---|
op_plugin/config/derivatives.yaml |
✅ 无问题 |
op_plugin/config/op_plugin_functions.yaml |
✅ 无问题 |
op_plugin/ops/opapi/GRUBackwardKernelNpuOpApi.cpp |
⚠️ 3 个问题(P1×1, P2×2) |
op_plugin/ops/opapi/GRUKernelNpuOpApi.cpp |
✅ 无问题 |
op_plugin/utils/KernelNpuOutputSize.cpp |
✅ 无问题 |
op_plugin/utils/KernelNpuOutputSize.h |
✅ 无问题 |
op_plugin/utils/op_api_common.cpp |
⚠️ 1 个建议(P3×1) |
test/core_tests/torch_npu_OpApi_schema_all.json |
✅ 无问题 |
test/test_base_ops/test_gru.py |
✅ 无问题 |
test/test_base_ops/test_gru_backward.py |
✅ 无问题 |
整体风险评估:中等。 最紧迫的问题是 GRUBackwardKernelNpuOpApi.cpp 中残留的 debug fprintf,它在每次反向传播时向 stderr 输出内部格式信息,应在上线前移除。grad_y_real 的 shape 错误在 input_size ≠ D*hidden_size 且 grad_y 未定义时会被触发,当前测试未覆盖该路径。dropout/train 参数在反向调用中的缺失需要与 CANN 算子 owner 确认是否需要补传。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 1 |
| 🟡 建议 | 2 |
⛔ 需要修改


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


【合入来源】
【修改方案】
【资料变更】
在原生api清单限制中添加dropout不生效的提示。已合入,pr见https://gitcode.com/Ascend/pytorch/pull/43385
【接口变更】
不涉及
【功能验证】
正向ut:


反向ut:
【CheckList】