Star1246
1370
代码
代码
Issues58
Pull Requests325
Actions
项目讨论
Wiki
分析
Star1246
1370
已合并
add aicpu:tile_with_axis&topkpqdistance #4174
sujunwei3创建于 23 天前
add aicpu:tile_with_axis&topkpqdistance #4174
已合并
sujunwei3创建于 23 天前
从sujunwei3/ops-mathdev合入到cann/ops-mathmaster
S
sujunwei3
23 天前

描述

新增四个 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 个算子条目

类型标签

  • Bug修复
  • 新特性
  • 性能优化
  • 文档更新
  • 其他，请描述：
likedislike
Pull Request已成功合入, 合并人@CANN-robot
（感谢 sujunwei3 的贡献）
Ssujunwei3
23 天前 创建了 pull request，commit f3c2b93a
atomgit-bot
atomgit-bot
23 天前 评论：

变更摘要

本 PR 新增了四个 AICPU 算子：TileWithAxis、TopKPQDistance、TopKV2D 和 Zeta。每个算子均包含算子原型定义（REG_OP）、AICPU 内核实现（继承 CpuKernel）、OpDef 注册、GE IR 集成测试示例以及单元测试。其中 TopKV2D 额外包含 Host 端 InferShape 逻辑。这些算子覆盖了张量沿轴复制、基于堆的 Top-K 距离筛选、沿指定维度 Top-K 查找以及 Hurwitz zeta 函数计算等场景。

主要改动

  • 新增 TileWithAxis 算子：在 conversion/tile_with_axis/ 下实现沿指定轴重复张量数据的 AICPU 内核，支持 float16/float/int8~int64/uint8~uint64 共 11 种数据类型，通过 Eigen 的 broadcast 机制完成沿 axis 轴的 tiles 次复制，支持 0~8 维输入。

  • 新增 TopKPQDistance 算子：在 math/top_k_pq_distance/ 下实现基于堆的 Top-K 距离筛选内核，接收动态输入 actual_count、pq_distance、grouped_extreme_distance、pq_ivf、pq_index，通过两级堆（分组极值堆与距离 Top-K 堆）筛选出 K 个最大/最小距离及其对应的 IVF 桶号和索引，支持 float32/float16，并支持多批次输入。

  • 新增 TopKV2D 算子：在 math/top_k_v2_d/ 下实现沿指定维度查找 K 个最大/最小值的 AICPU 内核，支持 float16/float/double/int8~int64/uint8~uint64 共 11 种数据类型，包含单调性检测优化、选择排序（小 n 小 k）、堆排序（通用路径）及全排序（k=n）四条计算路径，并通过 CpuKernelUtils::ParallelFor 实现多核并行。

  • 新增 Zeta 算子：在 math/zeta/ 下实现 Hurwitz zeta 函数的 AICPU 内核，基于 Eigen 的 numext::zeta 逐元素计算，支持 float32 和 double，内置并行阈值（64K 元素以上启用 ParallelFor），并对输入输出类型一致性、数据大小等进行校验。

  • 新增 TopKV2D 的 Host 端 InferShape：在 math/top_k_v2_d/op_host/top_k_v2_d_infershape.cpp 中实现根据 dim 属性和 k 输入张量推导输出形状的逻辑，替换 dim 轴维度为 k 值，并注册为 IMPL_OP_INFERSHAPE。

likedislike
atomgit-bot
atomgit-bot
23 天前 评论：

代码审查

审查总结

已完成对全部 33 个变更文件的逐一审查。发现共 8 个问题：

优先级 数量 说明
P0 2 test_top_k_pq_distance.cpp 中两处 CpuKernelUtils::CpuKernelUtils:: 双重限定符，可能导致编译失败
P1 2 zeta_aicpu.cpp 除零错误；top_k_pq_distance_aicpu.cpp VLA 栈溢出风险
P2 2 tile_with_axis_aicpu_def.cpp 中 tiles 属性 OPTIONAL/REQUIRED 不一致；tile_with_axis_proto.h 声明 BF16 但 def/kernel 未实现
P3 2 test_geir_top_k_v2_d.cpp 和 test_geir_zeta.cpp 示例代码中 fwrite 返回值未检查

各文件审查结果

