已关闭
[Requirement|需求建议] 为 top_k_pq_distance 与 zeta 算子补充 InferShape 实现 #2472
sujunwei3创建于  8月3日关闭于  8月3日
sujunwei3
8月3日 创建

一、背景信息 (必填)

top_k_pq_distance、zeta 两个 AICPU 算子此前缺少 host 侧 InferShape 实现,无法在 hostcpu infershape 链路上完成 shape/dtype 推导;同时 top_k_pq_distance、top_k_v2_d、zeta 三个算子仍在使用旧版 REGISTER_CPU_KERNEL 宏注册,未迁移到 ops-math 仓统一的 OPS_MATH_REGISTER_CPU_KERNELV2 注册方式。

二、价值/作用 (必填)

  1. 补齐 top_k_pq_distance、zeta 的 InferShape,使两个算子支持 hostcpu infershape 链路,shape/dtype 推导不再依赖 GE 默认逻辑。
  2. 统一三个算子(top_k_pq_distance、top_k_v2_d、zeta)的注册宏为 OPS_MATH_REGISTER_CPU_KERNELV2,与 ops-math 仓整体注册方式保持一致。
  3. 在 docs/zh/op_list.md 中将两个算子的 hostcpu infershape 支持状态更新为 √,文档与实现同步。

三、设计方案 (必填)

3.1 使能方式

AICPU 算子,通过 CMakeLists.txt 中 add_all_modules_sources(... HOSTCPU TRUE) 开启 hostcpu infershape 模块编译。

3.2 总体设计
3.2.1 算子支持的数据类型
  • top_k_pq_distance:输出 dtype 继承自对应动态输入的第 0 个实例(pq_distance/pq_ivf/pq_index)。
  • zeta:输出 z 的 dtype 继承自输入 x。
3.2.2 host侧设计
  • top_k_pq_distance:三个输出 topk_distance/topk_ivf/topk_index 的 shape 均为 {k}(k 来自属性),dtype 分别继承自 pq_distance[0]/pq_ivf[0]/pq_index[0]。
  • zeta:z 的 shape 为 x 与 q 的广播结果(复用 Ops::Base::InferShape4Broadcast),dtype 继承自 x;并实现 InferShapeRange 以支持动态 shape range 场景,按低维对齐、逐维广播/取交推导输出 range。
3.2.3 kernel侧设计

三个算子的 aicpu kernel 注册宏由 REGISTER_CPU_KERNEL 替换为 OPS_MATH_REGISTER_CPU_KERNELV2,引入 aicpu/math_aicpu_register.h 头文件,kernel 计算逻辑不变。

3.3 支持硬件

AI CPU(ascend910b 等)。

3.4 算子约束限制

  • top_k_pq_distance:k 为属性给定,输出固定为一维 {k}。
  • zeta:x 与 q 需满足广播规则;InferShapeRange 对含 0/1/-1 的 range 做了分类处理。

💡 备注
本 Issue 由 PR #4342 自动分析创建。

likedislike
sujunwei3
8月3日 评论:

/assign

likedislike
CANN-robotCANN-robot成员
8月3日 将 sujunwei3 设为负责人
CANN-robotCANN-robot成员
8月3日 关闭了 issue
CANN-robotCANN-robot成员
8月3日 添加了label:resolved