Merge blocked
变更摘要
本 PR 把 gather_for_rank1（1D 嵌入权重表按 index 查表，即 y = x[index]）迁移为 ops_rec 命名空间下的自定义算子，并同时补齐前向算子注册、NPU 与 Meta 双后端实现以及 autograd 反向绑定。PR 共涉及 6 个文件，其中新增的 gather_for_rank1_plugin.cpp 负责通过 TORCH_LIBRARY_FRAGMENT 与 TORCH_LIBRARY_IMPL 注册 aclnnGatherForRank1 调用；新增的 gather_for_rank1_autograd.py 负责把反向计算交给主仓独立算子 ops_rec::index_select_for_rank1_backward。其余 4 个文件没有内嵌补丁，因此不展开说明。
主要改动
- 新增前向算子 schema 与 NPU 实现：在
gather_for_rank1_plugin.cpp中以m.def("gather_for_rank1(Tensor x, Tensor index) -> Tensor")定义接口，并在PrivateUse1后端把该接口绑定到gather_for_rank1_impl，该实现内部通过EXEC_NPU_CMD(aclnnGatherForRank1, ...)完成查表计算并返回输出。 - 输入校验与执行约束：实现先用
check_tensor_non_empty、check_tensor_npu_device校验输入非空且位于同一 NPU 设备，再约束x、index均为 1D 且index必须是int64，随后对输入调用contiguous()，并以at::zeros_like(index, x_conti.options())构造输出张量。 - 新增 Meta 实现：通过匿名命名空间中的
gather_for_rank1_meta在Meta后端注册只推导 shape 与 dtype 的实现，供torch.compile图模式进行 shape 推断。 - 声明 vendor 依赖：在
#include "pytorch_npu_helper.hpp"之前定义OPS_REC_CUST_VENDORS "gather_for_rank1_forward"，用于指明 aclnn 接口所在的 op_api 库。 - autograd 公式注册：新增
gather_for_rank1_autograd.py，其中_setup_context保存x与index，_backward调用torch.ops.ops_rec.index_select_for_rank1_backward得到grad_x，并返回(grad_x, None)（index不可导）；register_autograd()加入重复注册的幂等保护，并说明该逻辑会经ops_rec插件加载机制自动执行。
- compile : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
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.
PR Approval Progress
⚠️ This PR does not yet meet the following requirements:lgtm (requires ≥ 2 person(s) per module)、approve (requires ≥ 1 person(s) per module)
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-Ascend/ops-rec | ❌ (0/2)(You can also ask: Fionazm, beau1, QianZH97, 文一晴, 汪泽南) | ❌ (0/1)(You can also ask: 蒋皓洋, 黄诗磊, imhandsome1, hjj12345678_, 项攀霄) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
suhaolong, thanks for your pull request. All authors of the commits have signed the CLA. 👍
ascend docs pipeline is running...
❌ 文档门禁未通过！
| 检查项 | 检查结果 | 详情 |
|---|---|---|
| markdownlint | ✅ 已通过 | 查看详情 |
| link-validity-check | ❌ 未通过 | 查看详情 |
| resource-existence-check | ✅ 已通过 | 查看详情 |
| tag-closed-check | ✅ 已通过 | 查看详情 |
💡 本次检查出 4 项错误，详情内容如下：
| 序号 | 错误详情 |
|---|---|
| 1 | [文件路径]：ascendc/embedding/gather_for_rank1/gather_for_rank1_forward/README.md [检查类型]：link-validity-check [错误位置]：第 63 行 [错误信息]：本地路径链接无法访问 [错误内容]： ./examples/aclnn/test_aclnn_gather_for_rank1.cpp |
| 2 | [文件路径]：ascendc/embedding/gather_for_rank1/gather_for_rank1_forward/README.md [检查类型]：link-validity-check [错误位置]：第 64 行 [错误信息]：本地路径链接无法访问 [错误内容]： ./examples/torch/test_torch_gather_for_rank1.py |
| 3 | [文件路径]：ascendc/embedding/gather_for_rank1/gather_for_rank1_forward/docs/aclnn_gather_for_rank1.md [检查类型]：link-validity-check [错误位置]：第 83 行 [错误信息]：本地路径链接无法访问 [错误内容]： ../examples/aclnn/test_aclnn_gather_for_rank1.cpp |
| 4 | [文件路径]：ascendc/embedding/gather_for_rank1/gather_for_rank1_forward/docs/torch_gather_for_rank1.md [检查类型]：link-validity-check [错误位置]：第 69 行 [错误信息]：本地路径链接无法访问 [错误内容]： ../examples/torch/test_torch_gather_for_rank1.py |
代码问题（按严重程度）：
1）index 取值未做任何范围校验（P0）
plugin 是唯一能拿到实际张量值的位置，但这里只校验了 dtype/shape，没有校验 index 的取值范围。kernel 里会把 int64 索引 Cast 成 int32 再乘 sizeof 当 uint32 偏移做 Gather：负值会回绕成巨大偏移，>= x.size(0) 会读到未拷贝区域，>= 2^31 还会被截断。建议在 EXEC_NPU_CMD 前加 TORCH_CHECK(index.min().item<int64_t>() >= 0 && index.max().item<int64_t>() < x.size(0))（或至少校验上界，与 op_host 的 MAX_DIM 约束对齐）。
2）输出用 zeros_like 预清零（P2）
kernel 对 Out 的 [0, indexLen) 区间是全量写，zeros_like 多了一次整块 GM 写；如果保留 zeros_like 是为了兼容 kernel 未覆盖的尾块，建议在 kernel 侧补齐写入，否则用 at::empty_like(index, x_conti.options())。
合入来源
问题/功能描述
修改方案描述
是否涉及UT/ST
开发自检