• 人工智能
• 深度学习
• 算子库
• CANN
• Ascend

导读

本文围绕 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 训练/推理系列产品。

样例集总览：排序操作的两类典型场景

两个样例的计算流程在结构上完全一致：输入数据先搬入片上存储（GM → Local），调用高阶 API 完成排序/选择计算，再将结果搬出（Local → GM）。差别仅在于所调用的高阶 API、输入输出规格与 tiling 参数的来源。下文分别对两个样例做源码级剖析。

Sort 样例详解：全量降序排序并保留索引

功能与输入输出规格

Sort 样例对输入张量按数值大小进行降序排序，同时保留原始索引信息，支持float数值与uint32_t索引的联合排序。样例规格如下（对应 sort/README_en.md）：

该样例使用固定 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 三段式流水：

1. CopyIn：通过AscendC::DataCopy将mValueGlobal、mIndexGlobal两个 GlobalTensor 搬入 VECIN 队列，入队等待计算。
2. 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。

3. 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 直调模板：

1. 初始化：aclInit→aclrtSetDevice(0)→aclrtCreateContext→aclrtCreateStream。
2. 内存准备：aclrtMallocHost分配 host 侧内存，aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_FIRST)分配设备侧内存。固定ELEMENTCOUNT = 128，workspace 固定16 * 1024 * 1024字节，tiling 数据为 2 个uint32_t。
3. 数据搬运：ReadFile读取./input/input_srcGmValue.bin、./input/input_srcGmIndex.bin，经aclrtMemcpy(ACL_MEMCPY_HOST_TO_DEVICE)上板；GenerateTiling(ELEMENTCOUNT)生成的 tiling 同样拷入设备。
4. 核函数启动：

sort_custom<<<numBlocks, 0, stream>>>( srcGmValueDevice, srcGmIndexDevice, dstGmValueDevice, dstGmIndexDevice, workspaceDevice, tilingDevice);

<<<>>>内核调用运算符传入 block 数（固定 1）、l2 模式（0）与 stream。

5. 结果回读与校验：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）：

与 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。

构建选项说明

构建层面，两个样例的 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 样例以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