CANN ops-math 算子解读MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathMaskedSelectV3 是 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产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√需要特别说明的差异点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参数名输入/输出/属性描述数据类型数据格式self输入张量输入张量shape 需要与 mask 满足 broadcast 关系支持非连续的 Tensor。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BOOLNDmask输入张量布尔掩码张量shape 要和 self 满足 broadcast 关系支持非连续的 Tensor。BOOLNDout输出张量输出一维张量元素个数为 mask 和 self 广播后的维度大小不支持非连续的 Tensor。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BOOLND补充说明Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16。在 aclnn 接口层self 与 out 的数据类型列表为BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLAtlas 推理/训练系列不含 BFLOAT16详见上文mask 在接口层支持UINT8、BOOL两种类型其中当 mask 数据类型为 UINT8 时值只能是 0 或 1。三个张量的维度shape范围均为 0–8 维。self 与 mask 均支持非连续non-contiguousTensorout 不支持非连续 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 接口背景两段式 APIMaskedSelectV3 对外以aclnnMaskedSelect形式暴露遵循 CANN 算子的两段式接口约定详见 两段式 API 说明先调用aclnnMaskedSelectGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器 executor再调用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 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入功能说明中的输入张量 self。shape 需要与 mask 满足 broadcast 关系。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√maskaclTensor*输入功能说明中的布尔掩码张量 mask。shape 需要与 self 满足 broadcast 关系数据类型是 UINT8 时值只能是 0 或 1。UINT8、BOOLND0-8√outaclTensor*输出功能说明中的输出一维张量。shape 为一维且元素个数为 mask 和 self 广播后的 shape size。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----5.4 第一段接口的返回值与错误码aclnnStatus返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、mask、out 是空指针。ACLNN_ERR_PARAM_INVALID161002self 和 mask 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 mask 的 shape 无法做 broadcast。ACLNN_ERR_PARAM_INVALID161002out 的 shape 不是一维时。ACLNN_ERR_PARAM_INVALID161002out 的元素个数不等于 self 和 mask 广播后的 shape size 时。这些校验在源码 CheckParams 中按顺序执行先检查空指针CheckNotNull再检查 dtype 合法性CheckDtypeValid最后检查 shape 合法性CheckShape包括最大维度检查、broadcast 推断与一维约束、输出元素个数比对。5.5 aclnnMaskedSelect 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMaskedSelectGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。六、源码级实现剖析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.hMaskedSelect与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.65fUB 使用率上限按 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最大限度填满 UBtilingKey 直接取sizeOfDataType1/2/4/8与 kernel 端TILING_KEY_IS(1/2/4/8)的模板分支一一对应workspace 大小计算为totalLengthAlignedWithBlock * sizeOfDataType numBlocks * 64再加上框架系统 workspace。另外TilingPrepareForMaskedSelectV3tiling 源码 L225-L258在编译期通过PlatformAscendC获取aivNum、ubSize、workSpaceSize与isRegbase信息并写入MaskedSelectV3CompileInfo结构体定义见 masked_select_v3_tiling.hRegBase 平台走独立实现TilingForMaskedSelectV3IsRegbaseSocVersion。6.5 Kernel 内核实现kernel 实现 中KernelMaskedSelectV3T的核心机制值得关注掩码位图化GenerateMask把 mask 元素经Cast转为 half 后与1.0做CompareScalar生成 bitMask将布尔判断转化为向量指令GatherMask 压缩选取GatherResult利用向量指令GatherMask按 bitMask 一次性从源数据中抽出有效元素rsvdCnt记录本次实际抽出的元素个数针对 1/2/4/8 字节类型分别特化处理8 字节类型按 int32 拆分搬运多核归并每个核把挑选结果写入用户 workspace 中的独立区域随后在 workspace 头部维护offsetGlobal偏移表blockIdx 3即OFFSET_SHIFT_BITS 3通过SyncAll()同步后各核累加前面所有核的 offset 计算出自己结果在最终输出中的写入位置Process 逻辑 L134-L164shapeout 写回最后一个核blockIdx numBlocks - 1负责把[1, 总元素数]写入 shapeout从而在运行时确定 out 的实际有效长度数据搬运8 字节类型使用DataCopyPadDoubleWord按 int32 双字搬运以规避对齐限制其余类型使用DataCopyPadDataCopyExtParams支持非对齐尾包。内核入口 masked_select_v3.cpp 通过GET_TILING_DATA读取 tiling 数据按TILING_KEY_IS(8/4/2/1)分别实例化uint64_t/uint32_t/uint16_t/uint8_t四种模板从而支持 INT64/DOUBLE8 字节、FLOAT/INT324 字节、FLOAT16/BF16/INT162 字节、BOOL/INT8/UINT81 字节等全部数据类型。七、完整调用示例以下示例来自仓库 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::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int64_t GetTrueElementNum(const std::vectorint8_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::vectorT hostData, const std::vectorint64_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::vectorint64_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::vectorint64_t selfShape {4, 2}; std::vectorint64_t maskShape {4, 2}; std::vectorint64_t outShape {8}; void* selfDeviceAddr nullptr; void* maskDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* mask nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorint8_t maskHostData {false,false,false,false,true,true,true,true}; std::vectorfloat outHostData{10,10,10,10,10,10,10,10}; // 创建self aclTensordtype为ACL_FLOATshape 4x2 ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建mask aclTensordtype为ACL_BOOLshape 4x2 ret CreateAclTensor(maskHostData, maskShape, maskDeviceAddr, aclDataType::ACL_BOOL, mask); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor一维长度按广播后size8预分配 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::vectorfloat 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 接口层 UTtests/ut/op_api/test_aclnn_masked_select.cpp共 948 行使用 gtest 框架覆盖多种 dtypefloat、half、int 系列、uint8 等与随机掩码genRandomBoolVector、genRandomUint8Vector并通过统计 mask 中 True 的个数与输出元素逐一比对完成正确性验证shape 推导 UTtests/ut/op_host/test_masked_select_v3_infershape.cpp 验证动态输出 shape 推导逻辑tiling UTtests/ut/op_host/test_masked_select_v3_tiling.cpp 与 tests/ut/op_host/arch35/test_masked_select_v3_tiling.cpp 验证切分参数的正确性kernel UTtests/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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN ops-math 算子解读：MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析
CANN ops-math 算子解读MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathMaskedSelectV3 是 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产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√需要特别说明的差异点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参数名输入/输出/属性描述数据类型数据格式self输入张量输入张量shape 需要与 mask 满足 broadcast 关系支持非连续的 Tensor。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BOOLNDmask输入张量布尔掩码张量shape 要和 self 满足 broadcast 关系支持非连续的 Tensor。BOOLNDout输出张量输出一维张量元素个数为 mask 和 self 广播后的维度大小不支持非连续的 Tensor。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BOOLND补充说明Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16。在 aclnn 接口层self 与 out 的数据类型列表为BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLAtlas 推理/训练系列不含 BFLOAT16详见上文mask 在接口层支持UINT8、BOOL两种类型其中当 mask 数据类型为 UINT8 时值只能是 0 或 1。三个张量的维度shape范围均为 0–8 维。self 与 mask 均支持非连续non-contiguousTensorout 不支持非连续 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 接口背景两段式 APIMaskedSelectV3 对外以aclnnMaskedSelect形式暴露遵循 CANN 算子的两段式接口约定详见 两段式 API 说明先调用aclnnMaskedSelectGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器 executor再调用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 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入功能说明中的输入张量 self。shape 需要与 mask 满足 broadcast 关系。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√maskaclTensor*输入功能说明中的布尔掩码张量 mask。shape 需要与 self 满足 broadcast 关系数据类型是 UINT8 时值只能是 0 或 1。UINT8、BOOLND0-8√outaclTensor*输出功能说明中的输出一维张量。shape 为一维且元素个数为 mask 和 self 广播后的 shape size。BFLOAT16、FLOAT16、FLOAT32、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----5.4 第一段接口的返回值与错误码aclnnStatus返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、mask、out 是空指针。ACLNN_ERR_PARAM_INVALID161002self 和 mask 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 mask 的 shape 无法做 broadcast。ACLNN_ERR_PARAM_INVALID161002out 的 shape 不是一维时。ACLNN_ERR_PARAM_INVALID161002out 的元素个数不等于 self 和 mask 广播后的 shape size 时。这些校验在源码 CheckParams 中按顺序执行先检查空指针CheckNotNull再检查 dtype 合法性CheckDtypeValid最后检查 shape 合法性CheckShape包括最大维度检查、broadcast 推断与一维约束、输出元素个数比对。5.5 aclnnMaskedSelect 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMaskedSelectGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。六、源码级实现剖析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.hMaskedSelect与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.65fUB 使用率上限按 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最大限度填满 UBtilingKey 直接取sizeOfDataType1/2/4/8与 kernel 端TILING_KEY_IS(1/2/4/8)的模板分支一一对应workspace 大小计算为totalLengthAlignedWithBlock * sizeOfDataType numBlocks * 64再加上框架系统 workspace。另外TilingPrepareForMaskedSelectV3tiling 源码 L225-L258在编译期通过PlatformAscendC获取aivNum、ubSize、workSpaceSize与isRegbase信息并写入MaskedSelectV3CompileInfo结构体定义见 masked_select_v3_tiling.hRegBase 平台走独立实现TilingForMaskedSelectV3IsRegbaseSocVersion。6.5 Kernel 内核实现kernel 实现 中KernelMaskedSelectV3T的核心机制值得关注掩码位图化GenerateMask把 mask 元素经Cast转为 half 后与1.0做CompareScalar生成 bitMask将布尔判断转化为向量指令GatherMask 压缩选取GatherResult利用向量指令GatherMask按 bitMask 一次性从源数据中抽出有效元素rsvdCnt记录本次实际抽出的元素个数针对 1/2/4/8 字节类型分别特化处理8 字节类型按 int32 拆分搬运多核归并每个核把挑选结果写入用户 workspace 中的独立区域随后在 workspace 头部维护offsetGlobal偏移表blockIdx 3即OFFSET_SHIFT_BITS 3通过SyncAll()同步后各核累加前面所有核的 offset 计算出自己结果在最终输出中的写入位置Process 逻辑 L134-L164shapeout 写回最后一个核blockIdx numBlocks - 1负责把[1, 总元素数]写入 shapeout从而在运行时确定 out 的实际有效长度数据搬运8 字节类型使用DataCopyPadDoubleWord按 int32 双字搬运以规避对齐限制其余类型使用DataCopyPadDataCopyExtParams支持非对齐尾包。内核入口 masked_select_v3.cpp 通过GET_TILING_DATA读取 tiling 数据按TILING_KEY_IS(8/4/2/1)分别实例化uint64_t/uint32_t/uint16_t/uint8_t四种模板从而支持 INT64/DOUBLE8 字节、FLOAT/INT324 字节、FLOAT16/BF16/INT162 字节、BOOL/INT8/UINT81 字节等全部数据类型。七、完整调用示例以下示例来自仓库 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::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int64_t GetTrueElementNum(const std::vectorint8_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::vectorT hostData, const std::vectorint64_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::vectorint64_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::vectorint64_t selfShape {4, 2}; std::vectorint64_t maskShape {4, 2}; std::vectorint64_t outShape {8}; void* selfDeviceAddr nullptr; void* maskDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* mask nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorint8_t maskHostData {false,false,false,false,true,true,true,true}; std::vectorfloat outHostData{10,10,10,10,10,10,10,10}; // 创建self aclTensordtype为ACL_FLOATshape 4x2 ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建mask aclTensordtype为ACL_BOOLshape 4x2 ret CreateAclTensor(maskHostData, maskShape, maskDeviceAddr, aclDataType::ACL_BOOL, mask); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor一维长度按广播后size8预分配 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::vectorfloat 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 接口层 UTtests/ut/op_api/test_aclnn_masked_select.cpp共 948 行使用 gtest 框架覆盖多种 dtypefloat、half、int 系列、uint8 等与随机掩码genRandomBoolVector、genRandomUint8Vector并通过统计 mask 中 True 的个数与输出元素逐一比对完成正确性验证shape 推导 UTtests/ut/op_host/test_masked_select_v3_infershape.cpp 验证动态输出 shape 推导逻辑tiling UTtests/ut/op_host/test_masked_select_v3_tiling.cpp 与 tests/ut/op_host/arch35/test_masked_select_v3_tiling.cpp 验证切分参数的正确性kernel UTtests/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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
