ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

CANN ops-math 算子接口实战:aclnnLinalgCross 向量叉乘(Cross)算子详解

CANN ops-math 算子接口实战:aclnnLinalgCross 向量叉乘(Cross)算子详解 CANN ops-math 算子接口实战aclnnLinalgCross 向量叉乘Cross算子详解【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本篇文章以 CANN ops-math 开源仓库中 Cross 算子的 aclnnLinalgCross 接口文档 为核心系统讲解该算子支持的硬件平台、数学定义、两段式调用接口aclnnLinalgCrossGetWorkspaceSize与aclnnLinalgCross、全部入参/出参约束、错误码语义并结合仓库内 op_api、op_host、op_kernel 的源码实现与 UT 测试用例剖析接口背后的校验逻辑、广播处理、BF16 升精度策略与昇腾 AI Core 上的算子内核计算流程。读完本文你将能够独立编写一个可编译、可运行的 aclnnLinalgCross 调用程序并理解从 API 入参校验到 NPU 内核执行的完整链路。一、算子功能与数学原理aclnnLinalgCross是 CANN 提供的一个数学类基础算子接口功能是对输入 Tensorself与other完成linear_cross线性叉乘运算即逐组地对dim指定维度上长度为 3 的向量求三维向量叉积对应 PyTorch 语义中的torch.linalg.cross。设self (x1, y1, z1)、other (x2, y2, z2)叉乘结果按三阶行列式展开为$$ 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 $$即输出三个分量为out[0] x1 * y2 - x2 * y1out[1] x2 * y0 - x0 * y2对应y1z2 - y2z1的通用分量式out[2] x0 * y1 - x1 * y0这一计算公式在仓库内核实现中可得到直接印证见 op_kernel/cross.h 的Compute()逻辑三路输出分别用AscendC::Mul与AscendC::Sub组合实现。二、产品支持情况产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√从算子定义源码看op_host/cross_def.cpp 中通过this-AICore().AddConfig(ascend910b)声明该算子适配的 AI Core 架构Atlas A2 系列对应昇腾 910B 架构输入x1、x2与输出y均要求 ND 格式。三、两段式接口调用模型每个 CANN 算子 API 都遵循两段式接口详见 两段式接口说明先调用aclnnLinalgCrossGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器executor再调用aclnnLinalgCross执行计算。两个接口的函数原型如下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)头文件声明位于 op_api/aclnn_linalg_cross.h接口域为aclnn_math。四、aclnnLinalgCrossGetWorkspaceSize 参数详解4.1 参数说明参数名输入/输出描述使用说明数据类型数据格式维度非连续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 执行器包含算子计算流程。-----4.2 返回值第一段接口返回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] 范围内。4.3 源码级校验链路上述校验在 op_api/aclnn_linalg_cross.cpp 的CheckParams中按序完成是理解各错误码触发条件的第一手依据空指针校验CheckNotNull对self、other、out逐一执行OP_CHECK_NULL任一为空即返回ACLNN_ERR_PARAM_NULLPTR161001。数据类型校验CheckDtypeValid源码中维护了两张支持列表DTYPE_SUPPORT_LISTFLOAT、FLOAT16、DOUBLE、INT8、INT16、INT32、INT64、COMPLEX64、COMPLEX128、UINT8DTYPE_SUPPORT_LIST_910B在上一列表基础上额外包含 BF16当当前 NPU 架构为DAV_2201或DAV_3510时使用含 BF16 的列表否则使用基础列表。同时通过OP_CHECK_DTYPE_NOT_MATCH强制self、other、out三者数据类型一致。维度校验CheckShape三个 Tensor 的维度都必须小于MAX_SUPPORT_DIMS_NUMS即 0-8 维。广播与输出 shape 校验CheckBroadcastShape先OP_CHECK_BROADCAST确认self/other满足广播关系再调用BroadcastInferShape推导广播后的 shape要求其与out的 shape 完全一致。dim 范围校验CheckDimdim必须落在[-dimSize, dimSize-1]dimSize 为 self 的维度数默认值 -1。关于 broadcast 的具体规则左对齐补 1、单维为 1 时拉伸、无可广播维时报错等参见 broadcast 关系。五、aclnnLinalgCross 参数详解5.1 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnLinalgCrossGetWorkspaceSize 获取。executor输入op 执行器包含算子计算流程。stream输入指定执行任务的 Stream。5.2 返回值返回aclnnStatus状态码具体参见 aclnn 返回码说明。从源码看第二段接口内部通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)调用框架能力完成实际计算见 op_api/aclnn_linalg_cross.cpp 第 233-238 行用户无需关心 executor 内部细节。六、约束说明确定性计算aclnnLinalgCross默认采用确定性实现即相同输入多次执行结果一致不会引入随机性偏差。维度约束指定进行叉乘的维度dim轴在广播后的 size 必须为 3。广播约束输入self、other的 shape 必须满足 broadcast 关系运算前较小的张量会被广播拉伸到与较大张量一致。类型约束self、other、out的数据类型必须保持一致Atlas A2 系列上支持 INT8、INT16、INT32、INT64、UINT8、FLOAT16、BFLOAT16、FLOAT、FLOAT64、COMPLEX64、COMPLEX128。仓库 README.md 亦明确指定维度的 size 必须为 3输入 shape 需满足 broadcast 关系输入数据类型需保持一致。七、接口内部实现原理源码纵深7.1 从 API 到算子内核的调用链aclnnLinalgCrossGetWorkspaceSize的实现op_api/aclnn_linalg_cross.cpp 的ExecLinalgCrossGetWorkspaceSize完成了以下关键步骤空 Tensor 短路处理if (out-IsEmpty())时直接返回 workspace 大小与 executor不进入内核计算对照 UT 中多个case_empty_tensor用例如{1, 0, 3}形状的输入可正常通过校验。连续化对非连续的self、other分别调用l0op::Contiguous转换为连续 Tensor因此接口支持非连续 Tensor 输入。广播展开当两者 shape 不同时通过BroadcastTensor调用l0op::BroadcastTo将两个输入统一广播到out的 shape若dim 0且self维度少于other还会把dim加上扩维数量以保持索引正确若dim为负数则加上广播后 Tensor 的维度数转为正索引。叉乘维度检查校验广播后dimInner维上的 shape 为 3否则返回ACLNN_ERR_PARAM_INVALID。BF16 升精度在非DAV_2201/DAV_3510架构上BF16 输入会被l0op::Cast为 FLOAT 计算计算结果再 Cast 回 BF16保证精度。内核调用与结果回写调用l0op::Cross(selfBroadCast, otherBroadCast, dimInner, executor)生成计算图最后用l0op::ViewCopy将结果拷贝到可能非连续的out上并汇总得到workspaceSize。7.2 shape 推导InferShapeop_host/cross_infershape.cpp 中的InferShapeCross将输出 shape 直接填充为输入x即 self的 shape。由于广播已在 API 层完成进入算子框架时输入已具备一致 shape。7.3 Tiling 切分策略op_host/cross_tiling.cpp 实现了编译期切分逻辑通过GetPlatformInfo获取 UB 大小与 AIV 核数GetCoreNumAiv依据输入 shape 与dim属性计算totalIdx总元素数、intervalNum与loopTimes叉乘轴 size 固定为 3kCrossAxisSize 3intervalNum totalIdx / dimProductloopTimes totalIdx / intervalNum / 3按数据类型计算单次搬入 UB 的tileDataNum区分 FP32/FP16/INT8 及 group 模式对应UB_DATA_NUM_FP32 19、UB_DATA_NUM_FP16 30、UB_DATA_NUM_INT8 18、UB_DATA_NUM_GROUP_MODE 6等缓冲数常量根据totalTileCount与核数计算blockDim实现多核并行切分按 dtype 设置 tiling keyFLOAT→0、INT32→1、INT8→2、FLOAT16→3、UINT8→4、INT16→5workspace 固定申请 16MBWS_SYS_SIZE。7.4 AI Core 内核计算AscendCop_kernel/cross.h 中的CrossT模板类实现了昇腾 AI Core 上的向量流水使用TPipe与TQue队列组织CopyIn → Compute → CopyOut三级流水Process()中交错执行 Compute 与下一次 CopyIn隐藏访存延迟非 group 模式下输入按x0/x1/x2、y0/y1/y2拆分为六个队列输出z0/z1/z2三个队列分别对应叉乘结果的三个分量对 FLOAT/INT32/INT16 使用向量指令AscendC::MulAscendC::Sub计算对 FLOAT16half先Cast到 FP32 计算再Cast回CAST_ROUND避免半精度溢出对 INT8/UINT8 提升到 INT32 计算后写回op_kernel/cross.cpp 中cross内核入口按 tiling key 模板分发到Crossfloat、Crosshalf、Crossint8_t等具体实例。八、完整调用示例以下示例代码完整演示了两段式接口的调用流程可直接参照 编译与运行样例 编译执行。代码逻辑为初始化 ACL 资源 → 构造 self3x3值为 0~8与 other3x3值为三组 1/2/3→ 调用aclnnLinalgCrossGetWorkspaceSize获取 workspace → 申请 Device 内存 → 调用aclnnLinalgCross执行 → 拷贝结果回 Host 并打印 → 释放资源。#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::vectordouble selfHostData {0, 1, 2, 3, 4, 5, 6, 7, 8}; std::vectordouble otherHostData {1, 1, 1, 2, 2, 2, 3, 3, 3}; std::vectordouble outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0}; // 创建 self aclTensor auto ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_DOUBLE, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建 other aclTensor ret CreateAclTensor(otherHostData, otherShape, otherDeviceAddr, aclDataType::ACL_DOUBLE, other); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建 out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_DOUBLE, 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::vectordouble 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); // 4. 释放 aclTensor aclDestroyTensor(self); aclDestroyTensor(other); aclDestroyTensor(out); // 5. 释放 device 资源 aclrtFree(selfDeviceAddr); aclrtFree(otherDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中self取{3, 3}dim1 上每组向量为(0,1,2)、(3,4,5)、(6,7,8)other取{3, 3}每组向量(1,1,1)、(2,2,2)、(3,3,3)逐行叉乘后可对照验证输出第一行(1*2-2*1, 2*1-0*2, 0*1-1*1) (0, 2, -1)。仓库另提供一份可独立编译的完整样例 examples/test_aclnn_linalg_cross.cpp。九、测试与验证仓库为 aclnnLinalgCross 提供了两层自动化测试可作为接口行为与边界条件的权威参考op_api 层 UTtests/ut/op_api/test_aclnn_linalg_cross.cpp 覆盖了不支持类型校验如ACL_BOOL输入返回ACLNN_ERR_PARAM_INVALID、空 Tensor{1, 0, 3}、{1, 1, 0, 3}等在有无广播时的通过/失败分支以及合法输入的TestGetWorkspaceSize调用op_host / op_kernel 层测试tests/ut/op_host/test_cross_tiling.cpp 验证 tiling 切分参数intervalNum、loopTimes、tileDataNum、blockDimtests/ut/op_kernel/test_cross.cpp 配合 gen_data.py 与 compare_data.py 完成内核数值比对。十、小结aclnnLinalgCross是 CANN ops-math 中实现三维向量叉乘的标准接口具有类型覆盖广含整数、浮点、复数与 BF16、支持广播与非连续 Tensor、默认确定性计算的特点。理解其两段式接口模型、参数与错误码语义再结合 op_api、op_host、op_kernel 三层的源码实现即可在 Atlas A2 系列产品上正确、高效地完成叉乘算子的接入与调试。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表