CANN ops-math 中 Tril 下三角算子的完整解析数学定义、aclnn 调用流程与 AICore 内核实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math在 ops-math 这个 CANN 数学类算子库中Trillower triangular下三角算子用于对张量最后两个维度执行保留下三角、右上部分置零的变换是注意力掩码attention mask、Cholesky 类数值算法、以及各类矩阵预处理流程中常见的基础转换原语。本篇基于仓库中 conversion/tril/README.md 及其配套 API 文档 conversion/tril/docs/aclnnTrilaclnnInplaceTril.md 展开结合算子的 host 端定义、AICore 内核源码与测试 golden讲清楚 Tril 的数学语义、diagonal参数行为、aclnn 两段式接口调用方式、错误码以及内核侧根据对角线位置自动选择的四类分块策略帮助你在 NPU 上正确接入并验证该算子。功能定义与 diagonal 语义Tril的算子功能将输入self张量的最后二维按 shape 从左向右数沿对角线的右上部分置零。参数diagonal可正可负默认为 0diagonal为正数主对角线向右上方向移动即被保留的下三角区域变大置零区域变小diagonal为负数主对角线向左下方向移动即被保留的下三角区域变小。用i表示倒数第二维行索引、j表示最后一维列索引、d表示diagonal在二维坐标中i d j表示恰在对角线上。计算公式为$$ \text{对角线及其左下方}id \ge j\text{} out_{i,j} self_{i,j} $$ $$ \text{对角线右上方}id j\text{} out_{i,j} 0 $$文档给出的三个示例直观展示了diagonal的影响以下示例矩阵为 conversion/tril/README.md 原文中的示例设self [[9, 6, 3], [1, 2, 3], [3, 4, 1]]diagonal0时[[9, 0, 0], [1, 2, 0], [3, 4, 1]]diagonal1时[[9, 6, 0], [1, 2, 3], [3, 4, 1]]diagonal-1时[[0, 0, 0], [1, 0, 0], [3, 4, 0]]说明conversion/tril/README.md 原文示例中函数名写作triu(...)结合上文功能说明与diagonal的行为可以确认这里指的是本算子tril本身triu是其姊妹算子见 conversion/triu/README.md。产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品√各产品支持的数据类型存在差异摘自 conversion/tril/README.mdAtlas 训练系列产品DOUBLE、FLOAT、FLOAT16、INT16、INT32、INT64、INT8、UINT16、UINT32、UINT64、UINT8、BOOLAtlas A2 训练/推理系列、Atlas A3 训练/推理系列在上述基础上增加 BFLOAT16Ascend 950PR/Ascend 950DT在上述基础上进一步增加 COMPLEX32、COMPLEX64仅该产品支持复数。参数说明算子级参数定义摘自 conversion/tril/README.md 参数说明表参数名输入/输出/属性描述数据类型数据格式self输入待进行 tril 计算的入参公式中的selfFLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BOOL、COMPLEX32、COMPLEX64NDdiagonal输入待进行 tril 计算的入参公式中的diagonalINT64NDout输出待进行 tril 计算的出参公式中的outFLOAT、FLOAT16、DOUBLE、BFLOAT16、INT8、INT16、INT32、INT64、UINT8、BOOLND需要注意out的可用数据类型比self少输出不支持 COMPLEX32/COMPLEX64、UINT32/UINT64。而算子在 GE 图层面的定义中输入x与输出y均声明了 15 种数据类型含复数见 conversion/tril/op_host/tril_def.cppthis-Input(x) .ParamType(REQUIRED) .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT64, ge::DT_UINT64, ge::DT_INT32, ge::DT_UINT32, ge::DT_INT16, ge::DT_UINT16, ge::DT_INT8, ge::DT_UINT8, ge::DT_DOUBLE, ge::DT_COMPLEX32, ge::DT_COMPLEX64, ge::DT_BOOL}) ... this-Attr(diagonal).AttrType(OPTIONAL).Int(0); // diagonal 为可选属性默认 0从 tril_def.cpp 可以看到diagonal是OPTIONAL属性默认值为 0与文档参数 diagonal 默认为零的描述一致同时 tril_def.cpp 中通过DynamicCompileStaticFlag(true)、DynamicShapeSupportFlag(true)声明了静态编译 动态形状支持并将内核文件注册为tril_apt同时为ascend950与ascend350两个架构分别注册了 AICore 配置——这与产品支持表中950 系列与 A2/A3 系列支持相吻合。约束说明方面conversion/tril/README.md 明确写为无即没有额外的调用约束aclnn 文档 conversion/tril/docs/aclnnTrilaclnnInplaceTril.md 中则补充了一条确定性说明aclnnTril 与 aclnnInplaceTril 默认为确定性实现。aclnn 两段式接口aclnnTril 与 aclnnInplaceTrilTril 面向用户暴露的 API 为两组两段式two-phase接口完整原型、逐参数说明和调用示例见 conversion/tril/docs/aclnnTrilaclnnInplaceTril.md两段式机制的通用背景可参考 docs/zh/context/two_phase_api.md。aclnnTril 与 aclnnInplaceTril 的选择aclnnTril需要新建一个输出张量对象out来存储计算结果aclnnInplaceTril无需新建输出张量直接在输入张量的内存中写入计算结果适合显存敏感或无需保留原矩阵的场景。函数原型如下头文件声明见 conversion/tril/op_api/aclnn_tril.h// 第一段入参校验 计算 workspace 大小产出 executor aclnnStatus aclnnTrilGetWorkspaceSize( const aclTensor* self, int64_t diagonal, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor); // 第二段在 stream 上执行计算 aclnnStatus aclnnTril( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream); aclnnStatus aclnnInplaceTrilGetWorkspaceSize( const aclTensor* selfRef, int64_t diagonal, uint64_t* workspaceSize, aclOpExecutor** executor); aclnnStatus aclnnInplaceTril( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream);aclnnTrilGetWorkspaceSize 参数与首段错误码参数名输入/输出描述数据类型维度(shape)非连续 tensorself (aclTensor*)输入待转换的目标张量公式中的 selfDOUBLE/FLOAT/FLOAT16/INT16/INT32/INT64/INT8/UINT16/UINT32/UINT64/UINT8/BOOL/BFLOAT16/COMPLEX32/COMPLEX642-8√diagonal (int64_t)输入对角线的位置int64_t--out (aclTensor*)输入输出张量需预先创建同 self2-8-workspaceSize (uint64_t*)出参需在 Device 侧申请的 workspace 大小---executor (aclOpExecutor**)出参op 执行器包含算子计算流程---其中输入self支持非连续非连续 stridesTensor张量维度要求rank 在 28 之间COMPLEX32/COMPLEX64 仅 Ascend 950PR/950DT 支持Atlas 推理/训练系列910/310p 代际不支持 BFLOAT16。第一段接口完成入参校验典型报错场景摘自 conversion/tril/docs/aclnnTrilaclnnInplaceTril.md返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型不在支持范围内ACLNN_ERR_PARAM_INVALID161002self 与 out 数据类型不一致ACLNN_ERR_PARAM_INVALID161002self、out 的 shape 不一致ACLNN_ERR_PARAM_INVALID161002self、out 的数据格式不一致ACLNN_ERR_PARAM_INVALID161002self 维度大于 8或小于 2aclnnInplaceTrilGetWorkspaceSize的入参校验逻辑类似但仅校验selfRef无 out空指针报 161001selfRef数据类型不在支持范围内、或维度大于 8/小于 2 时报 161002。第二段接口aclnnTril/aclnnInplaceTril只接收workspace、workspaceSize、executor、stream四个参数均返回aclnnStatus状态码返回码含义见 docs/zh/context/aclnn_return_code.md。从 aclnn_tril.h 中的计算图注释可以推断 aclnn 调用路径的内部编排self先经过l0::Contiguous这也是输入支持非连续 Tensor 的原因再进入l0::Tril执行下三角计算最后经l0::ViewCopy写到out。Level 0 接口l0op::Tril的声明位于 conversion/tril/op_api/tril.h。调用示例一个可复制的完整流程仓库提供了完整可编译的样例 conversion/tril/examples/test_aclnn_tril.cpp编译与运行方式参见 docs/zh/context/compile_and_run_sample.md。样例用INT32类型的 3×3 矩阵演示了aclnnTril与aclnnInplaceTril的完整调用链核心步骤如下省略部分注释完整代码见源文件#include acl/acl.h #include aclnnop/aclnn_tril.h #include iostream #include vector // ... Init(deviceId, stream)aclInit / aclrtSetDevice / aclrtCreateStream // ... CreateAclTensor()aclrtMalloc aclrtMemcpy(H2D) 计算连续 strides aclCreateTensor int main() { int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, return ret); // 1. 构造输入/输出3x3 INT32diagonal 0 std::vectorint64_t selfShape {3, 3}; std::vectorint64_t outShape {3, 3}; std::vectorint selfHostData {1, 2, 3, 4, 5, 6, 7, 8, 9}; int diagonal 0; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_INT32, self); ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_INT32, out); // 2. 两段式调用 aclnnTril uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnTrilGetWorkspaceSize(self, diagonal, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, return ret); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, return ret); } ret aclnnTril(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 两段式调用 aclnnInplaceTril直接改写 self 内存 uint64_t inplaceWorkspaceSize 0; aclOpExecutor* inplaceExecutor; ret aclnnInplaceTrilGetWorkspaceSize(self, diagonal, inplaceWorkspaceSize, inplaceExecutor); CHECK_RET(ret ACL_SUCCESS, return ret); // ... 同样按需申请 workspace 后调用 aclnnInplaceTril(...) // 4. 同步并拷回结果 ret aclrtSynchronizeStream(stream); // aclrtMemcpy(outDeviceAddr - host, ACL_MEMCPY_DEVICE_TO_HOST) 后打印 result[i] // aclrtMemcpy(selfDeviceAddr - host) 后打印 inplaceResult[i] // 5. 释放资源aclDestroyTensor / aclrtFree / aclrtDestroyStream / aclrtResetDevice / aclFinalize return 0; }该样例的几个值得注意的工程细节workspace 按需申请workspaceSize由第一段接口算出只有 0时才调用aclrtMallocaclnnTril的 workspace 地址允许为nullptr传入strides 按连续布局计算CreateAclTensor中从右向左累乘 shape 生成 strides若输入本身非连续则按实际 strides 构造aclTensor即可算子侧会用Contiguous兼容inplace 版本复用同一输入aclnnInplaceTril不接收out计算直接落在self的设备内存上因此样例最后分别拷回out和self两份结果用于对照。按diagonal0、输入[[1,2,3],[4,5,6],[7,8,9]]推演result[i]的输出序列应为1,0,0,4,5,0,7,8,9。内核实现按 diagonal 位置分流的四路 AICore 内核Tril 的 AICore 内核入口非常精简见 conversion/tril/op_kernel/tril_apt.cpp#include ../triu/arch35/triangulator_base.h using namespace Triangulator; extern C __global__ __aicore__ void tril(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) { KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); triangulatorEntrytrue(x, y, workspace, tiling); // IS_LOWER true }可以看到两点事实内核以KERNEL_TYPE_AIV_ONLY运行纯 AIV 核任务Tril 与 conversion/triu 的 Triu 共用同一套三角化内核基础设施 conversion/triu/op_kernel/arch35/triangulator_base.h仅通过模板参数IS_LOWERtrue/false区分上/下三角。tiling 注册Tiling4Trilufalse, falsehost 侧 tiling 在 conversion/tril/op_host/arch35/tril_tiling.cpp 中注册IMPL_OP_OPTILING(Tril) .Tiling(Tiling4Trilufalse, false) .TilingParseTriluCompileInfo(TilingPrepare4Trilu);Tril 与 Triu 共用Tiling4Trilu模板Triu 对应 conversion/triu/op_host/arch35/triu_tiling.cpp 中的注册各架构的静态内核清单由 conversion/tril/op_host/config/ascend950/tril_binary.json、conversion/tril/op_host/config/ascend350/tril_binary.json 描述int8/int16/int32/int64 各一个 bin编译键简化规则见 conversion/tril/op_host/config/ascend950/tril_simplified_key.ini。根据对角线位置选择分块策略triangulatorEntryIS_LOWER是真正的分发点见 triangulator_base.h。它的注释解释了 tiling key 的编码方式——个位表示位宽1/2/4/8 对应 int8~int64十位表示模板策略// 个位位宽 1 int8_t 2 int16_t 4 int32_t 8 int64_t // 十位模板 0 output_zero 1 output_input 2 output_normal 3 tiny_shape 4 split_outer即 host 端 tiling 会先根据diagonal与矩阵形状的关系做静态判定将计算退化或归类为以下四类之一内核按 tiling key 分流策略ten 位触发条件从源码结构看可按 diagonal 推断内核实例OUTPUT_ZERO1/2/4/8对角线落在矩阵最左侧之外diagonal足够负整行都位于右上区域时整个输出恒为 0内核直接写零无需读输入SplitAllT, OUTPUT_ZERO_MODEOUTPUT_INPUT11/12/14/18对角线落在矩阵最右侧之外diagonal足够正时全部元素都位于对角线左下方输出等于输入内核退化为纯拷贝SplitAllT, OUTPUT_INPUT_MODEOUTPUT_NORMAL21/22/24/28一般情况按行/列分块对角线两侧分别执行保留原值与写零两条数据通路SplitRowAndColT, IS_LOWERTINY_SHAPE31/32/34/38形状过小、分块收益有限时走专用小形状通路SplitTinyShapeTSPLIT_MEDIUM41/42/44/48中等形状的外层切分通路SplitMediumShapeT, IS_LOWER其中IS_LOWER模板参数决定保留哪一侧tril传true保留左下方i d jtriu传false保留右上方。这种先由 host 判定 diagonal 的相对位置、再把退化情况编译为独立 bin的设计意味着diagonal0之外的极端取值全零输出/恒等拷贝不会浪费算力做逐元素比较——这是对文档中diagonal 可正可负参数的一种高效落地。复数类型在内核侧按实部/虚部交错处理这一点可从测试 golden 直接印证见 conversion/tril/tests/assets/golden.pycomplex32输入被np.split(x, 2, axis-1)拆出实部与虚部后分别执行np.tril再沿最后一维stack回去其他类型直接np.tril(x, diagonal)。测试与验证仓库为 Tril 提供了多层测试资产kernel UTconversion/tril/tests/ut/op_kernel/test_tril.cpp配套 tiling 头文件 conversion/tril/tests/ut/op_kernel/tril_tiling.haclnn UTconversion/tril/tests/ut/op_api/test_aclnn_tril.cppinfershape UTconversion/tril/tests/ut/op_host/test_tril_infershape.cpp验证形状推导对应self维度 28 的约束arch35 tiling UTconversion/tril/tests/ut/op_host/arch35/test_tril_tiling.cpp验证 tiling 参数生成系统测试conversion/tril/tests/st/aclnnTril/executor_aclnnTril.py 与用例配置 conversion/tril/tests/st/aclnnTril/atk_aclnnTril.json、arch35 的 CSV 用例 conversion/tril/tests/st/arch35/ttk_kernel_tril_st.csv 等。golden 参考实现conversion/tril/tests/assets/golden.py明确以 NumPy/PyTorch 为基准def tril_golden(x, diagonal: int 0, **kwargs): if kwargs.get(input_dtypes) and kwargs[input_dtypes][0] complex32: real, imag np.split(x, 2, axis-1) real np.tril(np.squeeze(real, axis-1), diagonal) imag np.tril(np.squeeze(imag, axis-1), diagonal) golden np.stack((real, imag), axis-1) else: golden np.tril(x, diagonal) return golden def aclnn_tril_golden(self, diagonal0, outNone, **kwargs): return [torch.tril(self, diagonal)]也就是说kernel 测试以np.tril(x, diagonal)为期望值aclnn 测试以torch.tril(self, diagonal)为期望值——两者语义与文档中的i d j保留规则完全一致可作为数值正确性的权威参照。使用建议与注意事项接口选型结果需保留原矩阵时用aclnnTril需预建out且out与self的 dtype、shape、format 必须一致否则首段接口返回 161002原地改写、省一份显存时用aclnnInplaceTril维度约束输入 rank 必须在 28 之间对高维张量变换只作用于最后两个维度前面的 batch 维保持不变非连续输入aclTensor支持非连续 strides内部会先执行Contiguous见 aclnn_tril.h 的计算图注释但会增加一次拷贝开销能传连续张量时尽量传连续张量数据类型匹配产品BFLOAT16 在 Atlas 训练系列910上不可用COMPLEX32/64 仅在 Ascend 950PR/950DT 上可用且out参数不支持复数与 UINT32/UINT64见参数说明表diagonal 的极端取值从内核分路代码看diagonal超出矩阵范围时结果会退化为全零或恒等拷贝行为仍然确定该算子为默认确定性实现但请确保diagonal语义符合业务预期验证手段数值结果可用torch.tril/np.tril对照调用参数问题按错误码表161001/161002逐项排查即可。小结Tril 在 ops-math 中是一个语义清晰、实现高效的矩阵转换算子host 侧通过 tril_def.cpp 声明 15 种 dtype、ND 格式与 2~8 维支持aclnn 侧提供aclnnTril/aclnnInplaceTril两套两段式接口内核侧则由 triangulator_base.h 中的通用三角化框架按diagonal与形状的关系在全零输出、恒等拷贝、常规行列分块、小/中形状专用通路之间自动择优。配合 examples/test_aclnn_tril.cpp 的完整样例与多层测试 golden开发者可以在 NPU 上快速完成接入、调参与正确性验证。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考