企业官网建设流程全解析

CANN pyasc 排序算子 sort32 详解：一次迭代完成 32 个数排序的向量排序 API

【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc

导读

本文深入解析 CANN pyasc（Python 算子编程接口）中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。

几个关键点：

1. src1的语义：结合对应的 Ascend C 函数原型Sort32(const LocalTensor<T>& dst, const LocalTensor<T>& src0, const LocalTensor<uint32_t>& src1, const int32_t repeatTime)可见，src1在 C++ 侧被声明为LocalTensor<uint32_t>，即索引张量的元素类型应为无符号 32 位整数。仓库单元测试 python/test/unit/language/basic/test_common_api.py#L837-L847 正是用asc.uint32构造src1的，这与 C++ 原型严格一致。
2. dst与src0的数据类型一致性：dst与src0使用同一模板类型T（如float16），而src1固定为uint32_t，三者（dst、src0、src1）可以放置在不同的地址区间，TPosition 支持VECIN/VECCALC/VECOUT三种向量侧位置。
3. 返回值：None，结果直接写入dst。

底层 C++ 函数原型

sort32与 Ascend C 算子编程接口一一对应，其 C++ 侧原型如下：

template <typename T> __aicore__ inline void Sort32(const LocalTensor<T>& dst, const LocalTensor<T>& src0, const LocalTensor<uint32_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::Sort32<T>(dst, src0, src1, repeat_time)的 Ascend C 调用，最终经 CodeGen 生成可执行的昇腾算子二进制。

sort32方言算子在 include/ascir/Dialect/Asc/IR/Basic/OpProposal.td#L92-L96 中的定义如下：

def AscendC_Sort32Op : APIOp<"sort32", "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 CSort32<T>→ 向量单元执行。

约束说明

使用sort32时必须注意以下三条约束，它们在文档与 docstring（python/asc/language/basic/utils.py#L5835-L5875）中均有明确规定：

1. 相同分数时的排序稳定性：当存在score[i]与score[j]相同、且i > j时，score[j]（索引更小者）将首先被选出来、排在前面。也就是说，对于分数相等的元素，原始索引较小的元素排序后仍位于前面，这保证了排序结果可预测、可复现，便于在候选框抑制等场景中按索引回溯原始数据。
2. 迭代边界：每次迭代内的数据会进行排序，不同迭代间的数据不会进行排序。因此当待排序数据量超过 32 时，必须理解repeat_time只是"批次数"，sort32不会做跨批次的全局排序——若要整体有序，需要配合sort等其他方案或自行组织数据布局。
3. 地址对齐：操作数地址对齐要求请参见《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(dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=512) # 索引张量：uint32，放置在 VECIN，基地址 512（与 src0 错开） src1 = asc.LocalTensor(dtype=asc.uint32, pos=asc.TPosition.VECIN, addr=512, tile_size=512) # 目的张量：float16，放置在 VECOUT，基地址 0 dst = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECOUT, addr=0, tile_size=1024) # 重复迭代 4 次：对 4 * 32 = 128 个数分 4 组排序 asc.sort32(dst, src0, src1, repeat_time=4) 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 模块中还提供了其他排序相关算子，理解它们的差异有助于正确选型：

其中sort与rp_sort16的方言定义分别位于 OpProposal.td#L86-L104：AscendC_RpSort16Op每次排序 16 个 proposals，AscendC_SortOp支持is_full_sort布尔编译期属性（在 python/asc/language/basic/proposal.py#L156-L168 中会校验该参数必须为编译期常量）。如果你的场景是"候选框按置信度排序并回溯原始索引"，sort32是直接对应的接口。

注意事项与最佳实践

1. 数据量必须是 32 的整数倍：sort32以 32 个数为一组，repeat_time直接决定组数，跨组不排序，规划数据布局时务必据此切分。
2. 索引张量用uint32：与 C++ 原型LocalTensor<uint32_t>& src1对齐，避免类型不匹配导致的编译错误。
3. 相同分数行为可依赖：由于相同分数时索引更小者在前，建议在写入src0前为数据预置稳定的序号到src1，即可实现稳定的全序排列。
4. 必须在 JIT 内核中使用：sort32由@require_jit修饰（见 proposal.py#L176），只能在@asc.jit内核函数内调用，编译器在编译期将其展开为向量指令序列，因此repeat_time等参数在编译期即可确定。
5. 地址对齐：为满足向量访问约束，构造LocalTensor时注意各操作数基地址与tile_size的字节对齐，具体规则参见 Ascend C 的通用地址对齐约束。

参考资料

• API 文档源文件：docs/python-api/language/generated/asc.language.basic.sort32.md
• Python 绑定实现：python/asc/language/basic/proposal.py#L171-L180
• docstring 生成源：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