CANN ops-math 标量不等比较算子:aclnnNeScalar & aclnnInplaceNeScalar 两段式接口实战指南
【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math
本文围绕 CANN ops-math 仓库中experimental/math/not_equal模块的aclnnNeScalar与aclnnInplaceNeScalar两个 aclnn 接口展开,系统讲解"张量与标量逐元素比较是否不相等(!=)"这一基础数学算子的功能定义、两段式调用流程、参数与错误码约定,并结合仓库中的算子定义、tiling 调度、kernel 实现与单元测试源码,帮助开发者在 NPU 上正确、高效地完成标量比较计算。
算子功能与适用场景
aclnnNeScalar与aclnnInplaceNeScalar是 NotEqual(不等比较)算子的标量(Scalar)版本:将输入张量self中的每一个元素与标量other逐一比较,判断是否不相等,输出 BOOL 语义的结果。其计算公式为:
$$ out_i=(self_i \neq other)?[1]:[0] $$
inplace 版本则直接把结果写回输入张量所在内存:
$$ selfRef_i=(selfRef_i \neq other)?[1]:[0] $$
该算子在 CANN ops-math 仓库中的完整算子族还包括 NotEqual 双输入版本(aclnnNeTensor/aclnnInplaceNeTensor) 与 aclnnLogicalXor 等调用入口,具体调用方式可参考 not_equal 模块 README。标量比较是掩码生成、条件统计、数据清洗等场景中的高频基础操作。
产品支持情况
根据接口文档,aclnnNeScalar/aclnnInplaceNeScalar的产品适配情况如下:
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | × |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | × |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
需要说明的是,算子内核侧(AICore)的配置在源码中注册为ascend910b(见 not_equal_def.cpp),即当前算子内核面向昇腾 910B 系列芯片架构编译,接口文档中的产品支持矩阵是各软硬件栈组合下的最终适配结论,实际部署时请以当前 CANN 版本的兼容性矩阵为准。
两段式接口与函数原型
aclnnNeScalar与aclnnInplaceNeScalar实现相同的功能语义,区别仅在输出方式:
- aclnnNeScalar:需要新建一个输出张量对象
out存储计算结果,输入self保持不变; - aclnnInplaceNeScalar:无需新建输出张量对象,直接在输入张量
selfRef的内存中覆写计算结果。
两个算子均遵循 CANN aclnn 的两段式接口约定(基础概念参见 docs/zh/context/basic_concept.md):必须先调用GetWorkspaceSize版本接口获取入参校验结果、workspace 大小与 op 执行器,再调用执行版本接口真正下发计算任务。相关函数原型如下:
aclnnStatus aclnnNeScalarGetWorkspaceSize(const aclTensor *self, const aclScalar *other, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnNeScalar(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream) aclnnStatus aclnnInplaceNeScalarGetWorkspaceSize(aclTensor *selfRef, const aclScalar *other, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnInplaceNeScalar(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnNeScalarGetWorkspaceSize 参数说明
| 参数 | 输入/输出 | 说明 |
|---|---|---|
| self | 计算输入 | Device 侧aclTensor,对应公式中的self。shape 维度不高于 8 维,支持非连续 Tensor,数据格式支持 ND |
| other | 计算输入 | Host 侧aclScalar,对应公式中的other,为参与比较的标量值 |
| out | 计算输出 | Device 侧aclTensor,对应公式中的out,数据类型为 BOOL 可转换的数据类型,shape 与self的 shape 一致,支持非连续 Tensor,数据格式支持 ND |
| workspaceSize | 出参 | 返回需要在 Device 侧申请的 workspace 大小 |
| executor | 出参 | 返回 op 执行器,包含算子计算流程 |
数据类型支持情况(不同产品线存在差异,以下为文档列出的典型集合):
self(昇腾 910_95 系列):DOUBLE、FLOAT16、FLOAT、BFLOAT16、INT64、INT32、INT8、UINT8、BOOL、INT16、COMPLEX64、COMPLEX128、UINT64,且与other满足 TensorScalar 互推导关系;self(Atlas A2/A3 系列):DOUBLE、FLOAT16、FLOAT、BFLOAT16、INT64、INT32、INT8、UINT8、BOOL、INT16、COMPLEX64、COMPLEX128;other的数据类型集合与对应产品线下的self保持一致,并与self满足互推导关系(例如 FLOAT 的self配合 FLOAT 的other);out支持 BOOL 可转换的数据类型(昇腾 910_95 系列还额外支持 UINT32、UINT16),shape 与self严格一致。
从仓库算子定义 not_equal_def.cpp 可以看到,底层 NotEqual 算子注册的输入x1/x2支持FLOAT16、FLOAT、INT32、INT8、UINT8、BOOL、BF16,输出y固定为BOOL,格式统一为ND,aclnn 接口层的数据类型集合是这一算子定义在更多数据类型(含 DOUBLE、INT64、COMPLEX 等)上的扩展映射。
返回值与错误码
两段式接口均返回aclnnStatus状态码(参见 aclnn 返回码 相关说明)。第一段接口完成入参校验,出现以下场景时报错:
返回161001(ACLNN_ERR_PARAM_NULLPTR):1. 传入的self、other、out是空指针。 返回161002(ACLNN_ERR_PARAM_INVALID):1. self,other或out的数据类型不在支持的范围之内。 2. self和other数据类型不满足数据类型推导规则。 3. self和out的shape不同。 4. self和out的维度大于8。shape 校验规则与 not_equal_infershape.cpp 中的 InferShape 实现一致:输出 shape 直接拷贝输入x1的 shape,因此接口层强制out与self同 shape。
aclnnNeScalar 执行接口参数说明
第二段接口在完成 workspace 内存申请后调用,参数如下:
| 参数 | 输入/输出 | 说明 |
|---|---|---|
| workspace | 入参 | 在 Device 侧申请的 workspace 内存地址 |
| workspaceSize | 入参 | Device 侧 workspace 大小,由第一段接口aclnnNeScalarGetWorkspaceSize获取 |
| executor | 入参 | op 执行器,包含算子计算流程 |
| stream | 入参 | 指定执行任务的 Stream |
aclnnInplaceNeScalarGetWorkspaceSize 参数说明
inplace 版本将"输出"合并到输入selfRef上,因此入参更精简:
| 参数 | 输入/输出 | 说明 |
|---|---|---|
| selfRef | 计算输入/输出 | Device 侧aclTensor,对应公式中的selfRef。shape 维度不高于 8 维,支持非连续 Tensor,数据格式支持 ND;数据类型集合与other满足 TensorScalar 互推导关系(昇腾 910_95 系列额外支持 UINT32/UINT16 等) |
| other | 计算输入 | Host 侧aclScalar,数据类型集合与对应产品线下selfRef一致 |
| workspaceSize | 出参 | 返回需要在 Device 侧申请的 workspace 大小 |
| executor | 出参 | 返回 op 执行器,包含算子计算流程 |
错误码约定(第一段接口完成入参校验):
返回161001(ACLNN_ERR_PARAM_NULLPTR):1. 传入的selfRef、other是空指针时。 返回161002(ACLNN_ERR_PARAM_INVALID):1. selfRef和other的数据类型不在支持的范围之内。 2. selfRef和other的数据类型不满足数据类型推导规则。 3. selfRef的维度大于8。注意 inplace 版本由于输出即输入,无需校验out的 shape 与维度。执行接口aclnnInplaceNeScalar的参数与aclnnNeScalar一致(workspace、workspaceSize、executor、stream),此处不再赘述。
约束说明
接口文档明确约束为无。但算子底层(见 not_equal 模块 README 的约束说明)仍存在两条使用前提:底层 NotEqual 算子不支持广播,且算子内核定义层不支持 INT64 输入——这意味着使用aclnnNeScalar系列接口时,self与out必须保持完全一致的 shape(接口层已强制校验),而 INT64 等扩展类型由接口层的类型推导与转换逻辑承接,请勿将内核层约束误解为接口层能力。
调用示例
以下完整示例演示如何依次调用aclnnNeScalar与aclnnInplaceNeScalar(本示例与仓库 examples/test_aclnn_ne_scalar.cpp 内容一致,可直接参考;编译与运行样例的完整流程参见 docs/zh/context/basic_concept.md)。
#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnn_ne_scalar.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; } 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_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> otherShape = {4, 2}; std::vector<int64_t> outShape = {4, 2}; void *selfDeviceAddr = nullptr; void *otherDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor *self = nullptr; aclScalar *other = nullptr; aclTensor *out = nullptr; std::vector<float> selfHostData = {0, 1, 2, 3, 4, 5, 6, 7}; std::vector<float> outHostData = {0, 0, 0, 0, 0, 0, 0, 0}; float otherValue = 1.0f; // 创建self aclTensor ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建other aclScalar other = aclCreateScalar(&otherValue, aclDataType::ACL_FLOAT); 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); // aclnnNeScalar调用示例 // 3. 调用CANN算子库API,需要修改为具体的API名称 LOG_PRINT("test aclnnNeScalar\n"); uint64_t workspaceSize = 0; aclOpExecutor *executor; // 调用aclnnNeScalar第一段接口 ret = aclnnNeScalarGetWorkspaceSize(self, other, out, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNeScalarGetWorkspaceSize 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); } // 调用aclnnNeScalar第二段接口 ret = aclnnNeScalar(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNeScalar 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(selfShape); std::vector<float> resultData(size, 0); ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), 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]); } // aclnnInplaceNeScalar调用示例 // 3. 调用CANN算子库API,需要修改为具体的API名称 LOG_PRINT("\ntest aclnnInplaceNeScalar\n"); // 调用aclnnInplaceNeScalar第一段接口 ret = aclnnInplaceNeScalarGetWorkspaceSize(self, other, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceNeScalarGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 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); } // 调用aclnnInplaceNeScalar第二段接口 ret = aclnnInplaceNeScalar(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceNeScalar 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的接口定义修改 ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr, size * sizeof(resultData[0]), 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,需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyScalar(other); aclDestroyTensor(out); // 7. 释放device资源,需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中的关键点:
selfHostData = {0, 1, 2, 3, 4, 5, 6, 7}、otherValue = 1.0f时,非 inplace 版本的期望结果为{1, 0, 1, 1, 1, 1, 1, 1}(仅元素 1 与标量 1 相等,其余均不相等输出 1);inplace 版本执行后self所在内存被覆写为同样的结果;- 两次调用共用同一个
workspaceSize/executor变量,但每个接口都独立完成了"第一段取 workspace → 申请内存 → 第二段执行"的完整两段式流程; - workspace 为 0 时无需申请内存,直接传入
nullptr即可。
源码级实现剖析
算子定义与 shape 推导
not_equal_def.cpp 通过OpDef注册NotEqual算子:输入x1、x2与输出y均为 ND 格式,输出数据类型固定为 BOOL;not_equal_infershape.cpp 将输出 shape 直接赋值为输入x1的 shape,这正是接口层"out与self同 shape"校验的底层依据。
Tiling 调度策略
not_equal_tiling.cpp 中的TilingFunc展示了该算子的核数调度策略:以输入元素总数size为基准计算每个核分块,并结合数据类型做分级降核——FLOAT/BF16 类型下,元素数小于 8192(1<<13)时仅用 1 个核,小于 2M(1<<21)时核数减半;其他类型按数据字节数data_size在 64KB(1<<16)/16MB(1<<24)两个阈值处做同样的降核处理。同时通过GetLibApiWorkSpaceSize()申请算子库 API 所需的 workspace,这正是GetWorkspaceSize接口返回值中 workspace 大小的来源之一。
Kernel 实现原理
not_equal.h 中针对不同数据类型实现了多套NotEqual模板特化:
- half 类型(L77-L85):先
Compare(CMPMODE::EQ)得到相等掩码,再通过Duplicate构造 0/1 向量并用Select反转得到"不相等"结果,最后Cast到 BOOL; - float 类型(L88-L99):以 256 字节为步长循环调用
Compare,借助Select(VSEL_CMPMASK_SPR 模式)完成比较结果的 0/1 化; - int 类型(L102-L112)与 half 类似,但复用
ReinterpretCast后的 half 缓冲区完成 Select 与 Cast; - int8/uint8 类型(L115-L127):由于 8bit 类型不便于直接做向量比较,采用"先 Cast 到半精度 → 相减 → Abs → 与极小值 Mins → 放大回 BOOL"的数学技巧间接判断不等性;
- bfloat16 类型(L130-L145):先做无损 Cast,再走与 float 相同的
Compare循环路径。
主入口not_equal(L148-L188)通过TPipe/TQue双缓冲流水线,以MAX_TILE_SIZE(30KB)为粒度完成 GlobalMemory 到 LocalMemory 的 DataCopy、计算与结果回写三段流水;类型分发逻辑见 L191-L207,BOOL 输入会被映射为 uint8 参与计算,float/bfloat16 统一按 float 路径处理。
单元测试验证
tests/ut/op_kernel/test_not_equal.cpp 基于 gtest 与 tikicpulib 进行算子级 UT:以 94000 个元素(DTYPE_X1类型由编译宏决定)为输入,构造NotEqualTilingData(size = 94000)、block_dim = 20的 tiling 参数,通过ICPU_SET_TILING_KEY(0)与SetKernelMode(KernelMode::AIV_MODE)在 CPU 上模拟 AIV 内核执行,并与 gen_data.py 生成的期望数据比对,验证 kernel 计算结果的正确性。这为 aclnn 接口层的两段式调用提供了底层计算正确性保障。
总结
aclnnNeScalar与aclnnInplaceNeScalar为开发者提供了"张量 vs 标量"不等比较的标准 aclnn 入口:前者输出独立张量、语义清晰,后者就地覆写、省内存省拷贝。理解两段式接口的 workspace 获取与执行流程、掌握各产品线下数据类型与 shape 校验规则,是正确使用该类算子的关键。进一步地,通过阅读本仓库中 NotEqual 算子的 tiling 与 kernel 实现,还能为自定义比较类算子的开发提供可直接借鉴的分核调度与向量化实现范式。
【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考