Star0

Star0

cann-robotadd infershape for top_k_pq_distance&zeta

fdfab635创建于 8月3日历史提交

文件	最后提交记录	最后更新时间
examples	add aicpu:tile_with_axis&topkpqdistance Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4174 merge dev into master add aicpu:tile_with_axis&topkpqdistance Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述新增四个 AI CPU 算子： 1. tile_with_axis（conversion）：沿指定维度复制输入 Tensor 数据，扩展输出 Tensor。包含 AI Core 和 AI CPU 双实现。 2. top_k_pq_distance（math）：基于分组极值剪枝的 TopK 距离计算，用于向量检索场景。支持动态多组输入。 3. top_k_v2_d（math）：沿指定维度找出最大或最小的 k 个元素及其索引，兼容 TensorFlow TopKV2，包含 infershape 实现。 4. zeta（math）：计算 Hurwitz zeta 函数 ζ(x, q) = Σ(n=0→∞) 1/(q+n)^x。每个算子包含：op_kernel_aicpu 实现、op_graph 构图定义、examples 示例、UT 单元测试、README 文档。 ## 关联的Issue #2294 ## 测试根据代码变更，测试场景如下： 1. tile_with_axis 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=tile_with_axis --soc=ascend910b`，17 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 2. top_k_pq_distance 单元测试 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_pq_distance --soc=ascend910b`，5 个 UT 全部通过 3. top_k_v2_d 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_v2_d --soc=ascend910b`，UT 全部通过 - 执行编包 + GE IR 运行，验证通过 4. zeta 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=zeta --soc=ascend910b`，8 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 ## 文档更新 - 新增 4 个算子 README.md（tile_with_axis、top_k_pq_distance、top_k_v2_d、zeta） - 更新 docs/zh/op_list.md，新增 4 个算子条目 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4174	1 个月前
op_graph	add aicpu:tile_with_axis&topkpqdistance Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4174 merge dev into master add aicpu:tile_with_axis&topkpqdistance Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述新增四个 AI CPU 算子： 1. tile_with_axis（conversion）：沿指定维度复制输入 Tensor 数据，扩展输出 Tensor。包含 AI Core 和 AI CPU 双实现。 2. top_k_pq_distance（math）：基于分组极值剪枝的 TopK 距离计算，用于向量检索场景。支持动态多组输入。 3. top_k_v2_d（math）：沿指定维度找出最大或最小的 k 个元素及其索引，兼容 TensorFlow TopKV2，包含 infershape 实现。 4. zeta（math）：计算 Hurwitz zeta 函数 ζ(x, q) = Σ(n=0→∞) 1/(q+n)^x。每个算子包含：op_kernel_aicpu 实现、op_graph 构图定义、examples 示例、UT 单元测试、README 文档。 ## 关联的Issue #2294 ## 测试根据代码变更，测试场景如下： 1. tile_with_axis 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=tile_with_axis --soc=ascend910b`，17 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 2. top_k_pq_distance 单元测试 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_pq_distance --soc=ascend910b`，5 个 UT 全部通过 3. top_k_v2_d 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_v2_d --soc=ascend910b`，UT 全部通过 - 执行编包 + GE IR 运行，验证通过 4. zeta 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=zeta --soc=ascend910b`，8 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 ## 文档更新 - 新增 4 个算子 README.md（tile_with_axis、top_k_pq_distance、top_k_v2_d、zeta） - 更新 docs/zh/op_list.md，新增 4 个算子条目 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4174	1 个月前
op_host	add aicpu:tile_with_axis&topkpqdistance Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4174 merge dev into master add aicpu:tile_with_axis&topkpqdistance Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述新增四个 AI CPU 算子： 1. tile_with_axis（conversion）：沿指定维度复制输入 Tensor 数据，扩展输出 Tensor。包含 AI Core 和 AI CPU 双实现。 2. top_k_pq_distance（math）：基于分组极值剪枝的 TopK 距离计算，用于向量检索场景。支持动态多组输入。 3. top_k_v2_d（math）：沿指定维度找出最大或最小的 k 个元素及其索引，兼容 TensorFlow TopKV2，包含 infershape 实现。 4. zeta（math）：计算 Hurwitz zeta 函数 ζ(x, q) = Σ(n=0→∞) 1/(q+n)^x。每个算子包含：op_kernel_aicpu 实现、op_graph 构图定义、examples 示例、UT 单元测试、README 文档。 ## 关联的Issue #2294 ## 测试根据代码变更，测试场景如下： 1. tile_with_axis 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=tile_with_axis --soc=ascend910b`，17 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 2. top_k_pq_distance 单元测试 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_pq_distance --soc=ascend910b`，5 个 UT 全部通过 3. top_k_v2_d 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_v2_d --soc=ascend910b`，UT 全部通过 - 执行编包 + GE IR 运行，验证通过 4. zeta 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=zeta --soc=ascend910b`，8 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 ## 文档更新 - 新增 4 个算子 README.md（tile_with_axis、top_k_pq_distance、top_k_v2_d、zeta） - 更新 docs/zh/op_list.md，新增 4 个算子条目 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4174	1 个月前
op_kernel_aicpu	add infershape for top_k_pq_distance&zeta Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4342 merge dev into master add infershape for top_k_pq_distance&zeta Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述为 top_k_pq_distance 与 zeta 算子新增 host 侧 InferShape 实现，并将 top_k_pq_distance、top_k_v2_d、zeta 三个 AICPU 算子的注册宏从 `REGISTER_CPU_KERNEL` 迁移到 `OPS_MATH_REGISTER_CPU_KERNELV2`，同时在对应 CMakeLists.txt 中开启 `HOSTCPU TRUE` 以编译 host 侧 infershape 模块。主要改动： - 新增 `math/top_k_pq_distance/op_host/top_k_pq_distance_infershape.cpp`：三个输出 topk_distance/topk_ivf/topk_index 的 shape 均为 `{k}`（k 来自属性），dtype 分别继承自 pq_distance[0]/pq_ivf[0]/pq_index[0] 的动态输入。 - 新增 `math/zeta/op_host/zeta_infershape.cpp`：输出 z 的 shape 为 x 与 q 的广播结果（复用 `Ops::Base::InferShape4Broadcast`），dtype 继承自 x；并实现 InferShapeRange 以支持动态 shape range 场景。 - 修改 `math/{top_k_pq_distance,top_k_v2_d,zeta}/CMakeLists.txt`：`add_all_modules_sources(... HOSTCPU TRUE)`，启用 hostcpu infershape 编译。 - 修改 `math/{top_k_pq_distance,top_k_v2_d,zeta}/op_kernel_aicpu/_aicpu.cpp`：`REGISTER_CPU_KERNEL` → `OPS_MATH_REGISTER_CPU_KERNELV2`，引入 `aicpu/math_aicpu_register.h`。 - 更新 `docs/zh/op_list.md`：top_k_pq_distance、zeta 的 hostcpu infershape 支持状态由 × 更新为 √。 ## 关联的Issue #2472 https://gitcode.com/cann/ops-math/issues/2472 ## 测试根据代码变更，测试场景如下： 1. 编译构建测试* - 执行整仓编译 `bash build.sh --pkg --jit -j16`，验证三个算子 CMakeLists.txt 的 `HOSTCPU TRUE` 配置变更正确，hostcpu infershape 模块能正常编译产出。 2. InferShape 功能验证 - top_k_pq_distance：验证三个输出 topk_distance/topk_ivf/topk_index 的 shape 均为 `{k}`，dtype 分别与 pq_distance[0]/pq_ivf[0]/pq_index[0] 一致。 - zeta：验证输出 z 的 shape 为 x 与 q 广播后的 shape，dtype 与 x 一致；覆盖静态 shape 与动态 shape range 两种场景（含 0/1/-1 的 range 用例）。 3. 算子注册验证 - 验证三个算子使用 `OPS_MATH_REGISTER_CPU_KERNELV2` 宏注册成功，算子能正常加载与执行。 4. 二级冒烟 / 算子泛化 - 对 top_k_pq_distance、zeta 执行二级冒烟用例与泛化测试，确认 infershape 迁移后算子行为与预期一致。 ## 文档更新 - 更新 `docs/zh/op_list.md`：top_k_pq_distance、zeta 的 hostcpu infershape 支持状态由 × 更新为 √。 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4342	1 个月前
tests	add aicpu:tile_with_axis&topkpqdistance Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4174 merge dev into master add aicpu:tile_with_axis&topkpqdistance Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述新增四个 AI CPU 算子： 1. tile_with_axis（conversion）：沿指定维度复制输入 Tensor 数据，扩展输出 Tensor。包含 AI Core 和 AI CPU 双实现。 2. top_k_pq_distance（math）：基于分组极值剪枝的 TopK 距离计算，用于向量检索场景。支持动态多组输入。 3. top_k_v2_d（math）：沿指定维度找出最大或最小的 k 个元素及其索引，兼容 TensorFlow TopKV2，包含 infershape 实现。 4. zeta（math）：计算 Hurwitz zeta 函数 ζ(x, q) = Σ(n=0→∞) 1/(q+n)^x。每个算子包含：op_kernel_aicpu 实现、op_graph 构图定义、examples 示例、UT 单元测试、README 文档。 ## 关联的Issue #2294 ## 测试根据代码变更，测试场景如下： 1. tile_with_axis 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=tile_with_axis --soc=ascend910b`，17 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 2. top_k_pq_distance 单元测试 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_pq_distance --soc=ascend910b`，5 个 UT 全部通过 3. top_k_v2_d 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_v2_d --soc=ascend910b`，UT 全部通过 - 执行编包 + GE IR 运行，验证通过 4. zeta 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=zeta --soc=ascend910b`，8 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 ## 文档更新 - 新增 4 个算子 README.md（tile_with_axis、top_k_pq_distance、top_k_v2_d、zeta） - 更新 docs/zh/op_list.md，新增 4 个算子条目 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4174	1 个月前
CMakeLists.txt	add infershape for top_k_pq_distance&zeta Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4342 merge dev into master add infershape for top_k_pq_distance&zeta Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述为 top_k_pq_distance 与 zeta 算子新增 host 侧 InferShape 实现，并将 top_k_pq_distance、top_k_v2_d、zeta 三个 AICPU 算子的注册宏从 `REGISTER_CPU_KERNEL` 迁移到 `OPS_MATH_REGISTER_CPU_KERNELV2`，同时在对应 CMakeLists.txt 中开启 `HOSTCPU TRUE` 以编译 host 侧 infershape 模块。主要改动： - 新增 `math/top_k_pq_distance/op_host/top_k_pq_distance_infershape.cpp`：三个输出 topk_distance/topk_ivf/topk_index 的 shape 均为 `{k}`（k 来自属性），dtype 分别继承自 pq_distance[0]/pq_ivf[0]/pq_index[0] 的动态输入。 - 新增 `math/zeta/op_host/zeta_infershape.cpp`：输出 z 的 shape 为 x 与 q 的广播结果（复用 `Ops::Base::InferShape4Broadcast`），dtype 继承自 x；并实现 InferShapeRange 以支持动态 shape range 场景。 - 修改 `math/{top_k_pq_distance,top_k_v2_d,zeta}/CMakeLists.txt`：`add_all_modules_sources(... HOSTCPU TRUE)`，启用 hostcpu infershape 编译。 - 修改 `math/{top_k_pq_distance,top_k_v2_d,zeta}/op_kernel_aicpu/_aicpu.cpp`：`REGISTER_CPU_KERNEL` → `OPS_MATH_REGISTER_CPU_KERNELV2`，引入 `aicpu/math_aicpu_register.h`。 - 更新 `docs/zh/op_list.md`：top_k_pq_distance、zeta 的 hostcpu infershape 支持状态由 × 更新为 √。 ## 关联的Issue #2472 https://gitcode.com/cann/ops-math/issues/2472 ## 测试根据代码变更，测试场景如下： 1. 编译构建测试* - 执行整仓编译 `bash build.sh --pkg --jit -j16`，验证三个算子 CMakeLists.txt 的 `HOSTCPU TRUE` 配置变更正确，hostcpu infershape 模块能正常编译产出。 2. InferShape 功能验证 - top_k_pq_distance：验证三个输出 topk_distance/topk_ivf/topk_index 的 shape 均为 `{k}`，dtype 分别与 pq_distance[0]/pq_ivf[0]/pq_index[0] 一致。 - zeta：验证输出 z 的 shape 为 x 与 q 广播后的 shape，dtype 与 x 一致；覆盖静态 shape 与动态 shape range 两种场景（含 0/1/-1 的 range 用例）。 3. 算子注册验证 - 验证三个算子使用 `OPS_MATH_REGISTER_CPU_KERNELV2` 宏注册成功，算子能正常加载与执行。 4. 二级冒烟 / 算子泛化 - 对 top_k_pq_distance、zeta 执行二级冒烟用例与泛化测试，确认 infershape 迁移后算子行为与预期一致。 ## 文档更新 - 更新 `docs/zh/op_list.md`：top_k_pq_distance、zeta 的 hostcpu infershape 支持状态由 × 更新为 √。 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4342	1 个月前
README.md	add aicpu:tile_with_axis&topkpqdistance Co-authored-by: sujunwei3<sujunwei3@huawei.com> # message auto-generated for no-merge-commit merge: !4174 merge dev into master add aicpu:tile_with_axis&topkpqdistance Created-by: sujunwei3 Commit-by: sujunwei3 Merged-by: cann-robot Description: ## 描述新增四个 AI CPU 算子： 1. tile_with_axis（conversion）：沿指定维度复制输入 Tensor 数据，扩展输出 Tensor。包含 AI Core 和 AI CPU 双实现。 2. top_k_pq_distance（math）：基于分组极值剪枝的 TopK 距离计算，用于向量检索场景。支持动态多组输入。 3. top_k_v2_d（math）：沿指定维度找出最大或最小的 k 个元素及其索引，兼容 TensorFlow TopKV2，包含 infershape 实现。 4. zeta（math）：计算 Hurwitz zeta 函数 ζ(x, q) = Σ(n=0→∞) 1/(q+n)^x。每个算子包含：op_kernel_aicpu 实现、op_graph 构图定义、examples 示例、UT 单元测试、README 文档。 ## 关联的Issue #2294 ## 测试根据代码变更，测试场景如下： 1. tile_with_axis 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=tile_with_axis --soc=ascend910b`，17 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 2. top_k_pq_distance 单元测试 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_pq_distance --soc=ascend910b`，5 个 UT 全部通过 3. top_k_v2_d 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=top_k_v2_d --soc=ascend910b`，UT 全部通过 - 执行编包 + GE IR 运行，验证通过 4. zeta 单元测试 + GE IR 验证 - 执行 `bash build.sh -u --opkernel_aicpu --ops=zeta --soc=ascend910b`，8 个 UT 全部通过 - 执行编包 + GE IR 运行，`execute samples success`，输出 bin 文件正常 ## 文档更新 - 新增 4 个算子 README.md（tile_with_axis、top_k_pq_distance、top_k_v2_d、zeta） - 更新 docs/zh/op_list.md，新增 4 个算子条目 ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [x] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!4174	1 个月前

TopKV2D

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	√
Atlas 推理系列产品	√
Atlas 训练系列产品	√

功能说明

算子功能：沿指定维度找出输入张量中最大或最小的k个元素及其索引。与TopK相比，TopKV2D多了一个assist_seq辅助输入，兼容TensorFlow的TopKV2算子。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	输入张量。	FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64	ND
k	输入	要取出的元素个数。	INT32	ND
assist_seq	输入	辅助序列张量。	FLOAT16	ND
sorted	属性	是否对输出结果排序，默认为true。	Bool	-
dim	属性	指定沿哪个维度进行操作，默认为-1（最后一维）。	Int	-
largest	属性	是否取最大值，默认为true。若为false则取最小值。	Bool	-
values	输出	输出的top-k个元素值。	FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64	ND
indices	输出	输出的top-k个元素在输入张量中的索引。	INT32	ND

约束说明

dim必须在[-x维度数, x维度数)范围内
k必须大于等于0且小于等于x在dim维度上的大小

调用说明

调用方式	样例代码	说明
图模式接口	test_geir_top_k_v2_d	通过算子IR接口方式调用TopKV2D算子。