CANN ops-math TriangularSolve 算子 aclnn 接口详解从两段式调用到源码实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本篇文章基于 CANN 数学算子库 ops-math 仓库中 math/triangular_solve 目录及其配套接口文档 aclnnTriangularSolve.md 编写。该目录只包含 TriangularSolve 算子对应的 aclnnAscend CANN 算子库 Native 接口两层接口封装用于在 NPU 上求解形如AX b的三角线性方程组。读完本文你将掌握aclnnTriangularSolve两段式接口的函数原型、全部入参/出参语义、错误码与约束条件并能结合仓库源码理解其底层执行链路直接复用完整可运行的调用示例。功能说明与数学背景aclnnTriangularSolve的接口功能是求解一个具有方形上三角或下三角可逆矩阵 A 和多个右侧项 b 的线性方程组其计算公式为$$ AX b $$其中A 是一个上三角方阵当upper为 false 时为下三角方阵主对角线不含 0 元素b 与 A 为二维矩阵或者为二维矩阵的 batch批量当输入为 batch 时返回的输出 X 也是对应的 batch数值前提当 A 的主对角线含有 0或元素非常接近 0且unitriangular为 false 时输出结果可能包含NaN。该算子隶属于 CANN ops-math 仓库的数学类基础计算算子库当前目录 math/triangular_solve 仅提供 aclnn 接口层实现如果希望贡献该算子的 AscendC kernel 实现可参考仓库根目录的 CONTRIBUTING.md 贡献流程。产品支持情况根据 aclnnTriangularSolve.md 的声明该接口在以下产品形态上的支持情况如下产品形态支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品310P不支持Atlas 训练系列产品910支持从源码看该接口的 L0 层通过 AICPUAI CPUkernel 完成计算详见下文“源码实现剖析”因此上层产品能力取决于 AICPU 侧的MatrixTriangularSolvekernel 部署情况。两段式接口设计CANN 的 aclnn 算子接口普遍采用两段式Two-Phase设计TriangularSolve 也不例外。相关背景可参考 docs/zh/context/two_phase_api.md。使用流程必须遵循以下顺序先调用aclnnTriangularSolveGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器executor再调用aclnnTriangularSolve接口传入第一段接口返回的 workspace 与 executor 执行实际计算。之所以拆成两段是为了让调用方能够预先得知并申请 Device 侧的工作内存workspace从而把内存分配与算子执行解耦便于上层框架做内存复用与池化。函数原型两个接口的原型声明见 op_host/op_api/aclnn_triangular_solve.haclnnStatus aclnnTriangularSolveGetWorkspaceSize( const aclTensor* self, const aclTensor* A, bool upper, bool transpose, bool unitriangular, aclTensor* xOut, aclTensor* mOut, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnTriangularSolve( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)头文件中同时给出了关键约束的 C 接口注释self对应公式中的 bA对应系数矩阵二者数据类型支持 FLOAT、DOUBLE、COMPLEX64、COMPLEX128维度至少为 2 且不大于 8支持非连续 Tensor数据格式为 ND且shape(*,m,n)与A(*,m,m)除最后两维外需满足 broadcast 关系。第一段接口aclnnTriangularSolveGetWorkspaceSize 参数详解第一段接口的完整参数语义如下表所示信息继承自接口文档并结合头文件注释核对参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入公式中的 b方程右端项数据类型与 A 一致self[-2]A[-2]除最后两个维度之外A 和 self 的其余维度满足 broadcast 关系FLOAT、DOUBLE、COMPLEX64、COMPLEX128ND2-8√AaclTensor*输入公式中的 A系数矩阵数据类型与 self 一致最后两个轴相等方阵除最后两个维度之外A 和 self 的其余维度满足 broadcast 关系FLOAT、DOUBLE、COMPLEX64、COMPLEX128ND2-8√upperbool输入控制公式中的 A 按上三角或下三角参与计算的计算属性默认为 trueA 为上三角方阵当 upper 为 false 时A 为下三角方阵BOOL---transposebool输入控制公式中使用 A 还是 Aᵀ 参与计算的计算属性默认为 false当 transpose 为 true 时计算 AᵀXbBOOL---unitriangularbool输入控制公式中的 A 是否按单位三角矩阵处理的计算属性默认为 false当为 true 时A 的主对角线元素视为 1而不是从 A 引用当为 true 时输入 self 和 A、输出 xOut 和 mOut 的数据类型只支持 FLOATBOOL---xOutaclTensor*输出公式中的 X方程求解结果数据类型与 self 一致shape 需要与 broadcast 后的 A、b 满足 AXb 约束broadcast 之后的维度最后一根轴 dimself[-1]FLOAT、DOUBLE、COMPLEX64、COMPLEX128ND-√mOutaclTensor*输出broadcast 后 A 的上三角下三角拷贝数据类型与 self 一致broadcast 之后最后一根轴 dimA[-1]FLOAT、DOUBLE、COMPLEX64、COMPLEX128ND-√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----其中mOut是 broadcast 后 A 的上三角下三角拷贝其作用是将实际参与求解的系数矩阵回传方便调用方核对或后续使用。返回值与错误码接口返回aclnnStatus状态码具体参见 aclnn_return_code.md。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、A、xOut、mOut 中存在空指针ACLNN_ERR_PARAM_INVALID161002self、A、xOut、mOut 的数据类型和数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self、A、xOut、mOut 的 shape 不符合约束注意上述文档描述的mOut空指针场景在源码实现中做了放宽处理。查看 aclnn_triangular_solve.cpp 可见mOut为空时被允许if (mOut ! nullptr)分支跳过 dtype 与 shape 校验并且 L2 层 UTcase_nullptr也验证了OUTPUT(X_desc, nullptr)返回ACLNN_SUCCESS。第二段接口aclnnTriangularSolve 参数详解第二段接口的参数说明如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入由第一段接口aclnnTriangularSolveGetWorkspaceSize获取的 workspace 大小executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream第二段接口的返回值同样是aclnnStatus具体参见 aclnn_return_code.md。其实现非常简洁在 aclnn_triangular_solve.cpp 中只是调用框架的CommonOpExecutorRun完成计算aclnnStatus aclnnTriangularSolve(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream) { // 固定写法调用框架能力完成计算 L2_DFX_PHASE_2(aclnnTriangularSolve); return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); }约束说明确定性aclnnTriangularSolve默认为确定性实现多次运行同一输入可得到确定结果相关背景可参考 determinism_compute.md。数据类型常规模式下支持 FLOAT、DOUBLE、COMPLEX64、COMPLEX128unitriangular为 true 时仅支持 FLOAT。shape 约束A 必须为方阵最后两维相等self 与 A 维度范围为 2~8二者除最后两维外需满足 broadcast 关系。空 Tensor源码中 aclnn_triangular_solve.cpp 对A-IsEmpty() || self-IsEmpty()的情况直接返回ACLNN_SUCCESS且 workspace 置 0UTcase_empty也验证了空 batch 场景可以正常通过。源码实现深度剖析理解两段式接口背后的执行流程有助于正确使用参数。以下链路均可在 aclnn_triangular_solve.cpp 中逐行找到创建执行器与入参校验CREATE_EXECUTOR()创建OpExecutor随后CheckParams依次执行空指针检查CheckNotNull、数据类型检查CheckDtypeValid对照DTYPE_SUPPORT_LIST、shape 检查CheckShape要求 A 为方阵、self 与 A 的最后第二维相等、非末两维可 broadcast、xOut/mOut shape 与 broadcast 结果一致。连续化处理对 self 与 A 分别调用l0op::Contiguous转成连续 Tensor保证后续算子可处理对应“支持非连续 Tensor”的能力详见 non_contiguous_tensor.md。Broadcast 展开若输入 shape 与输出 shape 不同则调用l0op::BroadcastTo将 self broadcast 到 xOut 的 shape、将 A broadcast 到 mOut 的 shapeBroadcastTensor函数broadcast 后的 A 通过l0op::ViewCopy拷贝进 mOut 输出。unitriangular 归一化当unitriangular为 true 时CalcuUnitriangular通过组合 Eye、OnesLike、Sub、Mul、Add 等基础算子把 A 的主对角线强制归一为 1aUnitriangular A * (1 - eye) eye且仅支持 FLOAT 类型。核心求解调用l0op::TriangularSolve(selfBroadcast, aBroadcast, upper, transpose, xOut, executor)完成三角方程组求解。结果回写与 workspace 计算通过l0op::ViewCopy将结果拷贝到 xOut支持非连续输出最后uniqueExecutor-GetWorkspaceSize()返回 workspace 大小并ReleaseTo释放执行器句柄给第二段接口。核心求解的 L0 层实现在 triangular_solve.cpp通过ADD_TO_LAUNCHER_LIST_AICPU宏把任务投递到 AICPU 侧的MatrixTriangularSolvekernel并将三个计算属性做映射——lower !upper、adjoint transpose输入为 A 与 self输出为 Xstatic internal::AicpuTaskSpace space(MatrixTriangularSolve, ge::DEPEND_IN_SHAPE, true); auto ret ADD_TO_LAUNCHER_LIST_AICPU(TriangularSolve, OP_ATTR_NAMES({ lower, adjoint }), OP_INPUT(A, self), OP_OUTPUT(X), OP_ATTR(!upper, transpose));由此可以推断uppertrue等价于 AICPU kernel 的lowerfalsetransposetrue等价于adjointtrue两段接口之间通过属性映射完成了语义对齐。测试用例佐证仓库为该算子提供了 L0 与 L2 两层单元测试可作为行为契约的参考L0 层test_aclnn_l0_triangular_solve.cpp直接调用l0op::TriangularSolve覆盖upper为 true/false、transpose为 true、FLOAT/DOUBLE/COMPLEX64/COMPLEX128 四种数据类型、batch3 维、4 维 batch、单列{3,1}以及空 batch{0,3,4}等场景验证输出 shape 与预期一致。L2 层test_aclnn_l2_triangular_solve.cpp通过OP_API_UT框架调用完整两段接口覆盖正常求解{1,1,3,3}A 与{1,1,3,4}b、空指针返回ACLNN_ERR_INNER_NULLPTR、非法数据类型FLOAT16 返回ACLNN_ERR_PARAM_INVALID、类型不一致、维度小于 2 / 大于 8、A 非方阵、矩阵尺寸不匹配、broadcast 失败/成功、空 Tensor、transposetrue、unitriangulartrue及unitriangular与 DOUBLE 类型组合失败等边界场景。这些测试用例从参数校验和 shape 推导两个维度印证了接口文档中的约束表格是排查调用问题时最直接的参考。调用示例可直接编译运行的完整代码以下示例来自接口文档对应仓库中的可运行版本见 examples/test_aclnn_triangular_solve.cpp。示例求解selfb[1, 2, 3]ᵀshape 为{3, 1}A [[1,2,3],[0,4,5],[0,0,6]]shape 为{3, 3}上三角方阵upper truetranspose falseunitriangular false#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_triangular_solve.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; } 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::vectorint64_t selfShape {3, 1}; std::vectorint64_t aShape {3, 3}; std::vectorint64_t xOutShape {3, 1}; std::vectorint64_t mOutShape {3, 3}; void* selfDeviceAddr nullptr; void* aDeviceAddr nullptr; void* xOutDeviceAddr nullptr; void* mOutDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* A nullptr; aclTensor* xOut nullptr; aclTensor* mOut nullptr; bool upper true; bool transpose false; bool unitriangular false; std::vectorfloat selfHostData {1, 2, 3}; std::vectorfloat aHostData {1, 2, 3, 0, 4, 5, 0, 0, 6}; std::vectorfloat xOutHostData {-0.2500, -0.1250, 0.5000}; std::vectorfloat mOutHostData {1, 2, 3, 0, 4, 5, 0, 0, 6}; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建A aclTensor ret CreateAclTensor(aHostData, aShape, aDeviceAddr, aclDataType::ACL_FLOAT, A); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建xOut aclTensor ret CreateAclTensor(xOutHostData, xOutShape, xOutDeviceAddr, aclDataType::ACL_FLOAT, xOut); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建mOut aclTensor ret CreateAclTensor(mOutHostData, mOutShape, mOutDeviceAddr, aclDataType::ACL_FLOAT, mOut); CHECK_RET(ret ACL_SUCCESS, return ret); // 3.调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnTriangularSolve第一段接口 ret aclnnTriangularSolveGetWorkspaceSize(self, A, upper, transpose, unitriangular, xOut, mOut, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnTriangularSolveGetWorkspaceSize 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); } // 调用aclnnTriangularSolve第二段接口 ret aclnnTriangularSolve(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnTriangularSolve 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 xSize GetShapeSize(xOutShape); std::vectorfloat resultData(xSize, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), xOutDeviceAddr, xSize * 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 xSize; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } auto mSize GetShapeSize(mOutShape); std::vectorfloat mResultData(mSize, 0); ret aclrtMemcpy(mResultData.data(), mResultData.size() * sizeof(mResultData[0]), mOutDeviceAddr, mSize * sizeof(mResultData[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 mSize; i) { LOG_PRINT(result[%ld] is: %f\n, i, mResultData[i]); } // 6.释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(A); aclDestroyTensor(xOut); aclDestroyTensor(mOut); // 7.释放device资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(aDeviceAddr); aclrtFree(xOutDeviceAddr); aclrtFree(mOutDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }预期结果xOut 求解AXb得[-0.25, -0.125, 0.5]ᵀmOut 输出 broadcast 后 A 的上三角拷贝即[1,2,3,0,4,5,0,0,6]。示例中xOutHostData与mOutHostData仅为占位初始化实际值以算子执行后的回拷结果为准。编译与运行指引上述示例的具体编译与执行过程请参考仓库文档 compile_and_run_sample.md。仓库中该目录的构建入口为 math/triangular_solve/CMakeLists.txtop_api 源文件由 op_host/CMakeLists.txt 组织UT 用例的构建入口为 tests/ut/op_api/CMakeLists.txt。小结aclnnTriangularSolve是 CANN ops-math 仓库中三角方程组求解能力的标准 aclnn 封装具备完整的参数校验、broadcast 支持、unitriangular 归一化与确定性计算能力。通过本文介绍的两段式调用范式、参数语义表、错误码定义以及 aclnn_triangular_solve.cpp 的源码链路你可以快速将该接口集成进自己的推理或训练框架也可以在遇到 shape/类型报错时按 test_aclnn_l2_triangular_solve.cpp 的边界用例逐一排查入参合法性。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
