CANN asc-devkit 多路归并排序实战：Sort32 与 MrgSort 协同实现全局有序

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

本篇技术指南基于 asc-devkit 仓库中的 MrgSort 排序样例，完整讲解如何在昇腾 AI 处理器上利用Sort32与MrgSort两个向量指令接口，以"分组局部排序 + 多轮归并"的方式实现大规模数据的全局降序排序。读完本文后，你将掌握两个接口的调用方式、MrgSort4Info与MrgSortSrcList参数结构的完整语义、三种典型归并场景（4 队列、非 4 对齐、32 队列多轮归并）的工程实现，以及基于 CMake 的编译、运行与精度验证全流程。

为什么需要"排序 + 归并"两步组合

Sort32与MrgSort是昇腾向量计算单元（Vector Core）上的两条配套指令，两者单独使用都无法完成全局排序，必须组合成流水线：

• Sort32负责"局部有序"：每次迭代对 32 个元素执行降序排列，一次调用可连续处理多组 32 元素，把整段数据切分成若干条内部有序的子序列（队列）；
• MrgSort负责"全局有序"：将最多 4 条已排好序的队列按 score 域从大到小归并成 1 条队列，通过多次迭代或多轮调用，即可把全部子序列逐级合并为一条全局有序结果。

两者组合后的整体流程（见 mrg_sort.asc 头部注释）为：先调用Sort32将数据按每 32 个元素一组并行降序排序，形成以(score, index)交替结构存储的有序队列；再调用MrgSort将这些有序队列合并为一个全局有序结果。这正是多路归并排序（multi-way merge sort）在昇腾硬件上的落地形态。

支持的产品与 CANN 版本

本样例在以下产品与 CANN 版本上验证通过（与 README_en.md 及两个接口的官方 API 文档一致）：

需要说明的是，Sort32与MrgSort接口本身在 Sort32 API 文档 与 MrgSort API 文档 中还标注支持 Atlas 200I/500 A2 推理产品，但不支持 Atlas 推理系列产品 AI Core / Vector Core 及 Atlas 训练系列产品，实际部署前请以目标硬件型号为准。

样例目录结构

├── mrg_sort │ ├── scripts │ │ ├── gen_data.py // 输入数据和真值（golden）数据生成脚本 │ │ └── verify_result.py // 输出数据与真值数据的比对验证脚本 │ ├── CMakeLists.txt // 编译工程文件 │ ├── data_utils.h // 数据读入写出工具函数 │ ├── mrg_sort.asc // Ascend C 样例实现与调用 │ └── README.md // 样例说明文档（中文版）

核心接口速览：Sort32 与 MrgSort

Sort32：32 元素一次迭代降序排序

Sort32的原型（见 Sort32 API 文档）为：

template <typename T> __aicore__ inline void Sort32(const LocalTensor<T>& dst, const LocalTensor<T>& src0, const LocalTensor<uint32_t>& src1, const int32_t repeatTime)

其语义要点如下：

• 一次迭代完成 32 个数的降序排列；repeatTime为重复迭代次数，取值范围[0, 255]，每次迭代输入 src0 跳过 128 字节、src1 跳过 128 字节、输出 dst 跳过 256 字节；
• src0是待排序的 score，src1是对应索引 index，src1 固定为uint32_t类型；
• 排序结果以(score, index)结构写入 dst，该结构无论 score 是float还是half都固定占用 8 字节：当 src0 为float时，dst 中 index 存储在高 4 字节、score 存储在低 4 字节；当 src0 为half时，index 存储在高 4 字节、score 存储在低 2 字节，中间 2 字节保留；
• 支持的数据类型组合为half/uint32_t/half与float/uint32_t/float；
• 每次迭代输入 src0 要求 128 字节（float）/64 字节（half）对齐，src1 要求 128 字节对齐，输出 dst 要求 256 字节对齐。

下图展示了Sort32将 score 与 index 排序后落盘为(score, index)交替结构的过程（图片来源：Sort32 API 文档）：

MrgSort：最多 4 条有序队列的归并

MrgSort的原型为：

template <typename T> __aicore__ inline void MrgSort(const LocalTensor<T>& dst, const MrgSortSrcList<T>& src, const MrgSort4Info& params)

它将已经排好序的最多 4 条队列合并排列为 1 条队列，结果按 score 域从大到小排序。它不负责对无序输入做排序，通常处理经过Sort32预处理的数据。支持half/float两种操作数类型，流水类型为 PIPE_V。

