| Files | Last commit | Last update |
|---|---|---|
fix(radix_top_k): 修复 WS 变体 V→MTE3 同步缺失与 CreateVecIndex4TopK 越界，统一跨核前缀和累加为 int64 Co-authored-by: Nice_try<zhangzhiwei108@huawei.com> # message auto-generated for no-merge-commit merge: !6396 merge addut into master fix(radix_top_k): 修复 WS 变体 V→MTE3 同步缺失与 CreateVecIndex4TopK 越界，统一跨核前缀和累加为 int64 Created-by: Nice_try Commit-by: Nice_try Merged-by: cann-robot Description: ## 描述 <!--在这里详细描述你的改动，包括改动的原因和所采取的方法。--> ### 改动原因 RadixTopK 在 int64 大 shape（sortLen 超 int32 上限）与 WS 变体路径上存在三处缺陷，且 host UT 的 tiling key 期望值只编码 3 位，int64 主路径恒测不到： 1. WS 变体同步缺失： Add 在 V 流水线写 tileTopK，紧接着 MTE3 读同一块 UB；PipeBarrier<PIPE_V> 只保证 V 流水线内部顺序，SToMTE3Sync 只覆盖 S→MTE3，两者都管不到 V→MTE3，MTE3 可能读到 Add 之前的旧值，导致该段 k 配额错乱。同文件 AddTileHist2TileTopKInWs 对同样的 Add→CopyOut 序列用的是 VToMTE3Sync。 2. 索引序列越界：i64 变体的多段 regime 会命中 tileNum_ > MAX_TILE_NUM_IN_UB，此时 CreateVecIndex4TopK 中的无符号相减会下溢成巨大值，再被 tileLen_ 兜住，最终 GetWithOffset(maxIndexLen, tileNumAlign_ * 4) 的偏移会写出 6144 元素 buffer。 3. 跨核前缀和累加溢出：前驱 core 边界 bin 累计元素数总量可达 sortLen > 2^31，UB 变体用 int32 累加（cumSumValuePrev / remainCoreBoundaryNum）存在溢出风险；同时 WS 变体维护了一份只差累加器类型的 ComputeCumSumPrevI64。 4. UT 覆盖缺口：ComputeExpectedTilingKey 只编码 sorted/largest/isLargeShape 3 位，未编码 isInt64Shape（bit3），int64 主路径（sortLen > 2^31）恒测不到。 ### 改动方法 1. op_kernel/radix_top_k_ws.h - HandleLastRoundBoundary 中 Add→CopyTileTopKUb2Ws 的同步由 SToMTE3Sync() 改为 VToMTE3Sync()，与同文件 AddTileHist2TileTopKInWs 口径一致； - CreateVecIndex4TopK 增加前置防御：tileNumAlign_ >= MAX_TILE_NUM_IN_UB 时直接 return 0，下游走「现算索引」分支，与 maxIndexLen==0 的既有兜底行为一致； - 删除 ComputeCumSumPrevI64，改用基类 ComputeCumSumPrev（出参已是 int64），消除 UB/WS 两份仅差累加器类型的实现。 2. op_kernel/radix_top_k_common.h：ComputeCumSumPrev 出参由 int32_t& 改为 int64_t&（每 core 边界 bin 计数 < 2^31，int64 累加对 int32 形状同样安全）；IsNaNValue 的 float 位模式读取由指针 reinterpret 改为 union，规避浮点相等/自比较类告警；补充提取相与计数相 NaN 规范化口径不对称的说明。 3. op_kernel/radix_top_k_ub.h：cumSumValuePrev / remainCoreBoundaryNum 改为 int64，回写 tileTopK 时显式 static_cast<int32_t>（per-core 配额 < 该 tile 边界 bin 计数 < 2^31，int32 通道安全）。 4. op_host/radix_top_k_tiling.cpp：新增 static_assert(LARGE_SORT_LEN < INT32_SORT_LEN_MAX)，把「isInt64Shape_ 蕴含 isLargeShape_」钉死，避免后续调大 LARGE_SORT_LEN 时 i64 分支被漏掉；needWorkspace / GetWorkspaceSize 两处 || isInt64Shape_ 补充防御性冗余注释。 5. tests/ut/op_host/test_radix_top_k_tiling.cpp：ComputeExpectedTilingKey 补齐第 4 位 isInt64Shape（与 GET_TPL_TILING_KEY(sorted, largest, isLargeShape, isInt64Shape) 一一对应，此前 int64 主路径恒测不到）；新增两条用例：i64_shape_tiling_key_bit3（sortLen=2^31 ⇒ bit3+bit2、indices int64、必须申请 workspace）、i64_shape_reject_int32_indices（sortLen > 2^31 且 indices int32 必须 tiling 失败）；MakeRadixTopKTilingContext 支持指定输出 indices dtype。 6. math/radix_top_k/README.md：约束说明由「无」改为「本算子为选取类实现，输出为选中元素在排序轴上的索引序，不保证有序」。 7. math/topk/op_api/topk.cpp：补充 int64 大 shape 分支拦截条件的注释（说明放掉 k 比例门限与 MIN_K 门限、且 int64 索引无法通过 GatherElementsV2 补有序性的原因），不改变逻辑。 ## 关联的Issue <!-- 如果这个PR是为了解决特定的Issue，请在这里提供Issue链接。--> <!-- 如果这个PR是为了解决特定的问题单，请在这里描述问题单单号。--> - #3604 https://gitcode.com/cann/ops-math/issues/3604 ## 测试 <!--描述进行了哪些测试来验证你的改动。包括但不限于二级冒烟、算子泛化等。--> - 本次提交包含 host UT 变更（math/radix_top_k/tests/ut/op_host/test_radix_top_k_tiling.cpp）：修正 tiling key 期望值编码并新增 int64 主路径两条用例；本地未执行，预期由 CI 运行。 - kernel 侧改动（VToMTE3Sync、CreateVecIndex4TopK 防御、int64 累加）建议结合 RadixTopK int64 大 shape 场景回归验证（实施类用例未包含在本次提交内）。 ## 文档更新 <!--如果这个PR包含文档的更新，请在这里指出。例如：更新了README.md文件。--> 更新 math/radix_top_k/README.md 约束说明（补充输出无序约束）；另含 op_host/op_kernel 注释更新，不涉及其他文档。 ## 类型标签 <!-- [x] 表示选中 --> - [x] Bug修复 - [ ] 新特性 - [ ] 性能优化 - [ ] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!6396 | 29 days ago | |
新增RadixTopK算子实现 Co-authored-by: huang-wei-chen<chenhuangwei1@huawei.com> # message auto-generated for no-merge-commit merge: !3183 merge radix_top_k into master 新增RadixTopK算子实现 Created-by: huang-wei-chen Commit-by: huang-wei-chen Merged-by: cann-robot Description: ## 描述 <!--在这里详细描述你的改动，包括改动的原因和所采取的方法。--> 长序列TopK内存优化 x(1, 617025600), k=3179520，以该case为例，算子额外申请内存变化：18G→144M ## 关联的Issue <!-- 如果这个PR是为了解决特定的Issue，请在这里提供Issue链接。--> <!-- 如果这个PR是为了解决特定的问题单，请在这里描述问题单单号。--> [#1905](https://gitcode.com/cann/ops-math/issues/1905) ## 测试 <!--描述进行了哪些测试来验证你的改动。包括但不限于二级冒烟、算子泛化等。--> 500+用例精度通过 简单测试脚本如下： python import torch, torch_npu batch_list = [1, 2, 3, 4, 5, 6, 7, 8] for batch in batch_list: x = torch.randn(batch, 617025600, dtype=torch.bfloat16) values_npu, indices_npu = torch.topk(x.npu(), k=3179520, sorted=True, largest=True) values_cpu, indices_cpu = torch.topk(x, k=3179520, sorted=True, largest=True) values_match = torch.equal(values_npu.cpu(), values_cpu) print(f"batch={batch}: values_match={values_match}") ## 文档更新 <!--如果这个PR包含文档的更新，请在这里指出。例如：更新了README.md文件。--> ## 类型标签 <!-- [x] 表示选中 --> - [ ] Bug修复 - [x] 新特性 - [x] 性能优化 - [ ] 文档更新 - [ ] 其他，请描述： See merge request: cann/ops-math!3183 | 5 months ago |