TopK
产品支持情况
功能说明
获取最后一个维度的前k个最大值或最小值及其对应的索引。
如果输入是向量，则在向量中找到前k个最大值或最小值及其对应的索引；如果输入是矩阵，则沿最后一个维度计算每行中前k个最大值或最小值及其对应的索引。本接口最多支持输入为二维数据，不支持更高维度的输入。
如下图所示，对shape为(4, 32)的二维矩阵进行排序，k设置为1，输出结果为[[32] [32] [32] [32]]。
-
必备概念
基于如上样例，我们引入一些必备概念：行数称之为外轴长度（outter），每行实际的元素个数称之为内轴的实际长度（n）。本接口要求输入的内轴长度为32的整数倍，所以当n不是32的整数倍时，需要开发者将其向上补齐到32的整数倍，补齐后的长度称之为内轴长度（inner）。比如，如下的样例中，每行的实际长度n为31，不是32的整数倍，向上补齐后得到inner为32，图中的padding代表补齐操作。n和inner的关系如下：当n是32的整数倍时，inner=n；否则，inner > n。
-
接口模式
本接口支持两种模式：Normal模式和Small模式。Normal模式是通用模式；Small模式是为内轴长度固定为32（单位：元素个数）的场景提供的高性能模式。因为Small模式inner固定为32，可以进行更有针对性的处理，所以相关的约束较少，性能较高。内轴长度inner为32时建议使用Small模式。
-
附加功能：本接口支持开发者指定某些行的排序是无效排序。通过传入finishLocal参数值来控制，finishLocal对应行的值为true时，表示该行排序无效，此时排序后输出的dstIndexLocal的k个索引值会全部被置为无效索引n。
实现原理
-
MERGE_SORT算法
以float类型，ND格式，shape为[outter, inner]的输入Tensor为例，描述TopK高阶API内部算法框图，如下图所示。
根据TopKMode不同的模式选择，可分为两个分支。
-
计算TopK NORMAL模式，过程如下：
-
模板参数isInitIndex为false，需生成0到inner - 1的索引；
Atlas A3 训练系列产品/Atlas A3 推理系列产品采用方式二。
Atlas A2 训练系列产品/Atlas A2 推理系列产品采用方式二。
Kirin X90采用方式一。
Kirin 9030采用方式一。
- 方式一：使用CreateVecIndex生成0到inner - 1的索引。
- 方式二：使用Arange生成0到inner - 1的索引。
-
isLargest参数为false，由于Sort32指令默认为降序排序，则给数据乘以-1；
-
对输入数据完成全排序。
Atlas A3 训练系列产品/Atlas A3 推理系列产品采用方式二。
Atlas A2 训练系列产品/Atlas A2 推理系列产品采用方式二。
Kirin X90采用方式一。
Kirin 9030采用方式一。
方式一：
使用高阶API Sort对数据完成全排序。
方式二：
- 使用Sort32对数据排序，保证每32个数据是有序的。
- 使用MrgSort指令对所有的已排序数据块归并排序。
-
使用GatherMask指令提取前k个数据和索引；
-
finishLocal[i]为true时，则更新该行对应的排序结果为无效索引n；
-
isLargest参数为false，则给数据乘以-1还原数据。
-
-
计算TopK SMALL模式，过程如下：
-
模板参数isInitIndex为false，需生成0到inner - 1的索引，并使用Copy指令将数据复制为outter条；
Atlas A3 训练系列产品/Atlas A3 推理系列产品采用方式二。
Atlas A2 训练系列产品/Atlas A2 推理系列产品采用方式二。
Kirin X90采用方式一。
Kirin 9030采用方式一。
- 方式一：使用CreateVecIndex生成0到inner - 1的索引。
- 方式二：使用Arange生成0到inner - 1的索引。
-
isLargest参数为false，由于Sort32指令默认为降序排序，则给输入数据乘以-1；
-
使用Sort32对数据排序；
-
使用GatherMask指令提取前k个数据和索引；
-
isLargest参数为false，则给输入数据乘以-1还原数据。
-
-
函数原型
-
API内部申请临时空间
template <typename T, bool isInitIndex = false, bool isHasfinish = false, bool isReuseSrc = false, enum TopKMode topkMode = TopKMode::TOPK_NORMAL> __aicore__ inline void TopK(const LocalTensor<T>& dstValueLocal, const LocalTensor<int32_t>& dstIndexLocal, const LocalTensor<T>& srcLocal, const LocalTensor<int32_t>& srcIndexLocal, const LocalTensor<bool>& finishLocal, const int32_t k, const TopkTiling& tilling, const TopKInfo& topKInfo, const bool isLargest = true) -
通过tmpLocal入参传入临时空间
template <typename T, bool isInitIndex = false, bool isHasfinish = false, bool isReuseSrc = false, enum TopKMode topkMode = TopKMode::TOPK_NORMAL> __aicore__ inline void TopK(const LocalTensor<T>& dstValueLocal, const LocalTensor<int32_t>& dstIndexLocal, const LocalTensor<T>& srcLocal, const LocalTensor<int32_t>& srcIndexLocal, const LocalTensor<bool>& finishLocal, const LocalTensor<uint8_t>& tmpLocal, const int32_t k, const TopkTiling& tilling, const TopKInfo& topKInfo, const bool isLargest = true)
由于该接口的内部实现中涉及复杂的逻辑计算，需要额外的临时空间来存储计算过程中的中间变量。临时空间支持API接口申请和开发者通过tmpLocal入参传入两种方式。
-
API接口内部申请临时空间，开发者无需申请，但是需要预留临时空间的大小。
-
通过tmpLocal入参传入，使用该tensor作为临时空间进行处理，API接口内部不再申请。该方式开发者可以自行管理tmpLocal内存空间，并在接口调用完成后，复用该部分内存，内存不会反复申请释放，灵活性较高，内存利用率也较高。临时空间大小tmpLocal的BufferSize的获取方式如下：通过TopK Tiling中提供的GetTopKMaxMinTmpSize接口获取所需最大和最小临时空间大小。
参数说明
表 1 模板参数说明
|
Atlas A3 训练系列产品/Atlas A3 推理系列产品，支持的数据类型为：half、float。 |
|
|
|
Topk接口支持开发者通过finishLocal参数来指定某些行的排序是无效排序。该模板参数用于控制是否启用上述功能，true表示启用，false表示不启用。 isHasfinish参数和finishLocal的配套使用方法请参考表2中的finishLocal参数说明。 |
|
enum class TopKMode {
TOPK_NORMAL, // Normal模式
TOPK_NSMALL, // Small模式
};
|
表 2 接口参数说明
|
类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。
|
||
|
类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。
|
||
|
类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 |
||
|
类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 该参数和模板参数isInitIndex配合使用，isInitIndex为false时，srcIndexLocal只需进行定义，不需要赋值，将定义后的srcIndexLocal传入接口即可；isInitIndex为true时，开发者需要通过srcIndexLocal参数传入索引值。srcIndexLocal参数设置的规则如下：
|
||
源操作数。用于指定某些行的排序是无效排序，其shape为(outter, 1)。 类型为LocalTensor，支持的TPosition为VECIN/VECCALC/VECOUT。 该参数和模板参数isHasfinish配合使用，Normal模式下支持isHasfinish配置为true/false，Small模式下仅支持isHasfinish配置为false。
|
||
临时空间。接口内部复杂计算时用于存储中间变量，由开发者提供，临时空间大小的获取方式请参考TopK Tiling。数据类型固定uint8_t。 类型为LocalTensor，逻辑位置仅支持VECCALC，不支持其他逻辑位置。 |
||
Topk计算所需Tiling信息，Tiling信息的获取请参考TopK Tiling。 |
||
srcLocal的shape信息。TopKInfo类型，具体定义如下： struct TopKInfo {
int32_t outter = 1; // 表示输入待排序数据的外轴长度
int32_t inner; // 表示输入待排序数据的内轴长度，inner必须是32的整数倍
int32_t n; // 表示输入待排序数据的内轴的实际长度
};
|
||
返回值说明
无
约束说明
- 操作数地址偏移对齐要求请参见通用说明和约束。
- 不支持源操作数与目的操作数地址重叠。
- 当存在srcLocal[i]与srcLocal[j]相同时，如果i>j，则srcLocal[j]将首先被选出来，排在前面。
- inf在Topk中被认为是极大值。
- nan在topk中排序时无论是降序还是升序，均被排在前面。
调用示例
本样例实现了Normal模式和Small模式的代码逻辑。
if (!tmpLocal) { // 是否通过tmpLocal入参传入临时空间
if (isSmallMode) { // Small模式
AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NSMALL>(dstLocalValue,
dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, k, topKTilingData, topKInfo, isLargest);
} else {
AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NORMAL>(dstLocalValue,
dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, k, topKTilingData, topKInfo, isLargest);
}
} else {
if (tmplocalBytes % 32 != 0) {
tmplocalBytes = (tmplocalBytes + 31) / 32 * 32;
}
pipe.InitBuffer(tmplocalBuf, tmplocalBytes);
AscendC::LocalTensor<uint8_t> tmplocalTensor = tmplocalBuf.Get<uint8_t>();
if (isSmallMode) {
AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NSMALL>(dstLocalValue,
dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, tmplocalTensor, k, topKTilingData, topKInfo, isLargest);
} else {
AscendC::TopK<T, isInitIndex, isHasfinish, isReuseSrc, AscendC::TopKMode::TOPK_NORMAL>(dstLocalValue,
dstLocalIndex, srcLocalValue, srcLocalIndex, srcLocalFinish, tmplocalTensor, k, topKTilingData, topKInfo, isLargest);
}
}
表 3 Normal模式的样例解析
表 4 Small模式的样例解析