CANN ops-math 算子解读：MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析

CANN ops-math 算子解读：MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

MaskedSelectV3 是 CANN ops-math 数学算子库中用于"按布尔掩码挑选元素"的一维选择算子：它根据布尔掩码张量 mask 中为 True 的位置，从输入张量 self 中取出对应元素，输出一个一维张量。本指南以 算子 README 与 aclnnMaskedSelect 接口文档 为主体，结合仓库内 op_api、op_host、op_kernel 源码，完整讲解其产品支持情况、参数约束、两段式调用流程、源码级实现原理与可运行的完整样例，读者读完即可在 Ascend 系列产品上独立完成 MaskedSelectV3 的接入、编译、运行与结果验证。

一、算子功能与计算语义

1.1 功能说明

MaskedSelectV3 算子根据一个布尔掩码张量（mask）中的值选择输入张量（self）中的元素作为输出，形成一个新的一维张量。其行为等价于 PyTorch 中torch.masked_select：输出只保留 mask 对应位置为 True 的 self 元素，并按元素在广播后张量中的顺序排列。

计算公式如下：

$$ out = \left[ self[i] \right]_{i \in \mathcal{I}}, \quad \text{where } \mathcal{I} = \left{ i \mid mask[i] = \text{True} \right} $$

即输出张量 out 中的第 k 个元素，对应所有满足mask[i] == True的索引 i 按从小到大排列后的第 k 个self[i]。

1.2 典型使用场景

掩码选择是神经网络与数据预处理中的高频操作，典型场景包括：

• 注意力机制：根据 mask 挑选参与计算的 token 或特征；
• 稀疏化处理：按条件过滤梯度、权重或中间特征，得到稠密化的一维结果；
• 数据清洗：按阈值或标签过滤样本；
• 条件索引：动态路由、专家选择（MoE）中按门控条件抽取数据。

由于输出元素个数（True 的个数）在运行前不可预知，该算子属于输出 shape 依赖计算的动态 shape 算子，这一点在 算子定义源码 中以OutputShapeDependOnCompute()显式声明。

二、产品支持情况

MaskedSelectV3 在不同硬件平台上的支持情况如下表所示（来源：算子 README）：

需要特别说明的差异点：

• Kirin X90 / Kirin 9030 处理器系列产品：不支持 BFLOAT16 数据类型。
• 在 aclnn 接口层面（aclnnMaskedSelect 文档）：Atlas 推理系列产品、Atlas 训练系列产品同样不支持 BFLOAT16；而 Ascend 950PR/Ascend 950DT 的 self 与 out 支持全量数据类型（BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、UINT8、INT16、UINT16、INT32、UINT32、INT64、UINT64、BOOL）。

上述平台差异在源码中体现为两套 dtype 支持列表与芯片判断逻辑：aclnn_masked_select.cpp 定义了三个列表SELF_DTYPE_SUPPORT_LIST_NOT_SUPPORT_BF16、SELF_DTYPE_SUPPORT_LIST_SUPPORT_BF16、SELF_DTYPE_SUPPORT_LIST_SUPPORT_REGBASE，并由CheckSocVersionIsSupportBf16()依据当前 SoC 版本（IsRegBase()或ASCEND910B <= SocVersion <= ASCEND910E）动态选择。

三、参数说明

3.1 算子层参数（来源：算子 README）

补充说明：

• Kirin X90/Kirin 9030 处理器系列产品：不支持 BFLOAT16。
• 在 aclnn 接口层，self 与 out 的数据类型列表为：BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL（Atlas 推理/训练系列不含 BFLOAT16，详见上文）；mask 在接口层支持UINT8、BOOL两种类型，其中当 mask 数据类型为 UINT8 时，值只能是 0 或 1。
• 三个张量的维度（shape）范围均为 0–8 维。
• self 与 mask 均支持非连续（non-contiguous）Tensor；out 不支持非连续 Tensor。

3.2 各 dtype 的算子定义注册

在算子定义源码 masked_select_v3_def.cpp 中，MaskedSelectV3 的输入x、掩码mask与输出y均声明为 ND 格式、REQUIRED 参数；mask的数据类型在算子定义层为 BOOL。同时按芯片注册了不同 AICore 配置：

• ascend910b、ascend910_93：默认 AICore 配置；
• ascend950、ascend350：启用动态编译、动态 format、动态 rank、动态 shape 支持（DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)等）；
• kirinx90、kirin9030：使用GetKirinCoreConfig()配置，除动态能力外还额外开启PrecisionReduceFlag(true)，且其 x/y 数据类型列表不含 BFLOAT16，与 README 的平台约束一致。

