CANN pyasc 排序算子 sort32 详解一次迭代完成 32 个数排序的向量排序 API【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读本文深入解析 CANN pyascPython 算子编程接口中asc.language.basic.sort32的完整用法与实现原理。sort32是昇腾 AI 处理器上的向量排序函数一次迭代可完成 32 个数的排序其输入由分数张量 索引张量组成输出为排序后的 (score, index) 对广泛用于候选框/Region Proposal 排序等需要按值排序并携带原始索引的场景。读完本文你将掌握sort32的签名、参数语义、约束边界、底层 IR 生成与代码发射链路并能基于仓库中的单元测试模板快速编写可运行的算子内核。功能概述向量化排序与索引联动sort32属于 python/asc/language/basic 模块提供的矢量Vector基础算子其定位是按分数排序并同步输出对应的索引单次迭代处理32 个数通过repeat_time重复迭代即可排序更多数据排序结果以(score, index) 对的形式存储即不仅输出排序后的数值还输出每个数值在原输入中的索引位置排序在向量单元Vector Unit上执行属于典型的按迭代块批式处理语义每次迭代内部的数据会排序不同迭代之间的数据不会互相排序。这一数值 索引成对输出的设计在源码中有直接佐证。在 include/ascir/Dialect/Asc/IR/Basic/OpProposal.td#L92-L96 中AscendC_Sort32Op的描述明确写道AscendC::Sort32sorts 32 elements per iteration, storing results as (score, index) pairs.也就是说该算子本质上完成的是带索引的归并排序与单纯输出有序序列的sort形成功能互补详见后文与相邻排序 API 的对比。函数签名与参数详解asc.language.basic.sort32(dst: LocalTensor, src0: LocalTensor, src1: LocalTensor, repeat_time: int) - None对应地在 Python 绑定源码 python/asc/language/basic/proposal.py#L171-L180 中sort32同时提供了overload的静态类型声明与require_jit装饰的 JIT 实现说明它必须在asc.jit修饰的算子内核函数内调用由 JIT 编译器在编译期展开为底层 IR。参数类型说明dstLocalTensor目的操作数保存排序结果支持的 TPosition 为VECIN/VECCALC/VECOUTsrc0LocalTensor源操作数保存待排序的分数score数据支持的 TPosition 为VECIN/VECCALC/VECOUTsrc1LocalTensor源操作数保存与分数对应的原始索引index数据支持的 TPosition 为VECIN/VECCALC/VECOUTrepeat_timeint重复迭代次数即每次迭代完成 32 个数的排序总排序数据量为repeat_time * 32几个关键点src1的语义结合对应的 Ascend C 函数原型Sort32(const LocalTensorT dst, const LocalTensorT src0, const LocalTensoruint32_t src1, const int32_t repeatTime)可见src1在 C 侧被声明为LocalTensoruint32_t即索引张量的元素类型应为无符号 32 位整数。仓库单元测试 python/test/unit/language/basic/test_common_api.py#L837-L847 正是用asc.uint32构造src1的这与 C 原型严格一致。dst与src0的数据类型一致性dst与src0使用同一模板类型T如float16而src1固定为uint32_t三者dst、src0、src1可以放置在不同的地址区间TPosition 支持VECIN/VECCALC/VECOUT三种向量侧位置。返回值None结果直接写入dst。底层 C 函数原型sort32与 Ascend C 算子编程接口一一对应其 C 侧原型如下template typename T __aicore__ inline void Sort32(const LocalTensorT dst, const LocalTensorT src0, const LocalTensoruint32_t src1, const int32_t repeatTime)在 pyasc 的编译链路中该 Python API 通过 IR 构建器生成asc.sort32方言算子在 python/asc/language/basic/proposal.py#L179-L180 中调用builder.create_asc_Sort32Op(...)并将repeat_time通过_mat(repeat_time, KT.int32)转换为int32类型的编译期常量随后由lib/Target/AscendC/Basic/OtherOps.cpp中的代码发射逻辑OtherOps.cpp#L222-L234输出形如asc::Sort32T(dst, src0, src1, repeat_time)的 Ascend C 调用最终经 CodeGen 生成可执行的昇腾算子二进制。sort32方言算子在 include/ascir/Dialect/Asc/IR/Basic/OpProposal.td#L92-L96 中的定义如下def AscendC_Sort32Op : APIOpsort32, Sort32, [AscFunc] { let summary Call AscendC::Sort32 method; let description AscendC::Sort32 sorts 32 elements per iteration, storing results as (score, index) pairs.\n; let arguments (ins AnyType:$dst, AnyType:$src0, AnyType:$src1, AnyType:$repeatTime); }从该定义可以确认完整调用链Python 层asc.sort32→ MLIRasc.sort32Op携带 dst/src0/src1/repeatTime 四个操作数→ 代码发射为 Ascend CSort32T→ 向量单元执行。约束说明使用sort32时必须注意以下三条约束它们在文档与 docstringpython/asc/language/basic/utils.py#L5835-L5875中均有明确规定相同分数时的排序稳定性当存在score[i]与score[j]相同、且i j时score[j]索引更小者将首先被选出来、排在前面。也就是说对于分数相等的元素原始索引较小的元素排序后仍位于前面这保证了排序结果可预测、可复现便于在候选框抑制等场景中按索引回溯原始数据。迭代边界每次迭代内的数据会进行排序不同迭代间的数据不会进行排序。因此当待排序数据量超过 32 时必须理解repeat_time只是批次数sort32不会做跨批次的全局排序——若要整体有序需要配合sort等其他方案或自行组织数据布局。地址对齐操作数地址对齐要求请参见《Ascend C 算子开发接口》中的通用说明和约束—通用地址对齐约束在 pyasc 中构造LocalTensor时应确保各操作数满足向量访问的对齐要求。调用示例从测试用例到真实内核最小可运行示例根据仓库单元测试 python/test/unit/language/basic/test_common_api.py#L837-L847一个完整的sort32内核写法如下import asc asc.jit def sort32_kernel(): # 分数张量float16放置在 VECIN基地址 0 src0 asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) # 索引张量uint32放置在 VECIN基地址 512与 src0 错开 src1 asc.LocalTensor(dtypeasc.uint32, posasc.TPosition.VECIN, addr512, tile_size512) # 目的张量float16放置在 VECOUT基地址 0 dst asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECOUT, addr0, tile_size1024) # 重复迭代 4 次对 4 * 32 128 个数分 4 组排序 asc.sort32(dst, src0, src1, repeat_time4) sort32_kernel[1]()上述代码中的tile_size表示张量占用字节数addr用于规划各张量在统一内存Unified Buffer中的布局测试中src1与src0基地址错开 512 字节dst的tile_size为 1024容纳排序后的 score 与 index 输出。运行该内核后mock_launcher_run应被调用一次即验证了一次完整的编译 → 发射 → 运行链路。官方文档中的调用示例# repeat_time 4, 对128个数分成4组进行排序每次完成1组32个数的排序 asc.sort32(dst_local, src_local0, src_local1, 4)即当repeat_time 4时对 128 个数分成 4 组每组 32 个数独立排序。与相邻排序 API 的对比pyasc 的 basic 模块中还提供了其他排序相关算子理解它们的差异有助于正确选型API排序规模输出形式适用场景sort32每次迭代 32 个(score, index) 对需要按值排序 保留原始索引的成对排序sort按repeat_time重复支持is_full_sort仅数值降序排序纯数值降序排列不同平台布局有差异rp_sort16每次迭代 16 个 Region Proposal按 score 降序的 Proposal目标检测中 Region Proposal 的专用排序其中sort与rp_sort16的方言定义分别位于 OpProposal.td#L86-L104AscendC_RpSort16Op每次排序 16 个 proposalsAscendC_SortOp支持is_full_sort布尔编译期属性在 python/asc/language/basic/proposal.py#L156-L168 中会校验该参数必须为编译期常量。如果你的场景是候选框按置信度排序并回溯原始索引sort32是直接对应的接口。注意事项与最佳实践数据量必须是 32 的整数倍sort32以 32 个数为一组repeat_time直接决定组数跨组不排序规划数据布局时务必据此切分。索引张量用uint32与 C 原型LocalTensoruint32_t src1对齐避免类型不匹配导致的编译错误。相同分数行为可依赖由于相同分数时索引更小者在前建议在写入src0前为数据预置稳定的序号到src1即可实现稳定的全序排列。必须在 JIT 内核中使用sort32由require_jit修饰见 proposal.py#L176只能在asc.jit内核函数内调用编译器在编译期将其展开为向量指令序列因此repeat_time等参数在编译期即可确定。地址对齐为满足向量访问约束构造LocalTensor时注意各操作数基地址与tile_size的字节对齐具体规则参见 Ascend C 的通用地址对齐约束。参考资料API 文档源文件docs/python-api/language/generated/asc.language.basic.sort32.mdPython 绑定实现python/asc/language/basic/proposal.py#L171-L180docstring 生成源python/asc/language/basic/utils.py#L5835-L5875方言算子定义include/ascir/Dialect/Asc/IR/Basic/OpProposal.td#L92-L96代码发射实现lib/Target/AscendC/Basic/OtherOps.cpp#L222-L234单元测试python/test/unit/language/basic/test_common_api.py#L837-L847算子导出声明python/asc/language/basic/init.py#L139【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
