企业官网建设流程全解析

• 算子库
• 人工智能
• CANN

aclnnSWhere 是 CANN ops-math 数学算子库中 SelectV2 算子对外提供的 aclnn 单算子调用接口，功能是依据条件张量逐元素从两个输入张量中选择对应位置的值。本文以 experimental/math/select_v2/docs/aclnnSWhere.md 为骨架，完整展开接口原型、参数与返回值说明，并深入 op_api、op_host、op_kernel 源码，剖析两段式调用流程、入参校验、类型提升与 Kernel 计算原理，帮助读者在 Atlas A2 系列产品上正确编写可运行的 aclnnSWhere 调用程序。

产品支持情况

该能力在仓库内 SelectV2 算子级文档 experimental/math/select_v2/README.md 中有同样声明，算子注册配置（AICore().AddConfig("ascend910b")，见 select_v2_def.cpp）也与之对应。

功能说明

• 算子功能：实现张量的条件选择，根据条件张量从两个输入张量中选择对应位置的元素。
• 计算公式：

$$ out_i=if(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 参数说明

参数说明

说明：算子级文档 README.md 中 self/other/out 列出的支持类型为 FLOAT、FLOAT16、BFLOAT16、INT32、UINT32、INT16、UINT16、INT8、UINT8、BOOL，API 文档在此基础上额外覆盖 DOUBLE 与 INT64；实际支持范围以接口校验代码为准。从 aclnn_s_where.cpp 源码看，self/other/out的校验列表（DTYPE_SUPPORT_910B_LIST等）还包含 DOUBLE、INT64、UINT64 以及 COMPLEX64/COMPLEX128，而condition仅支持 BOOL 与 UINT8（CONDITION_DTYPE_SUPPORT），其中 UINT8 会在后续构图阶段被 Cast 成 BOOL 再参与计算。

返回值

aclnnStatus：返回状态码，具体参见 aclnn 返回码。

第一段接口会完成入参校验，出现以下场景时报错：

上述校验在源码中由CheckParams驱动（aclnn_s_where.cpp）：

1. CheckNotNull4Tensor检查四个 tensor 是否为空指针，失败返回ACLNN_ERR_PARAM_NULLPTR；
2. CheckDtypeValid检查self/other/out是否落在支持数据类型列表内、condition是否落在CONDITION_DTYPE_SUPPORT（BOOL、UINT8）内，失败返回ACLNN_ERR_PARAM_INVALID；
3. CheckShape通过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 参数说明

参数说明

返回值

aclnnStatus：返回状态码，具体参见 aclnn 返回码。

约束说明

无。

接口内部实现与计算流程

虽然接口文档标注“约束说明：无”，但从源码可以还原出 aclnnSWhere 的完整计算流水线（aclnn_s_where.cpp）：

1. 空 tensor 提前返回：若self/condition/other任一为空 tensor，直接置workspaceSize = 0并返回成功，无需下发计算；
2. 连续性处理：对self、condition、other分别调用l0op::Contiguous转换为连续 tensor，因此 API 层支持非连续 tensor 输入；
3. condition 类型统一：将 condition 通过l0op::Cast转为DT_BOOL；
4. 类型提升与处理：
   • 若self与other均为 BOOL，则创建数据类型为 INT8 的 View，执行 SelectV2 后再将结果 View 的数据类型恢复为 BOOL，最后经ViewCopy写入out；
   • 其余情况先按op::PromoteType得到的提升类型分别 Castself、other，执行 SelectV2，再 Cast 回out的目标数据类型，最后ViewCopy到out；
5. workspace 计算：通过uniqueExecutor->GetWorkspaceSize()汇总整个构图（含 Contiguous/Cast/SelectV2/ViewCopy）所需临时内存，输出workspaceSize，并将执行器通过ReleaseTo转移给调用方。

Tiling 与 Kernel 侧原理

• Tiling 计算（select_v2_tiling.cpp）：Tiling 阶段通过PlatformAscendC获取 UB 大小与核数，UB 侧固定预留 8KB（RESERVED_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.h）：Kernel 采用 AscendC 编程模型，使用双缓冲TQue（BUFFER_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::vector<int64_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::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初始化, 参考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::vector<int64_t> selfShape = {4, 2}; std::vector<int64_t> otherShape = {4, 2}; std::vector<int64_t> conditionShape = {4, 2}; std::vector<int64_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::vector<float> selfHostData = {0, 0, 0, 0, 0, 0, 0, 7}; std::vector<float> otherHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<int8_t> conditionHostData = {false, false, false, false, true, true, true, true}; std::vector<float> 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::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，需要根据具体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