芯片对应的二进制配置文件位于 op_host/config/ 下，按平台目录（ascend350、ascend910_93、ascend910b、ascend950、kirin9030、kirinx90）分别维护masked_select_v3_binary.json与masked_select_v3_simplified_key.ini。以 ascend950 的二进制配置 为例，它按 dtype 逐个声明bin_filename（如MaskedSelectV3_bfloat16、MaskedSelectV3_float32、MaskedSelectV3_int64等 13 种），每个条目声明输入 x、mask 与输出 y 的 dtype、ND 格式与动态 shape（-2表示 rank 动态）。

四、约束说明

• self 和 mask 的 shape 必须能够进行broadcast 操作；out 的 shape 必须是一维，且元素个数等于 self 和 mask 广播后的维度大小。
• 确定性计算：aclnnMaskedSelect 为默认确定性实现，多次运行同一输入可得到一致结果。
• broadcast 场景：功能支持，但不保证性能。从源码看，broadcast 场景是通过在计算图中额外插入 BroadcastTo 算子将 self 与 mask 先广播到同一 shape 后处理（见 aclnn_masked_select.cpp），该行为与接口文档"broadcast 场景通过额外插入 BroadcastTo 算子解决"的描述一致。
• 输入维度上限为 8 维（源码常量MAX_DIM_LEN = 8，见 aclnn_masked_select.cpp）。

五、aclnnMaskedSelect 两段式接口调用

5.1 接口背景：两段式 API

MaskedSelectV3 对外以aclnnMaskedSelect形式暴露，遵循 CANN 算子的两段式接口约定（详见 两段式 API 说明）：

1. 先调用aclnnMaskedSelectGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器 executor；
2. 再调用aclnnMaskedSelect执行计算。

5.2 函数原型

第一段接口：

aclnnStatus aclnnMaskedSelectGetWorkspaceSize( const aclTensor* self, const aclTensor* mask, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)

第二段接口：

