Star9
31
Code
Code
Issues8
Pull Requests17
Actions
Discussion
Wiki
Analysis
Star9
31
Open
gather_for_rank1迁移 #220
suhaolongAdd on 8 days ago
gather_for_rank1迁移 #220
Open
suhaolongAdd on 8 days ago
Merged fromsuhaolong/ops-recmaster_gather1intoAscend/ops-recmaster
Total: 6 files changed+402-0
Aascendc/embedding/gather_for_rank1/gather_for_rank1_forward/README.md+66-0
@@ -0,0 +1,66 @@
1+# gather_for_rank1
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| :--- | :------: |
7+| Ascend 950PR&950DT系列产品 | √ |
8+| Atlas A3系列产品 | √ |
9+| Atlas A2系列产品 | √ |
10+ 
11+## 功能说明
12+ 
13+对输入 x（一维张量）执行按索引取数（index_select / gather），即根据 index 中的索引值，从 x 中选择对应元素。
14+ 
15+计算公式：
16+ 
17+```text
18+y[i] = x[index[i]]
19+```
20+ 
21+对应接口：
22+ 
23+- aclnn 接口：`aclnnGatherForRank1`
24+- torch 接口：`ops_rec.ascendc.embedding.gather_for_rank1`
25+- 反向（自动求导）：`ops_rec.ascendc.embedding.index_select_for_rank1_backward`（独立反向算子，见聚合根 `index_select_for_rank1_backward/`）
26+ 
27+## 算子实现原理
28+ 
29+输入张量 x 为一维张量，包含 `embed_dim` 个元素；索引张量 index 为一维张量，包含 `index_num` 个索引值。算子根据 index 中的索引值，从 x 中选择对应的元素并写入输出 y，输出张量 y 的每个元素 `y[i] = x[index[i]]`。
30+ 
31+例如：
32+ 
33+```text
34+x = [x0, x1, x2, ..., x128] # shape: [129]
35+index = [3, 1, 4, 1, 5, ...] # shape: [index_num]
36+y = [x[3], x[1], x[4], x[1], x[5]] # shape: [index_num]
37+```
38+ 
39+算子实现分两步：
40+ 
41+1. 将整个 x 张量拷贝进每个核的 Unified Buffer
42+2. 按 index 分片均匀切分到各个核，逐块读取索引并执行 `Gather`，结果写回输出 y
43+ 
44+## 参数说明
45+ 
46+| 参数名 | 输入/输出 | 参数类型 | 数据类型 | 数据格式 | shape | 说明 |
47+|--------|---------|---------|---------|---------|-------|------|
48+| x | 输入 | Tensor | FLOAT16 / FLOAT32 | ND | [embed_dim] | 仅支持一维，`0 < embed_dim <= 20480` |
49+| index | 输入 | Tensor | INT64 | ND | [index_num] | 仅支持一维；index 内元素值不得超过 x 的长度（即 `< embed_dim`），否则可能导致算子执行失败 |
50+| y | 输出 | Tensor | FLOAT16 / FLOAT32 | ND | [index_num] | 从 x 中按 index 索引选取的元素 |
51+ 
52+## 约束说明
53+ 
54+- 仅支持一维输入 x 与一维索引 index
55+- 支持 FLOAT16 / FLOAT32 的 x，index 仅支持 INT64
56+- `embed_dim <= 20480`（每个核的 Unified Buffer 需要完整容纳 x）
57+- 该算子需要在 NPU 设备上运行
58+ 
59+> 通过 torch 接口调用时，扩展层还会额外校验：输入 tensor 必须已定义且非空、所有输入 tensor 必须位于同一 NPU 设备、x 与 index 必须为 1D、index 为 INT64。详见 [gather_for_rank1 接口文档](./docs/torch_gather_for_rank1.md)。
60+ 
61+## 调用示例
62+ 
63+- [aclnn 调用示例](./examples/aclnn/test_aclnn_gather_for_rank1.cpp)
64+- [torch 调用示例](./examples/torch/test_torch_gather_for_rank1.py)
65+- [aclnn 接口文档](./docs/aclnn_gather_for_rank1.md)
66+- [torch 接口文档](./docs/torch_gather_for_rank1.md)
Aascendc/embedding/gather_for_rank1/gather_for_rank1_forward/docs/aclnn_gather_for_rank1.md+83-0
@@ -0,0 +1,83 @@
1+# aclnnGatherForRank1
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| :--- | :------: |
7+| Ascend 950PR&950DT系列产品 | √ |
8+| Atlas A3系列产品 | √ |
9+| Atlas A2系列产品 | √ |
10+ 
11+## 功能说明
12+ 
13+对输入 x（一维张量）执行按索引取数，即 `y[i] = x[index[i]]`。
14+ 
15+计算公式：
16+ 
17+```text
18+y[i] = x[index[i]]
19+```
20+ 
21+## 函数原型
22+ 
23+两段式接口，必须先调用 `aclnnGatherForRank1GetWorkspaceSize` 接口获取 workspace 大小和执行器，再调用 `aclnnGatherForRank1` 接口执行计算。
24+ 
25+```c
26+aclnnStatus aclnnGatherForRank1GetWorkspaceSize(
27+ const aclTensor *x,
28+ const aclTensor *index,
29+ const aclTensor *y,
30+ uint64_t *workspaceSize,
31+ aclOpExecutor **executor);
32+ 
33+aclnnStatus aclnnGatherForRank1(
34+ void *workspace,
35+ uint64_t workspaceSize,
36+ aclOpExecutor *executor,
37+ aclrtStream stream);
38+```
39+ 
40+## aclnnGatherForRank1GetWorkspaceSize
41+ 
42+### 参数说明
43+ 
44+| 参数名 | 输入/输出 | 描述 | 数据类型 | 数据格式 |
45+|--------|---------|------|---------|---------|
46+| x（aclTensor） | 输入 | 输入数据，shape [embed_dim] | FLOAT16 / FLOAT32 | ND |
47+| index（aclTensor） | 输入 | 索引数据，shape [index_num] | INT64 | ND |
48+| y（aclTensor） | 输出 | 输出数据，shape [index_num] | FLOAT16 / FLOAT32 | ND |
49+| workspaceSize（uint64_t*） | 输出 | 返回需要在 Device 侧申请的 workspace 大小 | - | - |
50+| executor（aclOpExecutor**） | 输出 | 返回算子执行器 | - | - |
51+ 
52+### 返回值
53+ 
54+| 返回值 | 错误码 | 描述 |
55+|--------|-------|------|
56+| ACLNN_ERR_PARAM_NULLPTR | 161001 | 传入参数是必选输入/输出且为空指针 |
57+| ACLNN_ERR_PARAM_INVALID | 161002 | 参数数据类型或数据格式不在支持范围内 |
58+ 
59+## aclnnGatherForRank1
60+ 
61+### 参数说明
62+ 
63+| 参数名 | 输入/输出 | 描述 |
64+|--------|---------|------|
65+| workspace | 输入 | 在 Device 侧申请的 workspace 内存地址 |
66+| workspaceSize | 输入 | workspace 大小，由 `aclnnGatherForRank1GetWorkspaceSize` 获取 |
67+| executor | 输入 | 算子执行器 |
68+| stream | 输入 | 指定执行任务的 Stream |
69+ 
70+### 返回值
71+ 
72+aclnnStatus：返回状态码，成功返回 `ACLNN_SUCCESS`。
73+ 
74+## 约束说明
75+ 
76+- x 与 index 必须为一维
77+- x 支持 FLOAT16 / FLOAT32；index 仅支持 INT64，且 index 内元素值需小于 `embed_dim`
78+- `embed_dim <= 20480`
79+- 该算子需要在 NPU 设备上运行
80+ 
81+## 调用示例
82+ 
83+具体参考 [test_aclnn_gather_for_rank1.cpp](../examples/aclnn/test_aclnn_gather_for_rank1.cpp) 的完整实现。
Aascendc/embedding/gather_for_rank1/gather_for_rank1_forward/docs/torch_gather_for_rank1.md+83-0
@@ -0,0 +1,83 @@
1+# gather_for_rank1
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| :--- | :------: |
7+| Ascend 950PR&950DT系列产品 | √ |
8+| Atlas A3系列产品 | √ |
9+| Atlas A2系列产品 | √ |
10+ 
11+## 功能说明
12+ 
13+对输入 x（一维张量）执行按索引取数，即 `y[i] = x[index[i]]`。前向算子注册为 `ops_rec::gather_for_rank1`，其自动求导通过内置 autograd 反向绑定调用 `ops_rec::index_select_for_rank1_backward`。
14+ 
15+## 接口原型
16+ 
17+```python
18+ops_rec.ascendc.embedding.gather_for_rank1(x: Tensor, index: Tensor) -> Tensor
19+```
20+ 
21+## 参数说明
22+ 
23+| 参数名 | 输入/输出 | 参数类型 | 数据类型 | 数据格式 | shape | 说明 |
24+|--------|---------|---------|---------|---------|-------|------|
25+| x | 输入 | Tensor | FLOAT16 / FLOAT32 | ND | [embed_dim] | 仅支持一维，`0 < embed_dim <= 20480` |
26+| index | 输入 | Tensor | INT64 | ND | [index_num] | 仅支持一维；index 内元素值不得超过 x 的长度 |
27+| y | 输出 | Tensor | FLOAT16 / FLOAT32 | ND | [index_num] | 从 x 中按 index 索引选取的元素 |
28+ 
29+## 约束说明
30+ 
31+约束分为两层：算子自身语义适用于 aclnn 与 torch 两条调用路径；torch 扩展层额外加的入参校验仅在调用 `ops_rec.ascendc.embedding.gather_for_rank1` 时生效。
32+ 
33+### 算子约束
34+ 
35+- x 与 index 必须为一维
36+- x 支持 FLOAT16 / FLOAT32；index 仅支持 INT64
37+- `embed_dim <= 20480`，index 内元素值需小于 `embed_dim`
38+- 该算子需要在 NPU 设备上运行
39+ 
40+### torch 扩展层额外校验
41+ 
42+由 [gather_for_rank1_plugin.cpp](../../torch_extension/gather_for_rank1_plugin.cpp) 在调用算子前完成：
43+ 
44+- 输入 tensor 必须已定义且非空
45+- 所有输入 tensor 必须位于同一 NPU 设备
46+- x 与 index 必须为 1D
47+- index 必须为 INT64（当前迁移仅编译 int64 索引路径）
48+- 需要安装 `torch_npu`
49+ 
50+## 反向算子 index_select_for_rank1_backward
51+ 
52+前向算子的梯度通过 autograd 自动路由到反向算子，也可直接调用：
53+ 
54+```python
55+ops_rec.ascendc.embedding.index_select_for_rank1_backward(grad_y: Tensor, x: Tensor, index: Tensor) -> (Tensor, Tensor)
56+```
57+ 
58+- 返回 `(grad_x, grad_index)`，其中 `grad_index` 为未定义张量（index 不产生梯度）
59+- `grad_x` 的 shape 与 x 相同，数据类型为 FLOAT32
60+ 
61+## 安装方式
62+ 
63+```bash
64+bash build.sh --install -j16
65+```
66+ 
67+## 调用示例
68+ 
69+完整调用示例参考 [test_torch_gather_for_rank1.py](../examples/torch/test_torch_gather_for_rank1.py)。
70+ 
71+```python
72+import torch
73+import ops_rec
74+ 
75+x = torch.randn(129, device="npu")
76+index = torch.randint(0, 129, (10000,), dtype=torch.int64, device="npu")
77+x.requires_grad_()
78+y = ops_rec.ascendc.embedding.gather_for_rank1(x, index)
79+ 
80+loss = torch.mean(y)
81+loss.backward() # 自动调用 index_select_for_rank1_backward
82+print(x.grad)
83+```
Aascendc/embedding/gather_for_rank1/torch_extension/gather_for_rank1_autograd.py+83-0
@@ -0,0 +1,83 @@
1+#!/usr/bin/env python3
2+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
3+#
4+# Licensed under the Apache License, Version 2.0 (the "License");
5+# you may not use this file except in compliance with the License.
6+# You may obtain a copy of the License at
7+#
8+# http://www.apache.org/licenses/LICENSE-2.0
9+#
10+# Unless required by applicable law or agreed to in writing, software
11+# distributed under the License is distributed on an "AS IS" BASIS,
12+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+# See the License for the specific language governing permissions and
14+# limitations under the License.
15+# ==============================================================================
16+"""gather_for_rank1 前向算子的 autograd 注册（ops_rec/autograd 兼容）。
17+ 
18+前向算子 ops_rec::gather_for_rank1 计算 y = x[index]（1D 嵌入权重表按 index
19+查表）。反向梯度由主仓独立算子 ops_rec::index_select_for_rank1_backward
20+（scatter-add 语义累加 grad_x）计算，本模块用 torch.library.register_autograd
21+建立前反向绑定，使 loss.backward() 自动路由。
22+ 
23+本文件位于聚合根 torch_extension/ 下，由 build.sh sync_autograd() 同步进
24+ops_rec/autograd/，import ops_rec 时经 _load_plugins("autograd") 自动调用
25+register_autograd()。
26+"""
27+ 
28+import torch
29+ 
30+# 前向算子全名：ops_rec::gather_for_rank1
31+# 反向算子全名：ops_rec::index_select_for_rank1_backward（主仓独立算子）
32+_OP_NAMESPACE = "ops_rec"
33+_FORWARD_OP = f"{_OP_NAMESPACE}::gather_for_rank1"
34+_BACKWARD_OP = f"{_OP_NAMESPACE}::index_select_for_rank1_backward"
35+ 
36+ 
37+def _setup_context(ctx, inputs, output):
38+ """前向后保存反向所需张量。
39+ 
40+ Args:
41+ ctx: autograd context
42+ inputs: 前向输入 (x, index)
43+ output: 前向输出 y（反向算子内部自行推导，无需保存）
44+ """
45+ x, index = inputs
46+ # index 本身不可导，但反向算子需要 index 定位 grad_y 散射到 grad_x 的位置
47+ ctx.save_for_backward(x, index)
48+ 
49+ 
50+def _backward(ctx, grad_output):
51+ """反向传播。
52+ 
53+ Args:
54+ ctx: autograd context
55+ grad_output: loss 对 y 的梯度
56+ 
57+ Returns:
58+ 与前向输入 (x, index) 一一对应的梯度：(grad_x, None)
59+ """
60+ x, index = ctx.saved_tensors
61+ # 主仓反向算子返回 (grad_x, grad_index)，grad_index 恒为 undefined
62+ grad_x, _ = torch.ops.ops_rec.index_select_for_rank1_backward(grad_output, x, index)
63+ # index 不可导，其梯度为 None
64+ return grad_x, None
65+ 
66+ 
67+def register_autograd():
68+ """注册 gather_for_rank1 的 autograd 公式。
69+ 
70+ 在 ops_rec.autograd 插件包自动加载时被调用，也可手动调用。
71+ 重复注册会抛异常，因此做一次幂等保护。
72+ """
73+ try:
74+ torch.library.register_autograd(_FORWARD_OP, _backward, setup_context=_setup_context)
75+ except RuntimeError as exc:
76+ if "already" in str(exc).lower() or "registered" in str(exc).lower():
77+ return
78+ raise
79+ 
80+ 
81+if __name__ == "__main__":
82+ register_autograd()
83+
Aascendc/embedding/gather_for_rank1/torch_extension/gather_for_rank1_plugin.cpp+85-0
李振华李振华7 days ago

