| 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 | 3 months ago |