aclnnStatus aclnnMaskedSelect( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

5.3 aclnnMaskedSelectGetWorkspaceSize 参数说明

5.4 第一段接口的返回值与错误码

aclnnStatus返回状态码，具体参见 aclnn 返回码。第一段接口完成入参校验，出现以下场景时报错：

这些校验在源码 CheckParams 中按顺序执行：先检查空指针（CheckNotNull），再检查 dtype 合法性（CheckDtypeValid），最后检查 shape 合法性（CheckShape，包括最大维度检查、broadcast 推断与一维约束、输出元素个数比对）。

5.5 aclnnMaskedSelect 参数说明

六、源码级实现剖析

6.1 aclnn 层的完整计算流程

aclnn_masked_select.cpp 顶部注释清晰刻画了算子内部的计算图：

self mask | | \ / Contiguous(workspace_0) Contiguous(workspace_1) \ / \ Cast(workspace_2) \ / MaskedSelect(workspace_3) | Cast(workspace_4) | ViewCopy | result

即：self 与 mask 各自先经Contiguous转为连续张量 → mask 经Cast转为 BOOL → self 经Cast转为 out 的数据类型 → 广播（如需）→ 执行 MaskedSelect/MaskedSelectV3 内核 → 结果 Cast → ViewCopy 得到最终输出。workspace 由各阶段累加得到，最终通过uniqueExecutor->GetWorkspaceSize()返回给调用方（源码 L253）。

6.2 AiCore 内核选择逻辑

源码IsAiCoreSupport()（aclnn_masked_select.cpp）根据芯片类型与 dtype 决定走哪条内核路径：

• RegBase 芯片（如 Kirin 系列）：dtype 在SELF_DTYPE_SUPPORT_LIST_SUPPORT_REGBASE列表内时，走l0op::MaskedSelectV3（对应MaskedSelectV3算子）；
• ASCEND910B ~ ASCEND910E 平台：dtype 在SUPPORT_BF16列表内时，同样走l0op::MaskedSelectV3；
• 其余场景走l0op::MaskedSelect兼容路径。

两条 l0op 入口均声明于 masked_select.h：MaskedSelect与MaskedSelectV3。

6.3 shape 推导（InferShape）

由于输出长度依赖 mask 中 True 的个数，masked_select_infershape.cpp 实现了动态 shape 推导：

• InferShape4MaskedSelect：将输出 shape 置为一维（rank = 1），维度值设为-1（未知，编译期无法确定）；
• InferShapeRange4MaskedSelect：给出输出元素个数的范围——最大值为输入总元素数（number_range->GetMax()->GetShapeSize()），最小值为 1。

这解释了为什么 out 必须由用户预先申请一块"按广播后 size 大小"的内存：实际有效元素数在 kernel 运行时通过shapeout写回（见下文 6.5）。

6.4 Tiling 切分策略

tiling 源码 展示了该算子在 NPU 上的并行切分思路，核心参数如下：

• BLOCK_SIZE = 256（字节）、DOUBLE_BUFFER = 2（双缓冲）、UB_USAGE = 0.65f（UB 使用率上限）；
• 按 dtype 计算单个元素字节数sizeOfDataType，据此计算一个 block 可容纳元素数ALIGN_NUM = 256 / sizeOfDataType，并将总长度向上对齐到 block 边界得到totalLengthAlignedWithBlock；
• 依据 ubSize、UB_USAGE 与数据宽度推导每次可处理的 block 数ubBlockNum（保持偶数以适配双缓冲），进而得到ubLength与总 block 数ubNum；
• 实际运行核数numBlocks = min(ubNum, aivNum)，即以总工作量与可用 Vector 核数取小，通过tilingContext->SetBlockDim(numBlocks)下发；
• 数据切分为"前段大块 + 尾部小块"两段（formerNum/formerLength/formerTileNum...与tailNum/tailLength/tailTileNum...），每段再按 tile 细分并单独处理最后不完整 tile，最大限度填满 UB；
• tilingKey 直接取sizeOfDataType（1/2/4/8），与 kernel 端TILING_KEY_IS(1/2/4/8)的模板分支一一对应；
• workspace 大小计算为totalLengthAlignedWithBlock * sizeOfDataType + numBlocks * 64再加上框架系统 workspace。

另外TilingPrepareForMaskedSelectV3（tiling 源码 L225-L258）在编译期通过PlatformAscendC获取aivNum、ubSize、workSpaceSize与isRegbase信息并写入MaskedSelectV3CompileInfo（结构体定义见 masked_select_v3_tiling.h）；RegBase 平台走独立实现TilingForMaskedSelectV3IsRegbaseSocVersion。

6.5 Kernel 内核实现

kernel 实现 中KernelMaskedSelectV3<T>的核心机制值得关注：

1. 掩码位图化（GenerateMask）：把 mask 元素经Cast转为 half 后与1.0做CompareScalar，生成 bitMask，将"布尔判断"转化为向量指令；
2. GatherMask 压缩选取（GatherResult）：利用向量指令GatherMask按 bitMask 一次性从源数据中抽出有效元素，rsvdCnt记录本次实际抽出的元素个数；针对 1/2/4/8 字节类型分别特化处理（8 字节类型按 int32 拆分搬运）；
3. 多核归并：每个核把挑选结果写入用户 workspace 中的独立区域，随后在 workspace 头部维护offsetGlobal偏移表（blockIdx << 3，即OFFSET_SHIFT_BITS = 3），通过SyncAll()同步后，各核累加前面所有核的 offset 计算出自己结果在最终输出中的写入位置（Process 逻辑 L134-L164）；
4. shapeout 写回：最后一个核（blockIdx == numBlocks - 1）负责把[1, 总元素数]写入 shapeout，从而在运行时确定 out 的实际有效长度；
5. 数据搬运：8 字节类型使用DataCopyPadDoubleWord按 int32 双字搬运以规避对齐限制，其余类型使用DataCopyPad+DataCopyExtParams支持非对齐尾包。

内核入口 masked_select_v3.cpp 通过GET_TILING_DATA读取 tiling 数据，按TILING_KEY_IS(8/4/2/1)分别实例化uint64_t/uint32_t/uint16_t/uint8_t四种模板，从而支持 INT64/DOUBLE（8 字节）、FLOAT/INT32（4 字节）、FLOAT16/BF16/INT16（2 字节）、BOOL/INT8/UINT8（1 字节）等全部数据类型。

七、完整调用示例

以下示例来自仓库 examples/test_aclnn_masked_select.cpp，完整演示了从环境初始化、张量构造、两段式调用到结果回拷与资源释放的全流程（编译与运行样例详见 编译与运行样例）：

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_masked_select.h" #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vector<int64_t>& shape) { int64_t shapeSize = 1; for (auto i : shape) { shapeSize *= i; } return shapeSize; } int64_t GetTrueElementNum(const std::vector<int8_t>& data) { int64_t true_num = 0; for (auto i : data) { if (i!=0) { true_num ++; } } return true_num; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法，AscendCL资源初始化 auto ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); ret = aclrtSetDevice(deviceId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); return 0; } template <typename T> int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size = GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); // 计算连续tensor的strides std::vector<int64_t> strides(shape.size(), 1); for (int64_t i = shape.size() - 2; i >= 0; i--) { strides[i] = shape[i + 1] * strides[i + 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1.（固定写法）device/stream初始化，参考AscendCL对外接口列表 // 根据自己的实际device填写deviceId int32_t deviceId = 0; aclrtStream stream; auto ret = Init(deviceId, &stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); // 2.构造输入与输出，需要根据API的接口自定义构造 std::vector<int64_t> selfShape = {4, 2}; std::vector<int64_t> maskShape = {4, 2}; std::vector<int64_t> outShape = {8}; void* selfDeviceAddr = nullptr; void* maskDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor* self = nullptr; aclTensor* mask = nullptr; aclTensor* out = nullptr; std::vector<float> selfHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<int8_t> maskHostData = {false,false,false,false,true,true,true,true}; std::vector<float> outHostData={10,10,10,10,10,10,10,10}; // 创建self aclTensor（dtype为ACL_FLOAT，shape 4x2） ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建mask aclTensor（dtype为ACL_BOOL，shape 4x2） ret = CreateAclTensor(maskHostData, maskShape, &maskDeviceAddr, aclDataType::ACL_BOOL, &mask); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建out aclTensor（一维，长度按广播后size=8预分配） ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out); CHECK_RET(ret == ACL_SUCCESS, return ret); // 3.调用CANN算子库API，需要修改为具体的API名称 uint64_t workspaceSize = 0; aclOpExecutor* executor; // 调用aclnnMaskedSelect第一段接口，完成参数校验并计算workspace大小 ret = aclnnMaskedSelectGetWorkspaceSize(self, mask, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaskedSelectGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret); } // 调用aclnnMaskedSelect第二段接口，真正执行计算 ret = aclnnMaskedSelect(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMaskedSelect failed. ERROR: %d\n", ret); return ret); // 4.（固定写法）同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); // 5.获取输出的值，将device侧内存上的结果拷贝至host侧 // 实际有效元素个数 = mask中True的个数 auto size = GetTrueElementNum(maskHostData); std::vector<float> resultData(size, 0); ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size; i++) { LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]); } // 6.释放aclTensor和aclScalar aclDestroyTensor(self); aclDestroyTensor(mask); aclDestroyTensor(out); // 7.释放Device资源 aclrtFree(selfDeviceAddr); aclrtFree(maskDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

