- 人工智能
- 深度学习
- 算子库
- CANN
- Ascend
【免费下载链接】asc-devkit
本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。
导读
本文围绕 examples/01_simd_cpp_api/04_advanced_api/05_sort 样例集展开,讲解如何在 CANN asc-devkit 中使用 Sort 高阶 API 实现降序全排序、使用 TopK 高阶 API 实现沿最后一维取前 K 个最大/最小值及其索引的端到端算子工程。读完本文,你将掌握 Kernel 直调算子工程的目录组织、tiling 参数计算、核函数编写、host 端调用与精度校验的完整流程,并能理解GetSortMaxMinTmpSize、GetTopKMaxMinTmpSize、TopKTilingFunc等 tiling 接口在底层临时空间分配中的实际作用。
该样例集包含两个子样例:sort(全排序)与topk(TopK 选择),两者均基于 Kernel 直调算子工程,支持 Ascend 950PR/Ascend 950DT、Atlas A3 训练/推理系列、Atlas A2 训练/推理系列产品。
样例集总览:排序操作的两类典型场景
| 目录名称 | 功能描述 | 支持的产品 |
|---|---|---|
| sort | 基于 Kernel 直调算子工程,调用 Sort 高阶 API 实现 sort 单算子,按数值大小降序排序并保留原始索引信息 | Ascend 950PR/Ascend 950DT Atlas A3 训练系列产品/Atlas A3 推理系列产品 Atlas A2 训练系列产品/Atlas A2 推理系列产品 |
| topk | 基于 Kernel 直调算子工程,调用 TopK 高阶 API 实现 topk 单算子,获取最后一维的前 K 个最大值或最小值及其对应索引 | Ascend 950PR/Ascend 950DT Atlas A3 训练系列产品/Atlas A3 推理系列产品 Atlas A2 训练系列产品/Atlas A2 推理系列产品 |
两个样例的计算流程在结构上完全一致:输入数据先搬入片上存储(GM → Local),调用高阶 API 完成排序/选择计算,再将结果搬出(Local → GM)。差别仅在于所调用的高阶 API、输入输出规格与 tiling 参数的来源。下文分别对两个样例做源码级剖析。
Sort 样例详解:全量降序排序并保留索引
功能与输入输出规格
Sort 样例对输入张量按数值大小进行降序排序,同时保留原始索引信息,支持float数值与uint32_t索引的联合排序。样例规格如下(对应 sort/README_en.md):
| 项目 | 名称 | shape | 数据类型 | format |
|---|---|---|---|---|
| 输入 | srcLocalValue | [1, 128] | float | ND |
| 输入 | srcLocalIndex | [1, 128] | uint32_t | ND |
| 输出 | dstLocalValue | [1, 128] | float | ND |
| 输出 | dstLocalIndex | [1, 128] | uint32_t | ND |
| 核函数名 | sort_custom | - | - | - |
该样例使用固定 shape:输入srcLocalValue[1, 128]、srcLocalIndex[1, 128],输出dstLocalValue[1, 128]、dstLocalIndex[1, 128]。
目录结构
sort ├── scripts │ └── gen_data.py // 生成输入数据与期望(golden)数据 ├── CMakeLists.txt // 构建工程文件 ├── data_utils.h // 数据读写函数 ├── sort.asc // Ascend C 样例实现与调用示例 └── README.md // 样例文档其中 sort.asc 是核心实现文件,同时包含 tiling 计算、核函数类、host 端 main 调用三部分;data_utils.h 提供ReadFile/WriteFile二进制文件读写工具;CMakeLists.txt 以ASC语言将sort.asc编译为可执行文件demo。
Tiling 实现:GetSortMaxMinTmpSize 计算临时空间
Sort 的 tiling 在ComputeTiling中完成(sort.asc):
void ComputeTiling(uint32_t elementCount, optiling::SortCustomTilingData& tiling) { std::vector<int64_t> shapeVec = {elementCount}; AscendC::TensorShape srcShape(shapeVec); AscendC::SortConfig config; config.type = AscendC::SortType::RADIX_SORT; config.isDescend = false; uint32_t maxTmpSize; uint32_t minTmpSize; // 获取 Sort 操作所需的最大和最小临时空间大小 AscendC::GetSortMaxMinTmpSize( srcShape, AscendC::TensorDataType::DT_FLOAT, AscendC::TensorDataType::DT_UINT32, false, config, maxTmpSize, minTmpSize); uint32_t localWorkspaceSize = minTmpSize; tiling.set_elementCount(elementCount); tiling.set_sharedTmpBufferSize(localWorkspaceSize); }关键点:
- tiling 数据结构通过宏
BEGIN_TILING_DATA_DEF/TILING_DATA_FIELD_DEF/END_TILING_DATA_DEF定义,包含elementCount(待排序元素个数)与sharedTmpBufferSize(共享临时 buffer 大小)两个字段,并通过REGISTER_TILING_DATA_CLASS(SortCustom, SortCustomTilingData)注册。 GetSortMaxMinTmpSize的接口声明位于 include/adv_api/sort/sort_tiling_intf.h:srcShape为输入 shape,valueType/indexType为数值与索引数据类型,isReuseSource表示是否可复用源内存,config为排序配置,输出maxValue与minValue分别为最大/最小临时空间。SortConfig结构体(见 sort_tiling_intf.h 与 sort_utils.h)包含四个成员:type(排序算法类型,默认SortType::RADIX_SORT,即基数排序)、isDescend(是否降序)、hasSrcIndex、hasDstIndex。本样例取minTmpSize作为共享临时空间大小,以最小化片上 buffer 占用。
生成后的 tiling 数据通过SaveToBuffer序列化到内存,并由GenerateTiling返回给 host 端,最终经aclrtMemcpy拷入设备端 tiling 地址。
核函数实现:CopyIn → Sort/Extract → CopyOut
核函数类KernelSort<T, isFullSort>(sort.asc)遵循经典的 Ascend C 三段式流水:
- CopyIn:通过
AscendC::DataCopy将mValueGlobal、mIndexGlobal两个 GlobalTensor 搬入 VECIN 队列,入队等待计算。 - Compute:出队后先经
AscendC::Concat拼接数据(重复次数mConcatRepeatTimes = elementCount / 16),随后调用核心计算(sort.asc):
AscendC::Sort<T, isFullSort>(sortedLocal, concatLocal, indexLocal, sortTmp, mSortRepeatTimes);Sort模板参数T指定数据类型(本样例为float),isFullSort指定是否全排序(本样例为true)。排序完成后调用AscendC::Extract(dstValueLocal, dstIndexLocal, sortedLocal, mExtractRepeatTimes)将排序结果与索引分别抽取到输出张量。
- 临时工作空间有两路来源:当 tiling 给出的
sharedTmpBufferSize > 0时使用共享临时 buffer(dstQueSharedTmp),否则使用核内自建 buffer(sortTmpLocal),体现"以 tiling 结果驱动资源分配"的设计。 - 队列均通过
TPipe::InitBuffer初始化:输入队列深度为 2(双缓冲),输出队列与临时队列深度为 1。
- CopyOut:将 VECOUT 队列中的结果
DataCopy回 Global 内存(dstGmValue/dstGmIndex)。
核函数入口sort_custom(sort.asc)声明为extern "C" __global__ __vector__,接收源值、源索引、目标值、目标索引、workspace、tiling 六路__gm__地址参数,host 侧通过CopyTiling将 GM 上的 tiling 数据解析到结构体后调用op.Process()。
host 端调用与数据校验
main函数(sort.asc)的流程是典型的 Kernel 直调模板:
- 初始化:
aclInit→aclrtSetDevice(0)→aclrtCreateContext→aclrtCreateStream。 - 内存准备:
aclrtMallocHost分配 host 侧内存,aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_FIRST)分配设备侧内存。固定ELEMENTCOUNT = 128,workspace 固定16 * 1024 * 1024字节,tiling 数据为 2 个uint32_t。 - 数据搬运:
ReadFile读取./input/input_srcGmValue.bin、./input/input_srcGmIndex.bin,经aclrtMemcpy(ACL_MEMCPY_HOST_TO_DEVICE)上板;GenerateTiling(ELEMENTCOUNT)生成的 tiling 同样拷入设备。 - 核函数启动:
sort_custom<<<numBlocks, 0, stream>>>( srcGmValueDevice, srcGmIndexDevice, dstGmValueDevice, dstGmIndexDevice, workspaceDevice, tilingDevice);<<<>>>内核调用运算符传入 block 数(固定 1)、l2 模式(0)与 stream。
- 结果回读与校验:
aclrtSynchronizeStream同步后,将输出ACL_MEMCPY_DEVICE_TO_HOST回拷,WriteFile落盘为./output/output_dstGmValue.bin与./output/output_dstGmIndex.bin;CompareResult以绝对误差/相对误差均不超1e-4为阈值,与./output/golden_*.bin逐元素比对。
测试数据与 golden 生成
gen_data.py 使用 numpy 生成测试数据:x1为[-60000, 60000]均匀分布的 128 个 float32 随机数,x2为 128 个 0 组成的 int32 索引;golden 通过-np.sort(-x1)技巧实现降序排序并落盘。索引张量初始全 0 且 Sort 输出会重排索引,因此该样例侧重验证数值排序结果与索引随动关系。
TopK 样例详解:沿最后一维取前 K 个极值
功能与输入输出规格
TopK 样例在 Normal 模式下对输入张量进行 TopK 计算,沿最后一维获取前 K 个最大值或最小值及其对应索引,支持float数值与int32_t索引(样例规格见 topk/README_en.md):
| 项目 | 名称 | shape | 数据类型 | format |
|---|---|---|---|---|
| 输入 | srcLocalValue | [2, 32] | float | ND |
| 输入 | srcLocalIndex | [1, 32] | int32_t | ND |
| 输入 | srcLocalFinish | [1, 32] | int32_t | ND |
| 输出 | dstLocalValue | [2, 8] | float | ND |
| 输出 | dstLocalIndex | [2, 8] | int32_t | ND |
| 核函数名 | topk_custom | - | - | - |
与 Sort 不同,TopK 多了一个srcLocalFinish(finish 标志)输入:该标志用于标识某行是否已完成,若置位则输出索引被改写为总长度 N。样例固定 shape:srcLocalValue[2, 32]、srcLocalIndex[1, 32]、srcLocalFinish[1, 32],输出dstLocalValue[2, 8]、dstLocalIndex[2, 8](K=1 时按 8 对齐的 padding 输出)。
Tiling 实现:TopKTilingFunc + GetTopKMaxMinTmpSize
TopK 的 tiling 计算在 topk.asc 的GenerateTiling中完成:
AscendC::TopKTilingFunc( ascendCPlatform, inner, outter, k, dtypesize, false, AscendC::TopKMode::TOPK_NORMAL, isLargest, tiling.topKTilingData); AscendC::GetTopKMaxMinTmpSize( ascendCPlatform, inner, outter, false, false, AscendC::TopKMode::TOPK_NORMAL, isLargest, dtypesize, maxsize, minsize); tiling.set_minsize(minsize);关键点:
- tiling 数据类
TopKCustomTilingData内嵌TopkTiling topKTilingData结构体,外加k、outter、inner、n、minsize、isLargest六个字段,其中outter/inner分别对应输入张量第一维(行数)与第二维(每行长度),n为数据总长度,isLargest为 true 表示取最大值、false 表示取最小值。 TopKTilingFunc(声明见 include/adv_api/sort/topk_tiling.h)根据平台信息、inner/outter、K 值、数据类型与模式生成 TopK 接口所需的 tiling 信息。GetTopKMaxMinTmpSize(topk_tiling.h)计算 TopK 计算所需的最大/最小临时片上空间;样例取minsize作为局部临时 buffer 大小,并做 32 字节对齐修正(LOCAL_BYTES = 32)。- 平台对象
PlatformAscendC由PlatformAscendCManager::GetInstance获取:CPU 调试模式下按__NPU_ARCH__(2201 对应Ascend910B1,3510 对应Ascend950PR_9589)显式指定,NPU 模式则直接获取当前运行平台。
核函数实现:模板参数驱动的 TopK 调用
核函数类KernelTopK<T, U, isInitIndex, isHasfinish, isReuseSrc, topkMode, tmpLocal>(topk.asc)通过模板参数表达多种组合:是否初始化索引、是否携带 finish 标志、是否复用源数据、TopK 模式、是否使用局部临时空间。本样例实例化为KernelTopK<float, int32_t>(isHasfinish = true、tmpLocal = true)。
Compute 阶段的核心调用(topk.asc):
AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NORMAL>( dstLocalValue, dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, k, topKTilingData, topKInfo, isLargest);要点:
TopKInfo(定义见 include/adv_api/sort/topk_utils.h)携带outter、inner(32 字节对齐后的 inner)、n(张量实际长度)三个维度信息。TopKMode枚举(topk_utils_constants.h)包含TOPK_NORMAL(Normal 大数据量模式)与TOPK_NSMALL(小数据量模式);本样例固定TOPK_NORMAL,但核函数模板在topkMode == true时切换到TOPK_NSMALL分支,展示两种模式在统一模板下的适配能力。- 输出前先用
AscendC::Duplicate将dstLocalValue/dstLocalIndex清零,K 值按数据类型对齐为kPad(float 按 8 对齐、half 按 16 对齐,对应K_FLOAT = 8、K_HALF = 16),输出 shape 随之扩展为kPad × outter。 - 临时空间路径:
tmpLocal = true时在 VECCALC 位置初始化TBuf临时 buffer 并传入 TopK;tmpLocal = false时直接传入topKTilingData完成计算,展示两种资源供给方式。
核函数入口topk_custom(topk.asc)接收七路__gm__参数(源值、源索引、finish、目标值、目标索引、workspace、tiling),host 侧固定OUTTER = 2、INNER = 32、N = 32、K = 1、ISLARGE = false(即取每行最小值)。
host 端调用与 finish 语义验证
TopK 的 host 流程与 Sort 一致(aclInit→ 内存分配 → 读输入 → 搬 tiling →topk_custom<<<numBlocks, 0, stream>>>→ 同步 → 回拷 → 比对),差异在于多读一路input_finishGm.bin,且 tiling 大小按TILINGDATA_SIZE = 33 + 6个uint32_t计算(33 为内嵌TopkTiling结构体字段数)。golden 生成(topk/scripts/gen_data.py)通过np.argsort(x1, axis=-1, kind="stable")得到稳定排序索引,取前 K 列作为期望索引,再用np.take_along_axis取回对应数值;同时按 K 的 8 对齐补零 padding,并对 finish 置位的行将索引改写为 N——这正好呼应了核函数中 finish 标志的语义。
构建与运行:三种运行模式
在两个样例各自目录(sort/topk)下按以下步骤构建运行。
配置环境变量
根据当前环境的 CANN 开发套件安装方式配置环境变量:
source ${install_path}/cann/set_env.sh注意:
${install_path}为 CANN 包安装目录,未指定安装目录时默认安装路径为/usr/local/Ascend。
运行样例(NPU 模式,默认)
mkdir -p build && cd build; cmake -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # NPU 模式(默认) python3 ../scripts/gen_data.py # 生成测试输入数据 ./demo切换 CPU 调试 / NPU 仿真模式
cmake -DCMAKE_ASC_RUN_MODE=cpu -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # CPU 调试模式 cmake -DCMAKE_ASC_RUN_MODE=sim -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..;make -j; # NPU 仿真模式注意:切换构建模式前,需在 build 目录中删除
CMakeCache.txt清理 cmake 缓存后重新执行 cmake。
构建选项说明
| 选项 | 取值 | 说明 |
|---|---|---|
CMAKE_ASC_RUN_MODE | npu(默认)、cpu、sim | 运行模式:NPU 运行、CPU 调试、NPU 仿真 |
CMAKE_ASC_ARCHITECTURES | dav-2201(默认)、dav-3510 | NPU 架构:dav-2201 对应 Atlas A2 训练/推理系列与 Atlas A3 训练/推理系列产品;dav-3510 对应 Ascend 950PR/Ascend 950DT |
构建层面,两个样例的 CMakeLists.txt 结构相同:find_package(ASC REQUIRED)引入 ASC 语言工具链,project(kernel_samples LANGUAGES ASC CXX)声明工程语言,add_executable(demo sort.asc)(或topk.asc)生成可执行文件,并链接tiling_api、register、platform、unified_dlog、dl、m、graph_base等库;通过--npu-arch=${CMAKE_ASC_ARCHITECTURES}编译选项把目标架构传给 ASC 编译器。
执行结果
以下输出表示精度比对成功:
test pass!若比对失败,程序输出test failed!,且CompareResult会打印首个不一致元素的输出值与 golden 值,便于定位问题。
从样例到高阶 API 接口:关键能力清单
将样例中的调用点与仓库头文件对应,可以得到排序高阶 API 的完整能力视图:
| 接口 | 声明位置 | 样例中的用途 |
|---|---|---|
Sort<T, isFullSort> | 核函数高阶 API | 全量降序排序并携带索引(sort.asc) |
GetSortMaxMinTmpSize | include/adv_api/sort/sort_tiling_intf.h | 计算 Sort 最大/最小临时空间,取最小空间建 buffer |
SortConfig/SortType::RADIX_SORT | include/adv_api/sort/sort_utils.h | 指定基数排序算法与升降序、源/目标索引开关 |
TopK<T, ...> | include/adv_api/sort/topk.h | Normal 模式取前 K 个极值及索引(topk.asc) |
TopKTilingFunc | include/adv_api/sort/topk_tiling.h | 生成 TopK 接口所需的 tiling 信息 |
GetTopKMaxMinTmpSize | include/adv_api/sort/topk_tiling.h | 计算 TopK 最大/最小临时空间 |
TopKInfo | include/adv_api/sort/topk_utils.h | 携带 outter/inner/n 维度信息 |
TopKMode/TopKAlgo/TopKConfig | include/adv_api/sort/topk_utils_constants.h | Normal/NSMALL 模式、RADIX_SELECT/MERGE_SORT 算法、排序配置 |
值得注意的实现细节:
- 临时空间的两级决策:Sort 样例以
GetSortMaxMinTmpSize的minTmpSize为共享 buffer 大小,TopK 样例将minsize对齐到 32 字节后作为 VECCALC 局部 buffer;当最小空间为零或不可用时,核内自建 buffer 作为兜底。这套"tiling 定尺寸、核内按需取用"的模式是可复用的资源规划范式。 - 模板参数驱动行为差异:
KernelTopK的七个模板参数将"是否初始化索引、是否有 finish、是否复用源、大小模式、临时空间位置"全部编译期展开,避免运行期分支开销,这也是高阶 API 样例常见的组织方式。 - 数据生成与校验闭环:gen_data.py 生成的 golden 与核内逻辑严格对应(如 finish 改写索引、K 的 8 对齐 padding),保证样例可以自证正确性。
总结与扩展建议
本样例集以最小可运行的 Kernel 直调工程,完整示范了昇腾排序类算子的两种典型写法:
- Sort面向全量排序,
float + uint32_t数值/索引联合输出,适合需要"排完还要知道原来位置"的场景; - TopK面向极值选择,支持 Normal/NSMALL 双模式、最大值/最小值切换与 finish 语义,适合只关心前 K 个结果的高效检索场景。
实际开发中,你可以在此基础上做三类扩展:一是将固定 shape 改为动态 shape,让elementCount/inner/outter由运行时输入推导;二是通过修改SortConfig.isDescend与isLargest参数切换升/降序、最大/最小 K;三是参照TopKAlgo的RADIX_SELECT与MERGE_SORT两种算法枚举(topk_utils_constants.h)对比不同数据规模下的取舍。构建与调试时,建议先用cpu模式在宿主环境快速验证逻辑,再切sim仿真,最后上 NPU 实测,配合样例自带的 golden 比对机制即可完成全链路精度确认。
- 人工智能
- 深度学习
- 算子库
- CANN
- Ascend
【免费下载链接】asc-devkit
本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。
相关推荐
Wand-Enhancer 完整指南:本地补丁拆除 2 小时限制
Wand Enhancer 完整指南:本地补丁拆除 2 小时限制 Wand Enhancer 是面向 Wand(WeMod)客户端的开源本地增强工具。它把客户端
人工智能深度学习算子库CANNAscendCANN asc-devkit Math API 算子样例全解析:从 Acosh 到 Xor 的昇腾高阶数学算子实现与端到端验证
CANN asc devkit Math API 算子样例全解析:从 Acosh 到 Xor 的昇腾高阶数学算子实现与端到端验证 导读 本文以 CANN asc
人工智能深度学习算子库CANNAscend昇腾 Ascend C Arange 索引计算算子实战:基于 asc-devkit 高阶 API 的等差数列生成与端到端实现解析
昇腾 Ascend C Arange 索引计算算子实战:基于 asc devkit 高阶 API 的等差数列生成与端到端实现解析 导读 本文围绕 asc dev
人工智能深度学习算子库CANNAscend
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考