CANN ops-math 中 Crosslinalg.cross向量叉积算子的原理与 aclnn 调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN ops-math 仓库中experimental/math/cross目录下的 Cross 算子展开该算子实现数学上的向量叉积linear_cross运算支持与 NumPy/PyTorchlinalg.cross对齐的语义沿指定维度对两个张量逐元素求叉积并支持广播。读完本文你将掌握 Cross 算子的功能定义、参数与约束、基于两段式 aclnn 接口aclnnLinalgCross的完整调用方法以及从算子定义、shape 推导、tiling 到 AscendC kernel 的源码级实现原理。功能说明线性叉积linear_cross算子Cross 算子的功能是对输入 Tensorself与other沿指定维度执行向量叉积linear_cross运算。向量叉积也称向量积、外积其结果是一个与原向量所在平面垂直的新向量是三维几何与线性代数中的基础运算。其计算公式采用三阶行列式展开单位向量 i、j、k 对应叉积结果的三轴分量$$ out self\times other \begin{vmatrix}ijk\x_1y_1z_1\x_2y_2z_2\end{vmatrix} (y_1z_2-y_2z_1)i-(x_1z_2-x_2z_1)j(x_1y_2-x_2y_1)k $$其中self (x1, y1, z1)、other (x2, y2, z2)分别为进行叉积运算的两个三维向量输出out的三个分量分别对应 i、j、k 的系数。产品支持情况Cross 算子目前支持的硬件平台如下产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√支持的数据类型在 Atlas A2 训练/推理系列产品上Cross 算子支持的数据类型包括INT8、INT16、INT32、INT64、UINT8、FLOAT16、BFLOAT16、FLOAT、FLOAT64、COMPLEX64、COMPLEX128。参数说明Cross 算子的核心参数如下表所示参数名输入/输出/属性描述数据类型数据格式self输入计算公式中的输入 self。FLOAT16、FLOAT、INT32、INT8、UINT8、INT16、DOUBLE、INT64、UINT16、UINT32、UINT64、COMPLEX64、COMPLEX128NDother输入计算公式中的输入 other。FLOAT16、FLOAT、INT32、INT8、UINT8、INT16、DOUBLE、INT64、UINT16、UINT32、UINT64、COMPLEX64、COMPLEX128NDout输出计算公式中的输出 out。FLOAT16、FLOAT、INT32、INT8、UINT8、INT16、DOUBLE、INT64、UINT16、UINT32、UINT64、COMPLEX64、COMPLEX128ND说明上表中的数据类型为算子整体支持范围针对 Atlas A2 训练/推理系列产品实际支持范围以 aclnnLinalgCross 接口文档 中列出的INT8、INT16、INT32、INT64、UINT8、FLOAT16、BFLOAT16、FLOAT、FLOAT64、COMPLEX64、COMPLEX128为准。此外不同调用路径AI Core 内核路径与 AICPU 路径的支持范围存在差异详见后文源码剖析。约束说明使用 Cross 算子时需满足以下约束叉积维度大小必须为 3指定进行 cross 运算的维度dim参数指定的轴的 size 必须为 3。这是因为叉积只定义在三维向量上。输入需满足 broadcast 关系两个输入 tensor 的 shape 需要兼容能够进行广播运算具体广播规则参见 broadcast_relationship.md。数据类型需保持一致self、other、out三个 tensor 的数据类型必须完全一致。aclnn 接口调用说明两段式接口Cross 算子对外提供 aclnn 层接口aclnnLinalgCross采用 CANN 算子标准的两段式接口调用模式参见 two_phase_api.md先调用aclnnLinalgCrossGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnLinalgCross接口执行实际计算。调用示例位于 examples/test_aclnn_linalg_cross.cpp接口详细说明见 aclnnLinalgCross.md。函数原型aclnnStatus aclnnLinalgCrossGetWorkspaceSize( const aclTensor* self, const aclTensor* other, int64_t dim, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnLinalgCross( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)aclnnLinalgCrossGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度非连续 TensorselfaclTensor*输入公式中的 self。数据类型与 other 和 out 一致需与 other 满足 broadcast 关系且 shape 在 dim 指定的轴广播后的值为 3。INT8、INT16、INT32、INT64、UINT8、FLOAT16、BFLOAT16、FLOAT、FLOAT64、COMPLEX64、COMPLEX128ND0-8√otheraclTensor*输入公式中的 other。数据类型与 self 和 out 一致需与 self 满足 broadcast 关系。INT8、INT16、INT32、INT64、UINT8、FLOAT16、BFLOAT16、FLOAT、FLOAT64、COMPLEX64、COMPLEX128ND0-8√dimint64_t输入指定 self 进行 linear_cross 的轴。若不指定则默认为 -1取值范围为[-self维度数量, self维度数量-1]。----outaclTensor*输出公式中的 out。数据类型与 self 和 other 一致shape 需要与 self 和 other broadcast 后的 shape 一致。INT8、INT16、INT32、INT64、UINT8、FLOAT16、BFLOAT16、FLOAT、FLOAT64、COMPLEX64、COMPLEX128ND-√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----返回值aclnnStatus状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现如下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针。ACLNN_ERR_PARAM_INVALID161002self、other、out 的数据类型不一致或数据格式不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002self 和 other 的维度大于 8。ACLNN_ERR_PARAM_INVALID161002self 和 other 不符合 broadcast 关系。ACLNN_ERR_PARAM_INVALID161002self 和 other broadcast 后的 shape 与 out 不一致。ACLNN_ERR_PARAM_INVALID161002self 在对应 dim 维度上的 shape 不为 3。ACLNN_ERR_PARAM_INVALID161002dim 的值不在[-self的维度数量, self的维度数量-1]范围内。aclnnLinalgCross 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnLinalgCrossGetWorkspaceSize获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值aclnnStatus状态码具体参见 aclnn 返回码。确定性说明aclnnLinalgCross默认为确定性实现同一输入在多次运行中计算结果可复现。完整调用示例下面给出一个可编译运行的完整示例基于 examples/test_aclnn_linalg_cross.cpp以FLOAT类型、shape 为{3, 3}的输入沿dim 1计算叉积。具体编译与运行流程请参考 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_linalg_cross.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; } 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; } aclError InitAcl(int32_t deviceId, aclrtStream* stream) { auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); return ACL_SUCCESS; } aclError CreateInputs( std::vectorint64_t selfShape, std::vectorint64_t otherShape, std::vectorint64_t outShape, void** selfDeviceAddr, void** otherDeviceAddr, void** outDeviceAddr, aclTensor** self, aclTensor** other, aclTensor** out) { std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7, 8}; std::vectorfloat otherHostData {1, 1, 1, 2, 2, 2, 3, 3, 3}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0}; // 创建 self aclTensor auto 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); // 创建 out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); return ACL_SUCCESS; } aclError ExecOpApi( aclTensor* self, aclTensor* other, aclTensor* out, int64_t dim, void** workspaceAddrOut, uint64_t workspaceSize, void* outDeviceAddr, std::vectorint64_t outShape, aclrtStream stream) { aclOpExecutor* executor; // 调用 aclnnLinalgCross 第一段接口 auto ret aclnnLinalgCrossGetWorkspaceSize(self, other, dim, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnLinalgCrossGetWorkspaceSize 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); } *workspaceAddrOut workspaceAddr; // 调用 aclnnLinalgCross 第二段接口 ret aclnnLinalgCross(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnLinalgCross failed. ERROR: %d\n, ret); return ret); // 同步 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 从 device 拷贝结果到 host auto size GetShapeSize(outShape); std::vectorfloat 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: %lf\n, i, resultData[i]); } return ACL_SUCCESS; } int main() { // 1. device/stream 初始化 int32_t deviceId 0; aclrtStream stream; CHECK_RET(InitAcl(deviceId, stream) ACL_SUCCESS, return -1); // 2. 构造输入与输出 std::vectorint64_t selfShape {3, 3}; std::vectorint64_t otherShape {3, 3}; std::vectorint64_t outShape {3, 3}; void* selfDeviceAddr nullptr; void* otherDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* other nullptr; aclTensor* out nullptr; aclError ret CreateInputs( selfShape, otherShape, outShape, selfDeviceAddr, otherDeviceAddr, outDeviceAddr, self, other, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用 CANN 算子 API int64_t dim 1; uint64_t workspaceSize 0; void* workspaceAddr nullptr; ret ExecOpApi(self, other, out, dim, workspaceAddr, workspaceSize, outDeviceAddr, outShape, stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 6. 释放 aclTensor aclDestroyTensor(self); aclDestroyTensor(other); aclDestroyTensor(out); // 7. 释放 device 资源 aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例程序的执行流程可以概括为以下步骤资源初始化aclInit→aclrtSetDevice→aclrtCreateStream建立与设备的会话和任务流构造 aclTensor通过aclrtMalloc申请 Device 侧内存、aclrtMemcpy将 Host 数据拷贝到 Device再由aclCreateTensor携带 shape/strides/dtype/format 创建三个张量描述符self、other、out两段式调用先调用aclnnLinalgCrossGetWorkspaceSize获取 workspace 大小并申请内存再调用aclnnLinalgCross提交计算任务结果回拷与校验aclrtSynchronizeStream等待任务完成通过ACL_MEMCPY_DEVICE_TO_HOST将结果拷贝回 Host 并打印资源释放依次释放 aclTensor、Device 内存、workspace、Stream最后aclrtResetDevice与aclFinalize。以示例输入为例self [0 1 2; 3 4 5; 6 7 8]other [1 1 1; 2 2 2; 3 3 3]沿dim 1对每一行向量求叉积可验证第一行(0,1,2) × (1,1,1) (1*1-2*1, 2*1-0*1, 0*1-1*1) (-1, 2, -1)与公式一一对应。源码实现剖析Cross 算子完整的实现链路覆盖了 L0 API 层、Host 侧算子定义 / InferShape / Tiling与 AI Core Kernel 三层下面逐一展开。L0 API 层的多路径调度op_api/cross.cpp 实现了 L0 层Cross接口其核心逻辑是根据芯片型号与输入数据类型动态选择执行路径AI Core 路径Cross当数据类型属于AICORE_DTYPE_SUPPORT_LISTDT_FLOAT、DT_FLOAT16、DT_INT8、DT_INT16、DT_INT32、DT_UINT8时通过ADD_TO_LAUNCHER_LIST_AICORE宏将算子加入 AI Core 任务队列属性dim随任务下发AI Core 路径CrossV2在 Atlas A2 系列NpuArch::DAV_2201、DAV_3510上支持DT_FLOAT、DT_FLOAT16、DT_INT16、DT_INT32、DT_BF16类型走 CrossV2 内核AICPU 路径其余数据类型如 DOUBLE、INT64、COMPLEX 等回退到 AICPU 内核执行保证接口全类型可用。同时该文件通过OP_TYPE_REGISTER(Cross)与OP_TYPE_REGISTER(CrossV2)完成算子类型注册输出张量由executor-AllocTensor依据 self 的 shape 与 dtype 自动分配。Host 侧算子定义、shape 推导与 tiling算子定义位于 op_host/cross_def.cpp声明了两个输入x1、x2与一个输出y数据类型支持DT_FLOAT、DT_INT32、DT_INT8、DT_FLOAT16、DT_UINT8、DT_INT16格式均为FORMAT_ND并注册ascend910b的 AI Core 配置。shape 推导位于 op_host/cross_infershape.cpp实现非常简洁由于叉积不改变张量形状仅逐元素重算输出 shape 直接继承输入x1的 shape*yShape *xShape。值得注意的是这里的 InferShape 以x1为准而 API 层的校验已保证 self 与 other 广播后的 shape 与 out 一致。tiling 计算位于 op_host/cross_tiling.cpp负责把任务切分到多个 AI Core 并确定每次搬运的数据量平台信息通过PlatformAscendC获取 UB 大小与可用 AI Core 数量关键切分参数totalIdx输入总元素数、intervalNum叉积维之外的前缀切片数即totalIdx / (dim 及之前的维度乘积)、loopTimes沿叉积维切出的三元组组数即totalIdx / intervalNum / 3分组模式IsGroupMode即intervalNum 1当数据全部落在叉积维时把多个三元组整体作为一个 tile 搬运并针对 FLOAT16 增加了按核数限制的细粒度切分非分组模式按tileDataNum切分 interval并依据不同 dtypeFLOAT16/FLOAT/INT32/INT8 等选择不同的 UB buffer 数量UB_DATA_NUM_FP1630、UB_DATA_NUM_FP3219、UB_DATA_NUM_INT818以适配数据宽度每个 dtype 对应不同的TilingKeyELEMENTWISE_TPL_SCH_MODE_0到MODE_5kernel 据此选择实例化的模板分支workspace 固定申请16MBWS_SYS_SIZE块对齐粒度32BBLOCK_SIZE空输入totalIdx 0时直接输出空 tiling 数据并设置blockDim 1避免无效计算。AI Core KernelAscendC 实现与多核流水Kernel 实现位于 op_kernel/cross.h模板类NsCross::CrossT按T实例化不同 dtype 走不同计算分支常规分支float/int32_t/int16_t利用向量指令AscendC::Mul与AscendC::Sub将叉积展开为三组乘-减流水每组之间通过PipeBarrierPIPE_V保证向量计算单元的依赖顺序计算结果依次写入 z0/z1/z2 三个输出队列halfFLOAT16分支由于低精度乘法精度有限先将数据Cast到 FP32 进行计算再以CAST_ROUND舍入回 FLOAT16 输出中间结果使用 8 个独立的 FP32 临时 buffertmpBuf0~7承载int8_t/uint8_t 分支先将元素提升到int32_t做乘减再截断回原类型写回避免中间溢出分组模式GroupMode按三元组逐个取数、计算、写回逻辑与上述标量计算一致。数据流采用标准的CopyIn → Compute → CopyOut三段流水队列深度QUEUE_DEPTH 1、BUFFER_NUM 2实现双缓冲Process()中通过GetBlockNum()/GetBlockIdx()将总 tile 数均匀含余数分配切分到多个 AI Core 并行执行。搬运时以 32 字节对齐块为主路径DataCopy尾部不足 32 字节的片段使用DataCopyPad补齐确保内存访问对齐。测试验证算子配套的 Kernel 级单测位于 tests/ut/op_kernel/test_cross.cpp测试框架基于gtest与tikicpulibTIKI CPU 模拟器。测试用例通过gen_data.py生成输入数据如 shape(2, 3)的 FLOAT16 用例读取二进制输入后分配 GM 内存、装载CrossTilingData以单核blockDim 1启动 kernel 并与期望输出对比workspace 大小与 tiling 中的16MB保持一致。数据生成脚本位于 tests/ut/op_kernel/cross_data/gen_data.py对比脚本见 compare_data.py。此外tests/ut/op_host/test_cross_tiling.cpp 覆盖了 tiling 逻辑的单元测试。使用建议与注意事项dim 的选择dim 默认 -1最后一维指定其他轴时需保证该轴 size 为 3且取值在[-rank(self), rank(self)-1]之间否则第一段接口会返回ACLNN_ERR_PARAM_INVALID161002广播场景当 self 与 other shape 不同但满足广播关系时结果 shape 为广播后的 shapeout 需按该 shape 预先分配数据类型一致性三个张量 dtype 必须完全一致建议统一使用 FLOAT/FLOAT16 等 AI Core 路径支持的类型以获得最优性能DOUBLE、INT64、COMPLEX 等类型会走 AICPU 回退路径确定性aclnnLinalgCross为确定性实现无需额外配置即可获得可复现结果。贡献说明Cross 算子在 2026 年 4 月由个人开发者 hth810 贡献并适配进开源仓库贡献内容包括算子完整实现与测试用例。当前算子已纳入experimental/math/cross目录读者可结合上述源码路径深入研读与二次开发。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考