示例运行结果：self 为 4×2 的[0,1,2,3,4,5,6,7]，mask 后 4 个位置为 True，因此输出一维张量结果为[4, 5, 6, 7]，共 4 个元素。

代码中的关键点解读：

• outShape取{8}（广播后 size），因为实际输出长度（True 的个数）运行时才可知，必须按最大可能长度预分配；
• 结果回拷长度使用GetTrueElementNum(maskHostData)计算，与内核通过 shapeout 写回的实际有效长度对应；
• 两段式调用中，workspaceSize == 0时无需申请 workspace 内存，直接传空指针执行第二段接口亦可。

八、测试与验证

仓库为 MaskedSelectV3 提供了完整的多层级测试覆盖：

• aclnn 接口层 UT：tests/ut/op_api/test_aclnn_masked_select.cpp（共 948 行）使用 gtest 框架，覆盖多种 dtype（float、half、int 系列、uint8 等）与随机掩码（genRandomBoolVector、genRandomUint8Vector），并通过统计 mask 中 True 的个数与输出元素逐一比对完成正确性验证；
• shape 推导 UT：tests/ut/op_host/test_masked_select_v3_infershape.cpp 验证动态输出 shape 推导逻辑；
• tiling UT：tests/ut/op_host/test_masked_select_v3_tiling.cpp 与 tests/ut/op_host/arch35/test_masked_select_v3_tiling.cpp 验证切分参数的正确性；
• kernel UT：tests/ut/op_kernel/test_masked_select_v3.cpp 在核函数层面验证不同数据宽度（1/2/4/8 字节）模板实例的执行结果；
• ST 用例：tests/st/aclnnMaskedSelect/ 提供atk_aclnnMaskedSelect.json用例配置与 executor_aclnnMaskedSelect.py 执行脚本，另有 arch35 的 csv 用例 用于芯片侧的批量回归。

九、总结

MaskedSelectV3 是 ops-math 库中典型的"输出 shape 依赖计算"的数据选择算子，其工程实现完整覆盖了 aclnn 两段式接口（含参数校验、Contiguous/Cast/Broadcast 预处理与 workspace 计算）、动态 shape 推导、基于 Vector 核的块/尾两段式 tiling 切分、以及利用 GatherMask 位图压缩与多核偏移归并的高效 kernel 实现。开发者只需掌握本指南中的参数约束与两段式调用模式，即可在支持平台上稳定接入aclnnMaskedSelect，并通过仓库提供的示例与测试用例快速完成功能验证。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math