CANN ops-math 算子 aclnnDiagFlat 两段式接口详解对角线张量生成原理、参数约束与完整调用示例【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本篇技术指南聚焦 CANN ops-math 仓库中 conversion 类算子DiagFlat的 aclnn 接口aclnnDiagFlat完整讲解其先扁平化、再构对角线的算子语义、两段式接口aclnnDiagFlatGetWorkspaceSizeaclnnDiagFlat的每个参数与返回码、产品支持范围、源码级实现链路并给出可直接编译运行的单测风格示例代码。读完本文你将掌握在 Atlas/Ascend 系列产品上通过 aclnn 接口正确、高效地调用 DiagFlat 算子完成对角线张量构造的完整实战方法。一、算子功能与数学语义DiagFlat 的功能是生成对角线张量其语义定义位于 conversion/diag_flat/docs/aclnnDiagFlat.md 的功能说明中若输入self为一维张量则返回二维张量self中的元素作为对角线值依次填入输出张量的对角线位置其余位置填充 0若输入self为二维及以上张量则先将其扁平化flatten 为一维再按第一种场景处理即先扁平化、再构造对角线。举例而言输入self [[2, 5]]shape 为[1, 2]扁平化后为[2, 5]当diagonal 1时输出 shape 为[3, 3]2与5被依次放置在主对角线上方偏移 1 的对角线上得到0 2 0 0 0 5 0 0 0该语义与 PyTorch 的torch.diag(input, diagonal)保持一致仓库的算子级测试 conversion/diag_flat/tests/st/aclnnDiagFlat/executor_aclnnDiagFlat.py 正是以torch.diag(input_tensor, diagonal)作为期望结果来比对 NPU 输出。从图原型定义看conversion/diag_flat/op_graph/diag_flat_proto.h该算子与 TensorFlow 的Diag算子兼容输入x、输出y类型保持一致属性diagonal为可选 int默认值为 0。二、产品支持情况aclnnDiagFlat在不同产品上的支持情况如下来源于 conversion/diag_flat/docs/aclnnDiagFlat.md产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持此外算子 READMEconversion/diag_flat/README.md补充说明 Kirin X90、Kirin 9030 处理器系列产品也支持本算子。各产品在数据类型上存在差异详见本文约束与注意事项一节。三、两段式接口总览与其他 aclnn 算子一致aclnnDiagFlat采用两段式接口调用范式详见 docs/zh/context/two_phase_api.md第一段先调用aclnnDiagFlatGetWorkspaceSize完成入参校验、计算流程编排并返回执行计算所需的workspace 大小与包含了算子计算流程的op 执行器executor第二段再调用aclnnDiagFlat传入第一段申请好的 workspace 内存与 executor在指定 Stream 上异步执行计算。3.1 第一段接口原型aclnnStatus aclnnDiagFlatGetWorkspaceSize( const aclTensor* self, int64_t diagonal, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)3.2 第二段接口原型aclnnStatus aclnnDiagFlat( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)四、第一段接口 aclnnDiagFlatGetWorkspaceSize 参数详解4.1 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorselfaclTensor*输入表示填充到对角线的 Tensor-FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL、COMPLEX64、BFLOAT16ND1-8√diagonalint64_t输入用来指定对角线。diagonal 0 表示主对角线diagonal 0 表示主对角线上方的对角线diagonal 0 表示主对角线下方的对角线-INT64ND--outaclTensor*输出输出 Tensor-与 self 保持一致ND1-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----要点解读输入输出支持非连续 Tensorself与out均标注支持非连续 Tensor√。从实现看conversion/diag_flat/op_api/aclnn_diag_flat.cpp第一段接口会先用l0op::Contiguous将输入转换为连续张量计算后再通过l0op::ViewCopy将结果写回可能非连续的out因此调用方无需预先对张量做 contiguity 处理。维度范围 1-8 维self与out的 shape 支持 1~8 维维度超过上限会被第一段接口的 shape 检查拦截。diagonal 的取值语义0 为主对角线正数表示主对角线之上右上方向负数表示主对角线之下左下方向。其绝对值参与输出宽度计算见源码实现走读一节。4.2 返回值与错误码接口返回aclnnStatus状态码具体枚举含义参见 docs/zh/context/aclnn_return_code.md。第一段接口在入参校验阶段报错具体场景如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002diagonal 不在支持的数据类型范围之内从源码实现看conversion/diag_flat/op_api/aclnn_diag_flat.cppCheckParams依次执行空指针检查OP_CHECK_NULL、数据类型合法性检查OP_CHECK_DTYPE_NOT_SUPPORT依据不同 SoC 选择 dtype 支持列表、最大维度检查OP_CHECK_MAX_DIM上限为MAX_SUPPORT_DIMS_NUMS并在检测到 FRACTAL_NZ 存储格式时打印精度风险警告。五、第二段接口 aclnnDiagFlat 参数详解参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnDiagFlatGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同样为aclnnStatus状态码。第二段接口的实现非常简洁在 conversion/diag_flat/op_api/aclnn_diag_flat.cpp 中它直接调用框架公共能力CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成实际计算即真正执行第一段编排好的计算图。因此调用方必须保证两段接口成对使用且 workspace 内存必须根据第一段返回的workspaceSize在 Device 侧申请。六、约束与注意事项依据 conversion/diag_flat/docs/aclnnDiagFlat.md 及算子 READMEconversion/diag_flat/README.md使用aclnnDiagFlat需要注意确定性计算aclnnDiagFlat默认采用确定性实现多次运行同一输入可得到可重复的结果便于结果校验与调试。暂不支持反向当前接口aclnnDiagFlat暂不具备反向grad功能梯度场景需另行处理。数据类型平台差异Kirin X90 / Kirin 9030 处理器系列产品self、out数据类型不支持 COMPLEX、BFLOAT16Atlas 训练系列产品、Atlas 推理系列产品不支持 BFLOAT16。BFLOAT16 与 BOOL 的平台区分从 conversion/diag_flat/op_host/diag_flat_def.cpp 的算子定义可见Ascend 910/310P 配置中未包含 BF16 与 BOOL而 Ascend 910B含 910_93配置加入了 BF16Kirin 系列配置则仅支持整型与浮点类型不含 COMPLEX64。因此具体可用的 dtype 集合以实际部署的 SoC 版本为准。七、源码实现走读从接口到内核要深入理解aclnnDiagFlat的行为可以沿着仓库源码按API 层 → 图定义层 → Host Tiling 层 → 内核层逐层梳理。7.1 API 层两段式接口的实现conversion/diag_flat/op_api/aclnn_diag_flat.cpp 中第一段接口aclnnDiagFlatGetWorkspaceSize的执行流程如下创建 OpExecutorCREATE_EXECUTOR()执行CheckParams入参校验空指针、dtype、shape、format空输入特殊处理当self-IsEmpty()时若|diagonal| 0通过FillScalar(out, 0, ...)将输出整体填充为 0 后直接返回否则 workspace 置 0 直接返回成功将输入self转换为连续张量l0op::Contiguous调用底层 l0 算子执行核心计算l0op::DiagFlat(selfContiguous, diagonal, ...)将中间结果转换为输出out的数据类型l0op::Cast将结果拷贝到可能非连续的输出l0op::ViewCopy通过uniqueExecutor-GetWorkspaceSize()汇总整个计算链所需的 workspace 大小并返回。第二段接口aclnnDiagFlat则通过CommonOpExecutorRun在指定 Stream 上异步调度执行。整个流程体现了 aclnn 算子 API 层的通用编排模式校验 → 连续化 → 计算 → 类型转换 → 视图拷贝 → 获取 workspace。7.2 图定义层算子注册与 InferShape算子原型注册在 conversion/diag_flat/op_graph/diag_flat_proto.h输入x、输出y与输入同类型、可选属性diagonalint默认 0。输出 shape 的推导规则在 conversion/diag_flat/op_host/diag_flat_infershape.cpp 中实现这是理解算子语义的关键输出宽度 输入元素总数 |diagonal| 输出shape [输出宽度, 输出宽度] // 二维方阵输入元素总数由inputShape-GetShapeSize()计算得到即扁平化后的元素个数未知 rank 或未知 shape 的动态场景下输出推导为[-1, -1]的二维动态 shape输出数据类型与输入保持一致InferDataType直接透传。例如输入[1, 2]共 2 个元素diagonal 1时输出宽度为2 1 3即示例中的[3, 3]diagonal 0时输出为[2, 2]。算子定义 conversion/diag_flat/op_host/diag_flat_def.cpp 中为各 SoC 配置了 AICore 计算配置并声明了DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)、DynamicFormatFlag(true)、PrecisionReduceFlag(true)等动态能力说明该算子支持动态 shape/rank 场景Ascend 950 配置还通过ExtendCfgInfo(opFile.value, diag_flat_apt)指定了 APT算子性能模板内核实现。7.3 Host Tiling 层与内核层Tiling 逻辑conversion/diag_flat/op_host/arch35/diag_flat_tiling.cpp负责在多核场景下切分任务根据平台信息GetCoreNumAiv、UB 内存大小与输入规模计算每个核处理的数据量、tile 长度与本地内存大小并将结果写入 TilingData。值得注意的是arch35 的 tiling 函数以非静态接口导出见 conversion/diag_flat/op_host/arch35/diag_flat_tiling.h可被 DiagV2 算子在一维输入场景下复用体现了算子间代码复用的设计。内核实现分布在 conversion/diag_flat/op_kernel 目录下核心思路分为两步先对输出 GM 内存做多核并行清零conversion/diag_flat/op_kernel/diag_flat_common.h 中的DiagFlatMemSetZero再按对角线偏移量将输入数据写入对应的对角位置diag_flat_nd_to_2d系列变体针对不同数据类型与规模提供差异化实现。内核使用 AscendC 编程范式通过 TPipe 队列、SyncAll核间同步完成多核协作。八、完整调用示例以下示例来自 conversion/diag_flat/docs/aclnnDiagFlat.md与仓库中的可执行样例 conversion/diag_flat/examples/test_aclnn_diag_flat.cpp 基本一致输入self为 shape[1, 2]的 FLOAT 张量元素{2, 5}diagonal 1输出out为 shape[3, 3]的零张量。预期输出为[[0, 2, 0], [0, 0, 5], [0, 0, 0]]。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_diag_flat.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 {1, 2}; std::vectorint64_t outShape {3, 3}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {2, 5}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0}; int64_t diagonalVal 1; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); 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); // 3.调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnDiagFlat第一段接口 ret aclnnDiagFlatGetWorkspaceSize(self, diagonalVal, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnDiagFlatGetWorkspaceSize 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); } // 调用aclnnDiagFlat第二段接口 ret aclnnDiagFlat(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnDiagFlat 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(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: %f\n, i, resultData[i]); } // 6.释放aclTensor需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); // 7.释放device资源 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中的 7 个步骤是 aclnn 算子调用的通用模板资源初始化 → 构造 aclTensor → 两段式接口调用 → 同步等待 → 取回结果 → 释放 aclTensor → 释放 device 资源。其中CreateAclTensor模板同时承担了 device 内存申请、host→device 拷贝、连续 strides 计算与aclCreateTensor创建等职责可复用于其他算子示例。具体编译与执行过程可参考 docs/zh/context/compile_and_run_sample.md。九、测试与验证仓库为aclnnDiagFlat提供了算子级ST与单测UT两层验证算子级对照测试conversion/diag_flat/tests/st/aclnnDiagFlat/executor_aclnnDiagFlat.py 基于 ATK 框架实现读取kwargs中的self与diagonal在 NPU 上执行torch.diag(input_tensor, diagonal)作为期望输出用于与算子实际输出做数值比对测试数据配置见 conversion/diag_flat/tests/st/aclnnDiagFlat/atk_aclnnDiagFlat.json另有 conversion/diag_flat/tests/st/arch35/ttk_diag_flat_st.csv 提供 arch35 平台的用例参数。Host 侧单测conversion/diag_flat/tests/ut/op_host/test_diag_flat_infershape.cpp 验证 InferShape 推导逻辑conversion/diag_flat/tests/ut/op_host/arch22/test_diag_flat_tiling.cpp 验证 Tiling 计算。结果生成脚本conversion/diag_flat/tests/assets/golden.py 用于生成参考golden结果供比对使用。十、延伸阅读docs/zh/context/two_phase_api.mdaclnn 两段式接口通用说明docs/zh/context/aclnn_return_code.mdaclnn 返回码含义docs/zh/context/compile_and_run_sample.md样例编译与运行指引docs/zh/context/non_contiguous_tensor.md非连续 Tensor 的支持说明conversion/diag_flat/README.mdDiagFlat 算子 README含调用方式索引conversion/diag_flat/examples/test_aclnn_diag_flat.cpp可编译运行的调用样例。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-math 算子 aclnnDiagFlat 两段式接口详解：对角线张量生成原理、参数约束与完整调用示例
CANN ops-math 算子 aclnnDiagFlat 两段式接口详解对角线张量生成原理、参数约束与完整调用示例【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本篇技术指南聚焦 CANN ops-math 仓库中 conversion 类算子DiagFlat的 aclnn 接口aclnnDiagFlat完整讲解其先扁平化、再构对角线的算子语义、两段式接口aclnnDiagFlatGetWorkspaceSizeaclnnDiagFlat的每个参数与返回码、产品支持范围、源码级实现链路并给出可直接编译运行的单测风格示例代码。读完本文你将掌握在 Atlas/Ascend 系列产品上通过 aclnn 接口正确、高效地调用 DiagFlat 算子完成对角线张量构造的完整实战方法。一、算子功能与数学语义DiagFlat 的功能是生成对角线张量其语义定义位于 conversion/diag_flat/docs/aclnnDiagFlat.md 的功能说明中若输入self为一维张量则返回二维张量self中的元素作为对角线值依次填入输出张量的对角线位置其余位置填充 0若输入self为二维及以上张量则先将其扁平化flatten 为一维再按第一种场景处理即先扁平化、再构造对角线。举例而言输入self [[2, 5]]shape 为[1, 2]扁平化后为[2, 5]当diagonal 1时输出 shape 为[3, 3]2与5被依次放置在主对角线上方偏移 1 的对角线上得到0 2 0 0 0 5 0 0 0该语义与 PyTorch 的torch.diag(input, diagonal)保持一致仓库的算子级测试 conversion/diag_flat/tests/st/aclnnDiagFlat/executor_aclnnDiagFlat.py 正是以torch.diag(input_tensor, diagonal)作为期望结果来比对 NPU 输出。从图原型定义看conversion/diag_flat/op_graph/diag_flat_proto.h该算子与 TensorFlow 的Diag算子兼容输入x、输出y类型保持一致属性diagonal为可选 int默认值为 0。二、产品支持情况aclnnDiagFlat在不同产品上的支持情况如下来源于 conversion/diag_flat/docs/aclnnDiagFlat.md产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持此外算子 READMEconversion/diag_flat/README.md补充说明 Kirin X90、Kirin 9030 处理器系列产品也支持本算子。各产品在数据类型上存在差异详见本文约束与注意事项一节。三、两段式接口总览与其他 aclnn 算子一致aclnnDiagFlat采用两段式接口调用范式详见 docs/zh/context/two_phase_api.md第一段先调用aclnnDiagFlatGetWorkspaceSize完成入参校验、计算流程编排并返回执行计算所需的workspace 大小与包含了算子计算流程的op 执行器executor第二段再调用aclnnDiagFlat传入第一段申请好的 workspace 内存与 executor在指定 Stream 上异步执行计算。3.1 第一段接口原型aclnnStatus aclnnDiagFlatGetWorkspaceSize( const aclTensor* self, int64_t diagonal, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)3.2 第二段接口原型aclnnStatus aclnnDiagFlat( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)四、第一段接口 aclnnDiagFlatGetWorkspaceSize 参数详解4.1 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorselfaclTensor*输入表示填充到对角线的 Tensor-FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL、COMPLEX64、BFLOAT16ND1-8√diagonalint64_t输入用来指定对角线。diagonal 0 表示主对角线diagonal 0 表示主对角线上方的对角线diagonal 0 表示主对角线下方的对角线-INT64ND--outaclTensor*输出输出 Tensor-与 self 保持一致ND1-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----要点解读输入输出支持非连续 Tensorself与out均标注支持非连续 Tensor√。从实现看conversion/diag_flat/op_api/aclnn_diag_flat.cpp第一段接口会先用l0op::Contiguous将输入转换为连续张量计算后再通过l0op::ViewCopy将结果写回可能非连续的out因此调用方无需预先对张量做 contiguity 处理。维度范围 1-8 维self与out的 shape 支持 1~8 维维度超过上限会被第一段接口的 shape 检查拦截。diagonal 的取值语义0 为主对角线正数表示主对角线之上右上方向负数表示主对角线之下左下方向。其绝对值参与输出宽度计算见源码实现走读一节。4.2 返回值与错误码接口返回aclnnStatus状态码具体枚举含义参见 docs/zh/context/aclnn_return_code.md。第一段接口在入参校验阶段报错具体场景如下返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self 或 out 是空指针ACLNN_ERR_PARAM_INVALID161002self 和 out 的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002diagonal 不在支持的数据类型范围之内从源码实现看conversion/diag_flat/op_api/aclnn_diag_flat.cppCheckParams依次执行空指针检查OP_CHECK_NULL、数据类型合法性检查OP_CHECK_DTYPE_NOT_SUPPORT依据不同 SoC 选择 dtype 支持列表、最大维度检查OP_CHECK_MAX_DIM上限为MAX_SUPPORT_DIMS_NUMS并在检测到 FRACTAL_NZ 存储格式时打印精度风险警告。五、第二段接口 aclnnDiagFlat 参数详解参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnDiagFlatGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream返回值同样为aclnnStatus状态码。第二段接口的实现非常简洁在 conversion/diag_flat/op_api/aclnn_diag_flat.cpp 中它直接调用框架公共能力CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成实际计算即真正执行第一段编排好的计算图。因此调用方必须保证两段接口成对使用且 workspace 内存必须根据第一段返回的workspaceSize在 Device 侧申请。六、约束与注意事项依据 conversion/diag_flat/docs/aclnnDiagFlat.md 及算子 READMEconversion/diag_flat/README.md使用aclnnDiagFlat需要注意确定性计算aclnnDiagFlat默认采用确定性实现多次运行同一输入可得到可重复的结果便于结果校验与调试。暂不支持反向当前接口aclnnDiagFlat暂不具备反向grad功能梯度场景需另行处理。数据类型平台差异Kirin X90 / Kirin 9030 处理器系列产品self、out数据类型不支持 COMPLEX、BFLOAT16Atlas 训练系列产品、Atlas 推理系列产品不支持 BFLOAT16。BFLOAT16 与 BOOL 的平台区分从 conversion/diag_flat/op_host/diag_flat_def.cpp 的算子定义可见Ascend 910/310P 配置中未包含 BF16 与 BOOL而 Ascend 910B含 910_93配置加入了 BF16Kirin 系列配置则仅支持整型与浮点类型不含 COMPLEX64。因此具体可用的 dtype 集合以实际部署的 SoC 版本为准。七、源码实现走读从接口到内核要深入理解aclnnDiagFlat的行为可以沿着仓库源码按API 层 → 图定义层 → Host Tiling 层 → 内核层逐层梳理。7.1 API 层两段式接口的实现conversion/diag_flat/op_api/aclnn_diag_flat.cpp 中第一段接口aclnnDiagFlatGetWorkspaceSize的执行流程如下创建 OpExecutorCREATE_EXECUTOR()执行CheckParams入参校验空指针、dtype、shape、format空输入特殊处理当self-IsEmpty()时若|diagonal| 0通过FillScalar(out, 0, ...)将输出整体填充为 0 后直接返回否则 workspace 置 0 直接返回成功将输入self转换为连续张量l0op::Contiguous调用底层 l0 算子执行核心计算l0op::DiagFlat(selfContiguous, diagonal, ...)将中间结果转换为输出out的数据类型l0op::Cast将结果拷贝到可能非连续的输出l0op::ViewCopy通过uniqueExecutor-GetWorkspaceSize()汇总整个计算链所需的 workspace 大小并返回。第二段接口aclnnDiagFlat则通过CommonOpExecutorRun在指定 Stream 上异步调度执行。整个流程体现了 aclnn 算子 API 层的通用编排模式校验 → 连续化 → 计算 → 类型转换 → 视图拷贝 → 获取 workspace。7.2 图定义层算子注册与 InferShape算子原型注册在 conversion/diag_flat/op_graph/diag_flat_proto.h输入x、输出y与输入同类型、可选属性diagonalint默认 0。输出 shape 的推导规则在 conversion/diag_flat/op_host/diag_flat_infershape.cpp 中实现这是理解算子语义的关键输出宽度 输入元素总数 |diagonal| 输出shape [输出宽度, 输出宽度] // 二维方阵输入元素总数由inputShape-GetShapeSize()计算得到即扁平化后的元素个数未知 rank 或未知 shape 的动态场景下输出推导为[-1, -1]的二维动态 shape输出数据类型与输入保持一致InferDataType直接透传。例如输入[1, 2]共 2 个元素diagonal 1时输出宽度为2 1 3即示例中的[3, 3]diagonal 0时输出为[2, 2]。算子定义 conversion/diag_flat/op_host/diag_flat_def.cpp 中为各 SoC 配置了 AICore 计算配置并声明了DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)、DynamicFormatFlag(true)、PrecisionReduceFlag(true)等动态能力说明该算子支持动态 shape/rank 场景Ascend 950 配置还通过ExtendCfgInfo(opFile.value, diag_flat_apt)指定了 APT算子性能模板内核实现。7.3 Host Tiling 层与内核层Tiling 逻辑conversion/diag_flat/op_host/arch35/diag_flat_tiling.cpp负责在多核场景下切分任务根据平台信息GetCoreNumAiv、UB 内存大小与输入规模计算每个核处理的数据量、tile 长度与本地内存大小并将结果写入 TilingData。值得注意的是arch35 的 tiling 函数以非静态接口导出见 conversion/diag_flat/op_host/arch35/diag_flat_tiling.h可被 DiagV2 算子在一维输入场景下复用体现了算子间代码复用的设计。内核实现分布在 conversion/diag_flat/op_kernel 目录下核心思路分为两步先对输出 GM 内存做多核并行清零conversion/diag_flat/op_kernel/diag_flat_common.h 中的DiagFlatMemSetZero再按对角线偏移量将输入数据写入对应的对角位置diag_flat_nd_to_2d系列变体针对不同数据类型与规模提供差异化实现。内核使用 AscendC 编程范式通过 TPipe 队列、SyncAll核间同步完成多核协作。八、完整调用示例以下示例来自 conversion/diag_flat/docs/aclnnDiagFlat.md与仓库中的可执行样例 conversion/diag_flat/examples/test_aclnn_diag_flat.cpp 基本一致输入self为 shape[1, 2]的 FLOAT 张量元素{2, 5}diagonal 1输出out为 shape[3, 3]的零张量。预期输出为[[0, 2, 0], [0, 0, 5], [0, 0, 0]]。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_diag_flat.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 {1, 2}; std::vectorint64_t outShape {3, 3}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {2, 5}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0}; int64_t diagonalVal 1; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); 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); // 3.调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnDiagFlat第一段接口 ret aclnnDiagFlatGetWorkspaceSize(self, diagonalVal, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnDiagFlatGetWorkspaceSize 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); } // 调用aclnnDiagFlat第二段接口 ret aclnnDiagFlat(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnDiagFlat 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(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: %f\n, i, resultData[i]); } // 6.释放aclTensor需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(out); // 7.释放device资源 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中的 7 个步骤是 aclnn 算子调用的通用模板资源初始化 → 构造 aclTensor → 两段式接口调用 → 同步等待 → 取回结果 → 释放 aclTensor → 释放 device 资源。其中CreateAclTensor模板同时承担了 device 内存申请、host→device 拷贝、连续 strides 计算与aclCreateTensor创建等职责可复用于其他算子示例。具体编译与执行过程可参考 docs/zh/context/compile_and_run_sample.md。九、测试与验证仓库为aclnnDiagFlat提供了算子级ST与单测UT两层验证算子级对照测试conversion/diag_flat/tests/st/aclnnDiagFlat/executor_aclnnDiagFlat.py 基于 ATK 框架实现读取kwargs中的self与diagonal在 NPU 上执行torch.diag(input_tensor, diagonal)作为期望输出用于与算子实际输出做数值比对测试数据配置见 conversion/diag_flat/tests/st/aclnnDiagFlat/atk_aclnnDiagFlat.json另有 conversion/diag_flat/tests/st/arch35/ttk_diag_flat_st.csv 提供 arch35 平台的用例参数。Host 侧单测conversion/diag_flat/tests/ut/op_host/test_diag_flat_infershape.cpp 验证 InferShape 推导逻辑conversion/diag_flat/tests/ut/op_host/arch22/test_diag_flat_tiling.cpp 验证 Tiling 计算。结果生成脚本conversion/diag_flat/tests/assets/golden.py 用于生成参考golden结果供比对使用。十、延伸阅读docs/zh/context/two_phase_api.mdaclnn 两段式接口通用说明docs/zh/context/aclnn_return_code.mdaclnn 返回码含义docs/zh/context/compile_and_run_sample.md样例编译与运行指引docs/zh/context/non_contiguous_tensor.md非连续 Tensor 的支持说明conversion/diag_flat/README.mdDiagFlat 算子 README含调用方式索引conversion/diag_flat/examples/test_aclnn_diag_flat.cpp可编译运行的调用样例。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考