代码问题（按严重程度）：

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())。

likedislike
@@ -0,0 +1,85 @@
1+/* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
2+ 
3+Licensed under the Apache License, Version 2.0 (the "License");
4+you may not use this file except in compliance with the License.
5+You may obtain a copy of the License at
6+ 
7+ http://www.apache.org/licenses/LICENSE-2.0
8+ 
9+Unless required by applicable law or agreed to in writing, software
10+distributed under the License is distributed on an "AS IS" BASIS,
11+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+See the License for the specific language governing permissions and
13+limitations under the License.
14+==============================================================================*/
15+ 
16+// 声明本插件依赖的 vendor（aclnn 接口所在 op_api 库），必须在 include "pytorch_npu_helper.hpp" 之前定义
17+#define OPS_REC_CUST_VENDORS "gather_for_rank1_forward"
18+ 
19+#include <torch/library.h>
20+ 
21+#include "common_utils.h"
22+#include "pytorch_npu_helper.hpp"
23+ 
24+using namespace at;
25+ 
26+namespace npu_gather_for_rank1 {
27+ 
28+constexpr int64_t RANK1_DIM = 1;
29+ 
30+at::Tensor gather_for_rank1_impl(const at::Tensor& x, const at::Tensor& index)
31+{
32+ check_tensor_non_empty(x, "x");
33+ check_tensor_non_empty(index, "index");
34+ 
35+ // 检查NPU设备且设备ID一致
36+ std::vector<at::Tensor> tensors = {x, index};
37+ std::vector<std::string> names = {"x", "index"};
38+ check_tensor_npu_device(tensors, names);
39+ 
40+ TORCH_CHECK(x.dim() == RANK1_DIM, "The x should be 1D");
41+ TORCH_CHECK(index.dim() == RANK1_DIM, "The index should be 1D");
42+ // 当前迁移仅编译 int64 索引路径
43+ TORCH_CHECK(index.scalar_type() == at::kLong, "The index should be int64");
44+ 
45+ auto x_conti = x.contiguous();
46+ auto index_conti = index.contiguous();
47+ 
48+ at::Tensor y = at::zeros_like(index, x_conti.options());
49+ EXEC_NPU_CMD(aclnnGatherForRank1, x_conti, index_conti, y);
50+ return y;
51+}
52+ 
53+} // namespace npu_gather_for_rank1
54+ 
55+// 注册 gather_for_rank1 的 schema。
56+// autograd 反向不在此注册：由 ops_rec.autograd 模块经
57+// torch.library.register_autograd 绑定到 ops_rec::index_select_for_rank1_backward。
58+TORCH_LIBRARY_FRAGMENT(ops_rec, m)
59+{
60+ m.def("gather_for_rank1(Tensor x, Tensor index) -> Tensor");
61+}
62+ 
63+// NPU 设备实现（PrivateUse1 = Ascend NPU 后端标识）
64+TORCH_LIBRARY_IMPL(ops_rec, PrivateUse1, m)
65+{
66+ m.impl("gather_for_rank1", TORCH_FN(npu_gather_for_rank1::gather_for_rank1_impl));
67+}
68+ 
69+namespace {
70+ 
71+// Meta 实现：仅推导 shape/dtype，用于 torch.compile (inductor) 图模式的 shape 推断
72+at::Tensor gather_for_rank1_meta(const at::Tensor& x, const at::Tensor& index)
73+{
74+ TORCH_CHECK(x.dim() == npu_gather_for_rank1::RANK1_DIM, "The x should be 1D");
75+ TORCH_CHECK(index.dim() == npu_gather_for_rank1::RANK1_DIM, "The index should be 1D");
76+ return at::empty_like(index, x.options());
77+}
78+ 
79+} // anonymous namespace
80+ 
81+// Meta 实现，用于 torch.compile 的 shape 推断
82+TORCH_LIBRARY_IMPL(ops_rec, Meta, m)
83+{
84+ m.impl("gather_for_rank1", TORCH_FN(gather_for_rank1_meta));
85+}
Mdocs/zh/ops_rec_api.md+2-0
@@ -65,6 +65,7 @@ ops-rec 是面向大规模推荐场景的高性能算子库，基于昇腾 CANN
65 65 
66| 算子 | 接口签名 | 功能简介 | 硬件支持 |66| 算子 | 接口签名 | 功能简介 | 硬件支持 |
67| ---- | -------- | -------- | -------- |67| ---- | -------- | -------- | -------- |
68+| [gather_for_rank1](../../ascendc/embedding/gather_for_rank1/gather_for_rank1_forward/docs/torch_gather_for_rank1.md) | `ops_rec.ascendc.embedding.gather_for_rank1(x, index)` | 对一维张量 x 按索引 index 取数 `y[i]=x[index[i]]`，反向自动调用 `index_select_for_rank1_backward` | Atlas A2系列产品, Atlas A3系列产品, Ascend 950PR&950DT系列产品 |
68| [index_select_for_rank1_backward](../../ascendc/embedding/gather_for_rank1/index_select_for_rank1_backward/docs/torch_index_select_for_rank1_backward.md) | `ops_rec.ascendc.embedding.index_select_for_rank1_backward(grad_y, x, index)` | rank-1 嵌入查表 `y = x[index]` 的反向传播，按 scatter_add 语义累加权重表梯度，返回 `(grad_x, None)` | Atlas A2系列产品, Atlas A3系列产品, Ascend 950PR&950DT系列产品 |69| [index_select_for_rank1_backward](../../ascendc/embedding/gather_for_rank1/index_select_for_rank1_backward/docs/torch_index_select_for_rank1_backward.md) | `ops_rec.ascendc.embedding.index_select_for_rank1_backward(grad_y, x, index)` | rank-1 嵌入查表 `y = x[index]` 的反向传播，按 scatter_add 语义累加权重表梯度，返回 `(grad_x, None)` | Atlas A2系列产品, Atlas A3系列产品, Ascend 950PR&950DT系列产品 |
69 70 
70### mixing（混合算子）71### mixing（混合算子）
@@ -102,6 +103,7 @@ ascendc/
102│ └── fa_infer_with_arbitrary_func/ # Flash Attention 推理（任意区间掩码）103│ └── fa_infer_with_arbitrary_func/ # Flash Attention 推理（任意区间掩码）
103├── embedding/ # embedding 算子组（嵌入算子）104├── embedding/ # embedding 算子组（嵌入算子）
104│ └── gather_for_rank1/ # rank-1 嵌入查表（y = x[index]）105│ └── gather_for_rank1/ # rank-1 嵌入查表（y = x[index]）
106+│ ├── gather_for_rank1_forward/ # 前向（y = x[index]）
105│ └── index_select_for_rank1_backward/ # 反向（scatter_add 语义累加权重梯度）107│ └── index_select_for_rank1_backward/ # 反向（scatter_add 语义累加权重梯度）
106├── mixing/ # mixing 算子组（混合算子）108├── mixing/ # mixing 算子组（混合算子）
107│ └── token_mixing/ # Token 混合归一化109│ └── token_mixing/ # Token 混合归一化