两个关键参数结构体

MrgSort4Info是归并参数的载体，字段含义如下：

MrgSortSrcList<T>是源队列集合：包含src1、src2、src3、src4四个LocalTensor<T>成员，分别对应 4 条已排序队列的起始地址，每个起始地址要求 8 字节对齐，4 项都必须初始化（无效队列可复用合法 UB 地址并将对应elementLengths置 0）。

三个场景深入解读

样例通过编译参数SCENARIO_NUM在三个场景间切换（在 CMakeLists.txt 中通过add_definitions(-DSCENARIO_NUM=${SCENARIO_NUM})注入宏），所有场景数据格式均为 ND，核函数名统一为vec_mrgsort_kernel。核函数以srcDataSize、workSize两个模板参数驱动，main函数根据场景号推导输入尺寸：

constexpr uint32_t SRC_DATA_SIZE = (scenarioNum == 1) ? 128 : (scenarioNum == 2) ? 96 : (scenarioNum == 3) ? 1024 : 128; constexpr uint32_t WORK_SIZE = (scenarioNum == 1) ? 4 * SORT_SIZE * 2 : (scenarioNum == 2) ? 4 * SORT_SIZE * 2 : (scenarioNum == 3) ? (SRC_DATA_SIZE / SORT_SIZE) * SORT_SIZE * 6 : 4 * SORT_SIZE * 2;

场景 1：4 条队列合并为 1 条

• 输入：128 个float元素（score）+ 128 个uint32_t元素（index）；
• 输出：[1, 256]个float元素（128×2 个排序结果）；
• 实现：Sort32的repeat = 4，将 128 个元素分为 4 组分别降序排序；随后MrgSort以validBit = 0b1111、ifExhaustedSuspension = false、repeatTimes = 1将 4 条队列完全合并为 1 条有序队列。

对应源码（mrg_sort.asc）中，Sort32的输出workLocal被划分为 4 条队列，srcList依次取workLocal[0]、workLocal[SORT_SIZE*1*2]、workLocal[SORT_SIZE*2*2]、workLocal[SORT_SIZE*3*2]（乘以 2 是因为每个数据单元占 2 个 float 位置），4 条队列长度均为SORT_SIZE = 32。

场景 2：3 条队列非 4 对齐合并

• 输入：96 个float元素 + 96 个uint32_t元素；
• 输出：[1, 192]个float元素（96×2 个排序结果）；
• 实现：Sort32的repeat = 3，将 96 个元素分为 3 组降序排序；MrgSort以validBit = 0b0111、ifExhaustedSuspension = false、repeatTimes = 1将 3 条队列合并为 1 条。

该场景演示的是非 4 对齐下的归并：validBit = 0b0111表示前 3 条队列有效，第 4 条队列在elementLengths中被显式置为 0（源码中通过(validBit == 0b1111) ? SORT_SIZE : 0动态决定）。这正是处理"队列总数不能被 4 整除"时的手法——不存在的队列给出合法地址并将长度置 0，即可复用同一套归并逻辑。

场景 3：32 条队列多轮归并

• 输入：1024 个float元素 + 1024 个uint32_t元素；
• 输出：[1, 2048]个float元素（1024×2 个排序结果）；
• 实现：Sort32的repeat = 32，将 1024 个元素分为 32 组降序排序，得到 32 条有序队列；随后进行三轮MrgSort：
  • 第一轮：repeatTimes = 8，每 4 条队列合并为 1 条，得到 8 条有序队列；
  • 第二轮：repeatTimes = 2，每 4 条队列合并为 1 条，得到 2 条有序队列；
  • 第三轮：validBit = 0b11、repeatTimes = 1，将最后 2 条队列合并为 1 条全局有序队列（第 3、4 条队列长度置 0）。

该场景是多路归并排序的完整示范。源码对workLocal工作空间做了三段式划分（见 mrg_sort.asc）：workLocal[0:2048]存放Sort32后的 32 条有序队列；workLocal[2048:4096]作为第一轮归并结果存储区、同时是第二轮归并的源数据；workLocal[4096:6144]作为第二轮归并结果存储区、同时是第三轮归并的源数据。每轮归并之间以AscendC::PipeBarrier<PIPE_V>()保证向量流水同步，防止上一轮结果尚未就绪便被读取。

