算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载aclnnSWhere 是 CANN ops-math 数学算子库中 SelectV2 算子对外提供的 aclnn 单算子调用接口功能是依据条件张量逐元素从两个输入张量中选择对应位置的值。本文以 experimental/math/select_v2/docs/aclnnSWhere.md 为骨架完整展开接口原型、参数与返回值说明并深入 op_api、op_host、op_kernel 源码剖析两段式调用流程、入参校验、类型提升与 Kernel 计算原理帮助读者在 Atlas A2 系列产品上正确编写可运行的 aclnnSWhere 调用程序。产品支持情况产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√该能力在仓库内 SelectV2 算子级文档 experimental/math/select_v2/README.md 中有同样声明算子注册配置AICore().AddConfig(ascend910b)见 select_v2_def.cpp也与之对应。功能说明算子功能实现张量的条件选择根据条件张量从两个输入张量中选择对应位置的元素。计算公式$$ out_iif(condition_i)? self_i : other_i $$即对每一个位置i当condition[i]为真时取self[i]否则取other[i]。该语义与 PyTorch/TensorFlow 中常见的where条件选择一致可用于掩码选择、数据筛选等场景。两段式接口与函数原型aclnnSWhere 属于 CANN 单算子 API 中的两段式接口必须先调用第一段接口aclnnSWhereGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用第二段接口aclnnSWhere执行计算。两段式接口的通用约定参见 docs/zh/context/two_phase_api.md其中 workspace 指除输入/输出外算子在 NPU 上完成计算所需的临时内存且第二段接口不可重复调用。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)其中aclnn表示接口前缀SWhere表示 SelectV2 算子类型。两段式接口的具体实现位于 aclnn_s_where.cpp第一段接口完成参数校验、构图Contiguous/Cast/SelectV2/ViewCopy 等 L0 算子组合并计算出 workspace 大小第二段接口通过CommonOpExecutorRun驱动执行器在指定 Stream 上完成计算。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 计算的出参公式中的 out。shape 与 self 相同。FLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。-----executor输出返回 op 执行器包含了算子计算流程。-----说明算子级文档 README.md 中 self/other/out 列出的支持类型为 FLOAT、FLOAT16、BFLOAT16、INT32、UINT32、INT16、UINT16、INT8、UINT8、BOOLAPI 文档在此基础上额外覆盖 DOUBLE 与 INT64实际支持范围以接口校验代码为准。从 aclnn_s_where.cpp 源码看self/other/out的校验列表DTYPE_SUPPORT_910B_LIST等还包含 DOUBLE、INT64、UINT64 以及 COMPLEX64/COMPLEX128而condition仅支持 BOOL 与 UINT8CONDITION_DTYPE_SUPPORT其中 UINT8 会在后续构图阶段被 Cast 成 BOOL 再参与计算。返回值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 的数据形状不一致。上述校验在源码中由CheckParams驱动aclnn_s_where.cppCheckNotNull4Tensor检查四个 tensor 是否为空指针失败返回ACLNN_ERR_PARAM_NULLPTRCheckDtypeValid检查self/other/out是否落在支持数据类型列表内、condition是否落在CONDITION_DTYPE_SUPPORTBOOL、UINT8内失败返回ACLNN_ERR_PARAM_INVALIDCheckShape通过OP_CHECK_MAX_DIM限制最大 8 维源码中MAX_DIM 8并基于OP_CHECK_BROADCAST_AND_INFER_SHAPE推导self与other的广播 shape再将该 shape 与condition的 view shape 做广播推导最终要求与out的 shape 完全一致见 aclnn_s_where.cpp。aclnnSWhere 参数说明参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnSWhereGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值aclnnStatus返回状态码具体参见 aclnn 返回码。约束说明无。接口内部实现与计算流程虽然接口文档标注“约束说明无”但从源码可以还原出 aclnnSWhere 的完整计算流水线aclnn_s_where.cpp空 tensor 提前返回若self/condition/other任一为空 tensor直接置workspaceSize 0并返回成功无需下发计算连续性处理对self、condition、other分别调用l0op::Contiguous转换为连续 tensor因此 API 层支持非连续 tensor 输入condition 类型统一将 condition 通过l0op::Cast转为DT_BOOL类型提升与处理若self与other均为 BOOL则创建数据类型为 INT8 的 View执行 SelectV2 后再将结果 View 的数据类型恢复为 BOOL最后经ViewCopy写入out其余情况先按op::PromoteType得到的提升类型分别 Castself、other执行 SelectV2再 Cast 回out的目标数据类型最后ViewCopy到outworkspace 计算通过uniqueExecutor-GetWorkspaceSize()汇总整个构图含 Contiguous/Cast/SelectV2/ViewCopy所需临时内存输出workspaceSize并将执行器通过ReleaseTo转移给调用方。Tiling 与 Kernel 侧原理Tiling 计算select_v2_tiling.cppTiling 阶段通过PlatformAscendC获取 UB 大小与核数UB 侧固定预留 8KBRESERVED_UB_SIZE按数据类型区分 UB 分块数量8 字节类型对应UB_DATA_NUM_UINT8 16其他类型为 8以 32B 为基本块粒度BLOCK_SIZE 32计算单核数据量、tile 数与尾块数并将smallCoreDataNum/bigCoreDataNum/tileDataNum/smallTailDataNum/bigTailDataNum/finalSmallTileNum/finalBigTileNum/tailBlockNum写入SelectV2TilingData结构定义见 select_v2_tiling_data.h同时设置 BlockDim 与 TilingKey。Kernel 计算select_v2.hKernel 采用 AscendC 编程模型使用双缓冲TQueBUFFER_NUM 2流水化 CopyIn/Compute/CopyOut。Compute 的核心思路是把 condition 先Cast为 fp16再用CompareScalar与 0 比较生成 mask最后调用向量Select指令SELMODE::VSEL_TENSOR_TENSOR_MODE完成逐元素选择对于 INT8/UINT8/BOOL 等窄类型则先 Cast 到 fp16 做 Select 再 Cast 回原类型。算子入口 select_v2.cpp 通过REGISTER_TILING_DEFAULT获取 tiling 数据后实例化KernelSelectV2并执行Process()。调用示例示例代码如下仅供参考具体编译和执行过程请参考 编译与运行样例。仓库内对应可编译样例为 examples/test_aclnn_s_where.cpp。#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侧需要根据具体API的接口定义修改 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和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(other); aclDestroyTensor(condition); aclDestroyTensor(out); // 7.释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(conditionDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例预期结果示例输入为 shape{4, 2}的三路张量self {0, 0, 0, 0, 0, 0, 0, 7}other {0, 1, 2, 3, 4, 5, 6, 7}condition {false, false, false, false, true, true, true, true}根据out_i if(condition_i) ? self_i : other_i前 4 个位置 condition 为 false 取 other 的值后 4 个位置 condition 为 true 取 self 的值输出应为{0, 1, 2, 3, 0, 0, 0, 7}。测试用例验证仓库内为 SelectV2 提供了多级测试佐证接口行为API 层用例tests/ut/op_api/test_aclnn_select_v2.cpp覆盖 aclnnSWhere 两段式接口的调用与结果校验Tiling 层用例tests/ut/op_host/test_select_v2_tiling.cpp例如 fp16 用例期望 tiling 数据为8192 8208 1 1 15856 8192 8208 0对应smallCoreDataNum/bigCoreDataNum/finalBigTileNum/finalSmallTileNum/tileDataNum/smallTailDataNum/bigTailDataNum/tailBlockNum并校验 workspace 大小Kernel 层用例tests/ut/op_kernel 目录下提供gen_data.py/compare_data.py数据生成与比对脚本配合test_select_v2.cpp验证 Kernel 计算结果。总结aclnnSWhere 是 CANN ops-math 中 SelectV2 条件选择算子的标准两段式 aclnn 接口支持 ND 格式、0-8 维、非连续 tensor并自动完成 Contiguous、类型提升与结果 Cast 等预处理。通过第一段接口校验参数并获取 workspace再以第二段接口在指定 Stream 上异步执行底层由 Tiling 计算多核切分参数Kernel 以 fp16 中间精度 CompareScalar/Select 向量指令完成逐元素条件选择。结合 调用示例 与仓库内各级测试开发者可以快速在 Atlas A2 系列产品上落地该算子的实际调用与验证。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 算子实战aclnnBincount 两段式接口详解与源码剖析CANN ops math 算子实战aclnnBincount 两段式接口详解与源码剖析 本篇技术指南围绕 CANN ops math 数学算子库中的 Bin算子库人工智能CANNCANN ops-math 算子开发实战aclnnInplaceMaskedFillScalar 两段式接口详解与源码剖析CANN ops math 算子开发实战aclnnInplaceMaskedFillScalar 两段式接口详解与源码剖析 本文聚焦 CANN ops mat算子库人工智能CANNCANN ops-math 算子开发指南aclnnBitwiseNot 两段式接口详解与 AscendC 实现剖析CANN ops math 算子开发指南aclnnBitwiseNot 两段式接口详解与 AscendC 实现剖析 本指南以 CANN ops math 开源算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考