算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载SelectV2 是 CANN ops-math 实验性数学算子库中实现张量条件选择torch.where / select 语义的核心算子它依据布尔条件张量逐元素地从两个输入张量中选择输出。本文以 experimental/math/select_v2/README.md 及其配套的 aclnnSWhere 接口文档 为主体完整覆盖算子的功能定义、参数规格、两段式 aclnn 调用流程与可运行示例并深入 op_api / op_host / op_kernel 各层源码解析广播、类型提升、AICore/AICpu 分流与 tiling 实现原理。读完本文你将能够在 Atlas A2 训练/推理系列产品上独立完成 SelectV2 算子的 aclnn 单算子调用并理解其底层执行机制。一、算子功能与计算公式SelectV2 实现的是逐元素的条件选择对于输入张量的每一个位置 $i$以条件张量 condition 在该位置的取值决定输出取自 self 还是 other其计算公式为$$ out_i \begin{cases} self_i, condition_i true \ other_i, condition_i false \end{cases} $$即文档中给出的等价形式$$ out_i if(condition_i), ?, self_i : other_i $$该语义与 Python/Torch 中常见的where(condition, self, other)一致广泛用于掩码选择、梯度分流、数据过滤等场景。在 CANN ops-math 中该算子同时提供了底层算子接口SelectV2与面向单算子调用的 aclnn 接口aclnnSWheres_where 即 select where 的命名。产品支持情况产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√从算子定义源码 select_v2_def.cpp 可以看到算子注册时通过this-AICore().AddConfig(ascend910b)声明了 AICore 实现所适配的芯片配置从源码结构看与文档声明的 Atlas A2 系列产品支持范围一致。二、参数说明输入/输出规格SelectV2 共包含 3 个输入、1 个输出全部采用 ND 数据格式。各参数的规格如下表所示参数名输入/输出/属性描述数据类型数据格式condition输入待进行 select_v2 计算的入参公式中的 condition即布尔选择条件BOOLNDself输入待进行 select_v2 计算的入参公式中的 self条件为真时输出的取值来源FLOAT、FLOAT16、BFLOAT16、INT32、UINT32、INT16、UINT16、INT8、UINT8、BOOLNDother输入待进行 select_v2 计算的入参公式中的 other条件为假时输出的取值来源FLOAT、FLOAT16、BFLOAT16、INT32、UINT32、INT16、UINT16、INT8、UINT8、BOOLNDout输出待进行 select_v2 计算的出参公式中的 outFLOAT、FLOAT16、BFLOAT16、INT32、UINT32、INT16、UINT16、INT8、UINT8、BOOLND关键约定condition 仅支持 BOOL 类型其余输入与输出支持浮点、半精度含 BF16、各宽度有符号/无符号整型以及 BOOL 类型数据格式统一为 ND普通多维布局输出 out 的shape 与 self 相同而 condition、self、other 三者之间支持广播broadcast关系最终输出以三者广播后的形状为准这一点在接口层的CheckShape与CanBroadcast校验中实现详见第四节。约束说明文档声明该算子无特殊约束无。但需注意无约束并不等于无校验接口层仍会严格检查空指针、数据类型范围、维度上限与 shape 一致性具体错误场景见第三节的错误码说明。三、aclnn 两段式接口aclnnSWhere 调用说明CANN 单算子 API 采用两段式调用模式详见 两段式接口说明必须先调用第一段接口aclnnSWhereGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器 executor再调用第二段接口aclnnSWhere真正执行计算。aclnnSWhere的两段接口声明位于 aclnn_s_where.hdomain aclnn_math。3.1 函数原型aclnnStatus aclnnSWhereGetWorkspaceSize( const aclTensor *condition, const aclTensor *self, const aclTensor *other, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor);aclnnStatus aclnnSWhere( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream);说明workspace 是指除输入/输出之外算子在 NPU 上完成计算所需的临时内存workspaceSize 为其大小第二段接口aclnnSWhere(...)不可重复调用同一 executor 执行一次后即需重新获取。3.2 aclnnSWhereGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorcondition输入待进行 select_v2 计算的入参公式中的 condition无BOOLND0-8√self输入待进行 select_v2 计算的入参公式中的 self无FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√other输入待进行 select_v2 计算的入参公式中的 other无FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√out输出待进行 select_v2 计算的出参公式中的 outshape 与 self 相同FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----需要注意接口层aclnn API相比算子注册层README 参数表额外支持了 DOUBLE、INT64 等类型接口内部会通过类型提升与 Cast 将其转换到算子内核支持的范围内再计算见 4.2 节因此实际可用的输入类型以接口文档为准。三个输入及输出均支持 0~8 维 shape且支持非连续non-contiguousTensor——接口层会在内部通过Contiguous与ViewCopy完成连续化与结果回写。3.3 aclnnSWhere 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnSWhereGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream3.4 返回值与错误码两段接口均返回aclnnStatus状态码具体可参见 aclnn 返回码说明。第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针ACLNN_ERR_PARAM_INVALID161002condition、self 和 other 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002condition、self 和 other 的数据维度超过了 8 维ACLNN_ERR_PARAM_INVALID161002condition、self、other 和 out 的数据形状不一致四、调用示例完整的 aclnnSWhere 单算子样例官方提供的最小可运行样例位于 test_aclnn_s_where.cpp接口文档 aclnnSWhere.md 给出了完整注释版本编译与执行流程参考 编译与运行样例。以下示例使用 shape 为{4, 2}的输入演示从资源初始化、张量构造、两段式调用到结果回拷、资源释放的完整流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_s_where.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 shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 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初始化, 参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // check根据自己的需要处理 CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t selfShape {4, 2}; std::vectorint64_t otherShape {4, 2}; std::vectorint64_t conditionShape {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* otherDeviceAddr nullptr; void* conditionDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* other nullptr; aclTensor* condition nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 0, 0, 0, 0, 0, 0, 7}; std::vectorfloat otherHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorint8_t conditionHostData {false, false, false, false, true, true, true, true}; std::vectorfloat outHostData {10, 10, 10, 10, 10, 10, 10, 10}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建other aclTensor ret CreateAclTensor(otherHostData, otherShape, otherDeviceAddr, aclDataType::ACL_FLOAT, other); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建condition aclTensor ret CreateAclTensor(conditionHostData, conditionShape, conditionDeviceAddr, aclDataType::ACL_BOOL, condition); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnSWhere第一段接口 ret aclnnSWhereGetWorkspaceSize(condition, self, other, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSWhereGetWorkspaceSize 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;); } // 调用aclnnSWhere第二段接口 ret aclnnSWhere(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSWhere 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侧 auto size GetShapeSize(outShape); 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 aclDestroyTensor(self); aclDestroyTensor(other); aclDestroyTensor(condition); aclDestroyTensor(out); // 7. 释放device资源 aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(conditionDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }以示例数据为例condition 前 4 个元素为false、后 4 个为true因此输出前 4 个元素取自 other0, 1, 2, 3后 4 个元素取自 self0, 0, 0, 7即最终结果应为0, 1, 2, 3, 0, 0, 0, 7可在程序打印的result[i]中核对。五、源码级实现解析从 API 到 NPU 内核SelectV2 算子仓库内按标准 CANN 算子分层组织各层职责如下目录关键文件职责op_apiaclnn_s_where.cpp、select_v2.cpp两段式 aclnn 接口、参数校验、类型提升与执行器构图l0op 层op_hostselect_v2_def.cpp、select_v2_tiling.cpp算子信息注册OpDef与 tiling 计算分核、分块op_kernelselect_v2.cpp、select_v2.hAscendC 内核实现数据搬运、Select 计算指令tests/utop_api / op_host / op_kernel 三组测试接口、tiling、内核的单元测试与数据生成5.1 参数校验与广播推导op_api 层aclnnSWhereGetWorkspaceSizeaclnn_s_where.cpp的执行路径依次为创建 executorCREATE_EXECUTOR()固定写法失败返回ACLNN_ERR_INNER_CREATE_EXECUTOR参数校验CheckParams先CheckNotNull4Tensor判空对应错误码 161001再CheckDtypeValid校验类型对应 161002最后CheckShape校验维度上限与广播 shape对应 161002空 tensor 短路若 self/condition/other 任一为空 tensor直接返回workspaceSize 0非连续化处理对三个输入分别调用l0op::Contiguous转为连续 tensor接口文档注明支持非连续 Tensor 的原因所在类型提升与 Castcondition 统一Cast为DT_BOOLself 与 other 通过op::PromoteType推导公共提升类型后各自 Cast若两者均为 BOOL则借助 View 将数据类型视为DT_INT8参与计算再把结果 View 回 BOOL最终结果再Cast回 out 的目标类型执行器构图与 workspace 计算调用l0op::SelectV2将算子任务加入队列以ViewCopy回写非连续输出最后通过uniqueExecutor-GetWorkspaceSize()汇总整个计算图所需的 workspace 大小并将 executor 转移给调用方。第二段接口aclnnSWhere的实现非常精简aclnn_s_where.cpp仅调用框架统一的CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成异步下发执行。5.2 AICore/AICpu 双路径分发l0op 层select_v2.cpp 中的l0op::SelectV2是实际的构图入口关键逻辑包括广播推导CanBroadcast依次对condition self、other 结果做两次BroadcastInferShape得到最终广播形状broadcastShape并用executor-AllocTensor(broadcastShape, x1-GetDataType())分配输出AICore 判定IsAiCoreSupport依据芯片类型与 dtype 决定走 AICore 还是 AICpu 实现普通平台走 AICore 需要 dtype 属于{FLOAT, FLOAT16, INT32, INT8, UINT8, BF16, INT64}ASCEND610LITE610 Lite 平台仅支持{FLOAT, FLOAT16, INT32, INT8, UINT8}任务入队AICore 路径用ADD_TO_LAUNCHER_LIST_AICORE(SelectV2, OP_INPUT(condition, self, other), OP_OUTPUT(selectOut))AICpu 路径用ADD_TO_LAUNCHER_LIST_AICPU配合internal::AicpuTaskSpace space(SelectV2)注册任务空间。5.3 算子信息注册op_host 层select_v2_def.cpp 通过OpDef注册算子原型condition输入仅ge::DT_BOOL10 个占位项对应 10 种 self 类型的组合格式统一FORMAT_NDself/other输入与out输出{DT_FLOAT16, DT_FLOAT, DT_BF16, DT_INT8, DT_UINT8, DT_BOOL, DT_INT16, DT_UINT16, DT_INT32, DT_UINT32}共 10 种 dtype均通过UnknownShapeFormat声明动态 shape 场景仍保持 ND最后以AICore().AddConfig(ascend910b)声明 AICore 适配配置OP_ADD(SelectV2)加入算子信息库。该注册列表与 README 参数表中的数据类型完全对应是参数表内容的源码级印证。5.4 Tiling 策略op_host 层select_v2_tiling.cpp 实现了基于数据量的均衡分核分块策略核心步骤平台信息获取从platform_ascendc::PlatformAscendC读取 UB 大小与核数coreNumUB 容量规划固定预留RESERVED_UB_SIZE 8KB块大小BLOCK_SIZE 32B按 dtype 区分分块粒度——DT_FLOAT/INT32/UINT32每块容纳UB_DATA_NUM_UINT8 16个 32B 块其余类型为UB_DATA_NUM_OTHER 8个核数决策若单块 tile 即可容纳全部数据则coreNum 1否则在核数上限与inputLengthAlgin32 / BLOCK_SIZE按 32B 对齐后的块数之间取小保证每个核至少处理 32B 数据负载均衡CalculateCoreBlockNums计算smallCoreDataNum / bigCoreDataNum前tailBlockNum个核多分一块、smallTailDataNum / bigTailDataNum最后一个 tile 的尾数以及finalSmallTileNum / finalBigTileNum每核 tile 数将结果写入SelectV2TilingDatatiling key 与 block dimcontext-SetTilingKey(GET_TPL_TILING_KEY(ELEMENTWISE_TPL_SCH_MODE_0))context-SetBlockDim(coreNum)下发分核数。SelectV2TilingData的 8 个字段定义在 select_v2_tiling_data.h。5.5 AscendC 内核实现op_kernel 层内核模板KernelSelectV2TYPE_X, TYPE_Yselect_v2.h采用标准流水结构Init根据核号与 tiling 数据计算本核的起始偏移globalBufferIndex与数据量初始化三路队列inQueueCcondition、inQueueXself/other 合并进同一块 buffer以tileDataNum为界、outQueueY输出CopyInDataCopy将 condition、self、other 三段数据搬入 Vector 侧队列Compute核心计算分三类 dtype 分支32bit 类float/uint32_t/int32_tcondition 经Cast转 fp16 后与 0 做CompareScalarNE 模式生成选择掩码再调用AscendC::Select(..., SELMODE::VSEL_TENSOR_TENSOR_MODE, ...)完成逐元素选择16bit 类half/bfloat16_t/uint16_t/int16_t同样经 CompareScalar Select 处理8bit/BOOL 类uint8_t/int8_t/bool先将两个输入 Cast 到 fp16 临时 bufferSelect 后再 Cast 回原类型规避向量指令对 8bit 类型的选择限制CopyOutDataCopy将结果写回 GM整体按tileNum次循环最后一次处理tailDataNum尾块。模板入口 select_v2.cpp 通过REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT从 GM 侧 tiling 地址反序列化SelectV2TilingData随后op.Init(...)op.Process()启动流水。六、测试验证仓库为该算子提供了三层单元测试可用于验证接口正确性与 tiling 结果op_api 层test_aclnn_select_v2.cpp 使用OP_API_UT(aclnnSWhere, INPUT(...), OUTPUT(...))框架在ASCEND910B平台下构造{2, 2}的 BOOL condition 与 FLOAT self/other验证TestGetWorkspaceSize返回ACLNN_SUCCESSop_host 层test_select_v2_tiling.cpp 通过TilingContextPara模拟{1,64,2,64}的 FP16 输入断言 tiling key 为 0、tiling 数据8192 8208 1 1 15856 8192 8208 0以及 workspace 大小16777216可直接验证 5.4 节 tiling 策略的计算结果op_kernel 层select_v2_data 目录 提供gen_data.py输入数据生成与compare_data.py结果比对配合test_select_v2.cpp完成内核级精度验证。七、贡献说明贡献者贡献方贡献算子贡献时间贡献内容infinity个人开发者SelectV22025/11/25SelectV2 算子适配开源仓小结SelectV2 作为 CANN ops-math 中条件选择类算子的代表对外以aclnnSWhere两段式接口提供简洁易用的单算子调用能力支持广播、类型提升与非连续 Tensor对内则通过 op_api → op_host → op_kernel 三层结构将参数校验、AICore/AICpu 分流、均衡 tiling 与 AscendC 向量化Select计算完整落地。开发者可直接复用 test_aclnn_s_where.cpp 作为模板接入自己的业务也可借助各层单元测试深入理解算子的分核分块与精度实现细节。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 算子解读MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析CANN ops math 算子解读MaskedSelectV3 掩码选择的 aclnn 两段式接口与 NPU 实现剖析 MaskedSelectV3 是 C算子库人工智能CANNCANN ops-math TruncateDiv 算子实战指南截断除法语义、aclnn 两段式调用与 NPU 源码实现解析CANN ops math TruncateDiv 算子实战指南截断除法语义、aclnn 两段式调用与 NPU 源码实现解析 TruncateDiv 是 CA算子库人工智能CANNCANN ops-math LogicalNot 算子aclnn 两段式调用接口、实现原理与源码深度解析CANN ops math LogicalNot 算子aclnn 两段式调用接口、实现原理与源码深度解析 本篇技术指南围绕 CANN 开源算子库 ops ma算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考