需要特别注意：当MrgSort的repeatTimes > 1时，validBit必须为 15 且 4 条队列长度一致、连续存储，因此第一、二轮都使用validBit = 0b1111；只有最后一轮只剩 2 条队列时才改用validBit = 0b11并将repeatTimes降回 1。

样例输入输出规格

表 1：场景 1 规格

表 2：场景 2 规格

表 3：场景 3 规格

编译与运行

在样例根目录examples/01_simd_cpp_api/03_basic_api/01_memory_vector_compute/mrg_sort下依次执行以下步骤。

1. 配置环境变量

请先根据当前环境中 CANN 开发套件包的安装方式完成安装，然后执行：

source ${install_path}/cann/set_env.sh

说明：${install_path}为 CANN 包安装目录，未指定安装目录时默认安装至/usr/local/Ascend下。

2. 编译、运行与验证

SCENARIO_NUM=1 mkdir -p build && cd build; # 创建并进入 build 目录 cmake -DSCENARIO_NUM=$SCENARIO_NUM -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..; make -j; # 编译工程，默认 npu 模式 python3 ../scripts/gen_data.py -scenarioNum=$SCENARIO_NUM # 生成测试输入数据与 golden 数据 ./demo # 运行编译生成的可执行程序 python3 ../scripts/verify_result.py -scenarioNum=$SCENARIO_NUM output/output.bin output/golden.bin # 比对输出与真值

3. CPU 调试与 NPU 仿真模式

使用 CPU 调试或 NPU 仿真模式时，在 cmake 阶段追加-DCMAKE_ASC_RUN_MODE=cpu或-DCMAKE_ASC_RUN_MODE=sim参数：

cmake -DSCENARIO_NUM=$SCENARIO_NUM -DCMAKE_ASC_RUN_MODE=cpu -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..; make -j; # CPU 调试模式 cmake -DSCENARIO_NUM=$SCENARIO_NUM -DCMAKE_ASC_RUN_MODE=sim -DCMAKE_ASC_ARCHITECTURES=dav-2201 ..; make -j; # NPU 仿真模式

注意：切换编译模式前需清理 cmake 缓存，可在 build 目录下执行rm CMakeCache.txt后重新运行 cmake。

4. 编译选项说明

CMAKE_ASC_ARCHITECTURES的默认值dav-2201定义在 CMakeLists.txt 中，并通过target_compile_options以--npu-arch=${CMAKE_ASC_ARCHITECTURES}传给 ASC 编译器；SCENARIO_NUM通过add_definitions编译进核函数源码（constexpr uint32_t scenarioNum = SCENARIO_NUM;），因此切换场景后需要重新编译。

5. 执行结果

验证脚本比对通过时输出如下，表示精度对比成功：

test pass!

数据生成与精度验证脚本

gen_data.py：输入与 golden 数据生成

gen_data.py 以固定随机种子np.random.seed(9)生成np.random.uniform(1, 100, [total_elements])的 float32 score 数据，index 初始化为 0，写入./input/input_x.bin与./input/input_y.bin。真值由 Host 端np.argsort(input_x)[::-1]得到全局降序索引，再按(score, index)交替布局写入./output/golden.bin——其中偶数位golden[2*i]存排序后的 score，奇数位由索引映射填充，模拟硬件上每个数据单元占 8 字节的存储格式。

verify_result.py：输出与真值比对

verify_result.py 读取output/output.bin与output/golden.bin，按场景截取期望的 float 数量（场景 1/2/3 分别为 256/192/2048），使用np.isclose逐元素比对，容差设置为相对误差rtol = 1e-3、绝对误差atol = 1e-5，并以整体误差比例error_tol = 1e-3判定结果是否通过；不通过时打印前 100 个差异点的期望值、实际值与相对偏差，便于定位排序结果中的问题元素。

底层实现原理：从接口到硬件指令

MrgSort与Sort32的实现集中在 impl/basic_api/kernel_operator_proposal_intf_impl.h（接口声明位于 include/basic_api/kernel_operator_proposal_intf.h）。阅读源码可以印证样例中的参数语义：