文件 结果
conversion/tile_with_axis/README.md 无问题
conversion/tile_with_axis/examples/test_geir_tile_with_axis.cpp 无问题
conversion/tile_with_axis/op_graph/tile_with_axis_proto.h P2: BF16 类型声明未实现
conversion/tile_with_axis/op_kernel_aicpu/tile_with_axis_aicpu.cpp 无问题
conversion/tile_with_axis/op_kernel_aicpu/tile_with_axis_aicpu.h 无问题
conversion/tile_with_axis/op_kernel_aicpu/tile_with_axis_aicpu_def.cpp P2: tiles 属性标记不一致
conversion/tile_with_axis/tests/ut/op_kernel_aicpu/test_tile_with_axis.cpp 无问题
docs/zh/op_list.md 无问题
math/top_k_pq_distance/CMakeLists.txt 无问题
math/top_k_pq_distance/README.md 无问题
math/top_k_pq_distance/examples/test_geir_top_k_pq_distance.cpp 无问题
math/top_k_pq_distance/op_graph/top_k_pq_distance_proto.h 无问题
math/top_k_pq_distance/op_kernel_aicpu/top_k_pq_distance_aicpu.cpp P1: VLA 栈溢出风险
math/top_k_pq_distance/op_kernel_aicpu/top_k_pq_distance_aicpu.h 无问题
math/top_k_pq_distance/op_kernel_aicpu/top_k_pq_distance_aicpu_def.cpp 无问题
math/top_k_pq_distance/tests/ut/op_kernel_aicpu/test_top_k_pq_distance.cpp P0: 双重限定符×2
math/top_k_v2_d/CMakeLists.txt 无问题
math/top_k_v2_d/README.md 无问题
math/top_k_v2_d/examples/test_geir_top_k_v2_d.cpp P3: fwrite 返回值未检查
math/top_k_v2_d/op_graph/top_k_v2_d_proto.h 无问题
math/top_k_v2_d/op_host/top_k_v2_d_infershape.cpp 无问题
math/top_k_v2_d/op_kernel_aicpu/top_k_v2_d_aicpu.cpp 无问题
math/top_k_v2_d/op_kernel_aicpu/top_k_v2_d_aicpu.h 无问题
math/top_k_v2_d/op_kernel_aicpu/top_k_v2_d_aicpu_def.cpp 无问题
math/top_k_v2_d/tests/ut/op_kernel_aicpu/test_top_k_v2_d.cpp 无问题
math/zeta/CMakeLists.txt 无问题
math/zeta/README.md 无问题
math/zeta/examples/test_geir_zeta.cpp P3: fwrite 返回值未检查
math/zeta/op_graph/zeta_proto.h 无问题
math/zeta/op_kernel_aicpu/zeta_aicpu.cpp P1: total=0 除零
math/zeta/op_kernel_aicpu/zeta_aicpu.h 无问题
math/zeta/op_kernel_aicpu/zeta_aicpu_def.cpp 无问题
math/zeta/tests/ut/op_kernel_aicpu/test_zeta.cpp 无问题

整体风险评估：该 PR 新增了 4 个 AICPU 算子（TileWithAxis、TopKPQDistance、TopKV2D、Zeta）及其配套的测试、示例和文档。核心 kernel 逻辑基本正确，但存在 2 个 P0 编译阻断问题（测试代码双重限定符）、2 个 P1 运行时风险（除零和栈溢出）以及 2 个 P2 接口一致性问题。建议优先修复 P0 和 P1 项后再合入。

类型 数量
🔴 阻塞 6
🟡 建议 3

⛔ 需要修改

likedislike
CANN-robotCANN-robot成员
23 天前 添加了label：cann-cla/yes
CANN-robot
CANN-robot成员
23 天前 评论：

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

PR Approval Progress

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

Module Approval Details

module lgtm status approve status
*/*/README.md ✅ 陈娇, wangrui_, 宋恺 (3/2) ✅ 陈娇 (1/1)
*/*/op_graph/*_proto.h ✅ 汤磊, wangrui_, 宋恺 (3/2) ✅ 汤磊 (1/1)
repo-cann/ops-math ✅ wangrui_, 宋恺 (2/2) ✅ wangrui_ (1/1)

💡 Tip:

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

CLA Signature Pass

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

likedislike
此处折叠了70条消息 查看更多
S
songkai111成员
20 天前 评论：

/lgtm

likedislike
CANN-robot
CANN-robot成员
20 天前 评论：

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
刘伟

likedislike
CANN-robotCANN-robot成员
20 天前 添加了label：lgtm
CANN-robotCANN-robot成员
20 天前 关闭了关联的issue
CANN-robotCANN-robot成员
20 天前 合入了pull request