• 参数校验：MrgSort4/MrgSort在入口处校验操作数类型必须是half/float组合，elementLengths必须落在[0, 4095]，validBit必须是3 / 7 / 15（见 kernel_operator_proposal_intf_impl.h），与 API 文档的约束完全一致；
• 指令封装：MrgSort4将repeatTimes、4 个elementLengths、ifExhaustedSuspension、validBit逐位打包进一个 64 位config寄存器字，MrgSort则将elementLengths打包进src1、把validBit/ifExhaustedSuspension写入config的低位域（见 kernel_operator_proposal_intf_impl.h），最终统一调用向量指令内建函数Vmrgsort4Cal下发到硬件；
• 队列寻址：4 条源队列通过src.src1~src4.GetPhyAddr()组成物理地址数组，与样例中srcList.srcN = workLocal[...]的取址方式一一对应；
• CPU 调试兼容：在ASCENDC_CPU_DEBUG宏下，MrgSort会依据validBit将未被选中的无效队列地址置为nullptr，保证 CPU 模拟路径下的行为安全（见 kernel_operator_proposal_intf_impl.h）。

此外，该头文件还提供了MrgSort4（直接封装Vmrgsort4Cal的变体）、Sort32以及GetMrgSortResult等关联接口，读者可在了解本样例后继续阅读源码探索完整接口族。

核函数内部的执行细节

以场景 1 为例，核函数vec_mrgsort_kernel（mrg_sort.asc）的执行链路清晰展示了"搬运 → 排序 → 归并 → 回写"的典型模式：

1. UB 空间分配：通过AscendC::LocalMemAllocator<AscendC::Hardware::UB>在统一缓冲区（UB）中分配srcLocal0（score）、srcLocal1（index）、workLocal（Sort32 输出工作区，float）、dstLocal（归并结果区，大小为srcDataSize * 2）；
2. GM→UB 搬运：DataCopy将两份输入拷入 UB，随后SetFlag/WaitFlag<HardEvent::MTE2_V>同步搬运完成（确保 MTE2 与向量计算流水之间的数据依赖）；
3. 局部排序：AscendC::Sort32<float>(workLocal, srcLocal0, srcLocal1, repeat)完成分组降序排序，repeat = srcDataSize / SORT_SIZE；
4. 流水同步：PipeBarrier<PIPE_V>()确保Sort32写出的(score, index)队列对后续MrgSort可见；
5. 多路归并：按场景配置MrgSort4Info与MrgSortSrcList，调用AscendC::MrgSort<float>(dstLocal, srcList, params)；
6. UB→GM 回写：SetFlag/WaitFlag<HardEvent::V_MTE3>同步向量计算完成后，DataCopy(dstGlobal, dstLocal, dstDataSize)将结果写回全局内存，最后PipeBarrier<PIPE_ALL>()收尾。

Host 端main函数则使用aclrtSetDevice/aclrtCreateStream/aclrtMalloc/aclrtMemcpy等 ACL 运行时 API 完成设备初始化、显存分配、H2D/D2H 拷贝，并通过ReadFile/WriteFile（data_utils.h）读写input_x.bin、input_y.bin与output.bin。

约束与注意事项总结

结合 API 文档与样例源码，使用Sort32+MrgSort时必须注意以下边界条件：

• 地址对齐：MrgSort的 dst 起始地址需 32 字节对齐，4 条源队列起始地址需 8 字节对齐；Sort32每次迭代输入 src0 需 128/64 字节对齐、src1 需 128 字节对齐、输出需 256 字节对齐；
• 地址重叠：MrgSort不支持源操作数与目的操作数地址重叠，样例中dstLocal与workLocal是独立分配的 UB 空间；
• 稳定序语义：存在score[i] == score[j]且i > j时，接口优先将score[j]（索引更小者）排在前面，即保持相同 score 元素按原索引升序的稳定行为；
• 迭代间不排序：每次迭代内部的数据参与排序，不同迭代间的数据不会被排序——这正是场景 3 需要"多轮归并"而不是把 32 条队列一次性交给MrgSort的原因；
• repeatTimes > 1 的硬性条件：4 条源队列长度一致、连续存储、ifExhaustedSuspension = false、validBit = 15，缺一不可；
• repeatTimes = 0 的行为差异：在 Atlas A3 与 Atlas A2 系列上视为 NOP（不写目的操作数），而在 Ascend 950PR&950DT 系列上不保证为 NOP，生产代码应避免依赖该取值；
• 无效队列处理：srcList的 4 项必须全部初始化，无效队列复用合法 UB 地址并将对应elementLengths置 0（样例场景 2、场景 3 第三轮均如此处理）。

通过本文介绍的三个场景与完整编译运行流程，读者可以在此基础上扩展出任意队列数量的多轮归并排序方案——只要数据量满足 UB 容量约束，即可用"Sort32 分组 + MrgSort 逐轮归并"的模板覆盖更大规模的排序需求。

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