CANN ops-math FillDiagonalV2 算子深度解析对角线填充的算法原理与 aclnnInplaceFillDiagonal 调用实践【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathFillDiagonalV2 是 CANN ops-math 数学算子库中用于「以指定填充值填充 Tensor 对角线」的原地in-place算子其宿主侧通过aclnnInplaceFillDiagonal两段式接口对外暴露广泛服务于矩阵初始化、注意力掩码、单位阵构造等 NPU 加速场景。本文以 conversion/fill_diagonal_v2/README.md 为主线结合算子定义、shape 推导、tiling 与 kernel 源码完整讲解该算子的功能语义、参数约束、接口调用与底层实现读完即可在自有工程中正确配置并调用该算子。产品支持情况FillDiagonalV2 算子在当前仓库中注册了多套 AICore 配置支持情况如下表数据源自 fill_diagonal_v2_def.cpp 中AddConfig的注册结果与 README 声明产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品×Kirin X90 处理器系列产品√Kirin 9030 处理器系列产品√从源码可以进一步印证op_host/fill_diagonal_v2_def.cpp中为ascend910b对应 Atlas A2、ascend910_93对应 Atlas A3、ascend950注册了同一套支持 BFLOAT16 的 AICore 配置为ascend310pAtlas 推理系列、kirinx90、kirin9030注册了另一套不含 BFLOAT16的配置。这与 README 中「Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16」的说明完全一致因此在使用 Kirin 平台时需注意避开 BFLOAT16 输入。功能说明与填充公式算子功能FillDiagonalV2 的功能是以fillValue填充 Tensor 的主对角线区域。它属于原地算子输入selfRef同时作为输出即填充操作直接修改原 Tensor 的数据。填充位置计算公式README 给出了二维场景下的精确定义设矩阵行数为row、列数为col记m min(col, row)wrap为 False 时填充位置为[r, r]其中0 r m。即只填充主对角线元素数量等于min(col, row)。wrap为 True 时填充位置为[r (m 1) * i, r]其中0 r m0 i col // m。即对于高矩阵行数大于列数场景每经过N min(col, row)行就形成一条新的对角线实现“缠绕”式填充最终填充的物理位置由展开后的一维索引(col 1) * i决定。接口文档 aclnnInplaceFillDiagonal.md 给出了等价的一维索引表述wrap为 True 时对于满足(col 1) * i row * col的非负整数i填充位置为[floor((col 1) * i / col), ((col 1) * i) % col]。两种表述在数学上等价后者更贴近 kernel 中按一维步长寻址的实现方式。多维 Tensor 的处理READNE 与接口文档均强调当selfRef维度大于 2 时各维度长度必须相同见下文约束说明。此时可将多维 Tensor 视为等边长超立方体填充逻辑按行主序展开后仍然落到“对角线”位置上。参数说明算子的三个入参定义如下表格内容继承自 README并结合 aclnnInplaceFillDiagonal.md 与 aclnn_fill_diagonal.cpp 中的校验逻辑补充说明参数名输入/输出/属性描述数据类型数据格式selfRef输入/输出张量表示输入/输出张量支持非连续的 TensorBFLOAT16、FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOLNDfillValue输入属性表示填充值数据类型需要是可转换为 FLOAT 的数据类型可转换为 FLOAT 的数据类型-wrap输入属性表示填充方式对于高矩阵行数 row 大于列数 col若 wrap 为 True每经过 N 行形成一条新的对角线其中 N min(col, row)BOOL-补充说明Kirin X90/Kirin 9030 处理器系列产品不支持 BFLOAT16。selfRef在接口文档中支持的数据类型还包含COMPLEX64这是 aclnn 宿主侧AscendCL 接口层放宽后的集合算子定义层fill_diagonal_v2_def.cpp的 kernel 输入仍限定为 FLOAT16/FLOAT/DOUBLE/UINT8/BOOL/INT8/INT16/INT32/BF16/INT64接入时以实际运行的平台为准。wrap为可选属性算子定义中Attr(wrap).AttrType(OPTIONAL).Bool(false)表明其默认值为false即不传时按普通主对角线填充。fillValue在接口层通过aclScalar传递必须能无损转换为selfRef的数据类型见下文约束说明aclnn_fill_diagonal.cpp 中通过CheckCanCastType与CheckNotOverflow两个函数完成这两项校验。约束说明使用 FillDiagonalV2 前必须满足以下约束README 接口文档 源码三重印证selfRef的维度必须大于 1即至少是二维 Tensor。在 aclnn_fill_diagonal.cpp 的CheckShapeValid中通过OP_CHECK_MIN_DIM(selfRef, 2, ...)强制校验维度小于等于 1 会返回ACLNN_ERR_PARAM_INVALID。当selfRef的维度大于 2 时各维度的长度必须相同源码中IsSameDimLength函数逐一比较所有维度长度不相等则校验失败。fillValue必须能转换为 FLOAT 类型并且在转换为selfRef的数据类型时不能发生溢出接口层分别通过CheckCanCastType能否转 FLOAT与CheckNotOverflow转换是否溢出把关溢出场景会明确报错如把 300 填充到 INT8 矩阵会触发溢出检查。非连续 Tensor 支持selfRef支持非连续 Tensor如切片、转置视图。宿主侧在 aclnn_fill_diagonal.cpp 中先通过l0op::Contiguous将输入转为连续 Tensor再执行填充最后用ViewCopy将结果拷贝回原可能非连续的张量。数据规模提示接口文档指出当selfRef总字节数超过 2^31 字节即超过 2GB时会触发算子执行超时超大矩阵需要拆分处理。两段式 aclnn 接口调用说明接口原型FillDiagonalV2 通过aclnnInplaceFillDiagonal两段式接口调用第一段获取 workspace 大小并完成入参校验第二段执行计算完整协议参见 两段式接口说明aclnnStatus aclnnInplaceFillDiagonalGetWorkspaceSize( aclTensor* selfRef, // 输入/输出张量 const aclScalar* fillValue, // 填充值 bool wrap, // 填充方式 uint64_t* workspaceSize, // 输出的 workspace 大小 aclOpExecutor** executor) // 输出的 op 执行器aclnnStatus aclnnInplaceFillDiagonal( void *workspace, // Device 侧申请的 workspace 内存 uint64_t workspaceSize, // 由第一段接口返回 aclOpExecutor *executor, // op 执行器 aclrtStream stream) // 执行 Stream第一段接口 aclnnInplaceFillDiagonalGetWorkspaceSize 参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续 TensorselfRefaclTensor*输入/输出表示需要填充的输入、输出 TensorselfRef 最大维度不能超过 2总字节数超过 2GB 会触发执行超时FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL、COMPLEX64、BFLOAT16ND1、2√fillValueaclScalar*输入表示填充值数据类型需可转换为 FLOAT 且转换到 selfRef 类型时不溢出FLOAT、FLOAT16、DOUBLE、UINT8、INT8、INT16、INT32、INT64、BOOL--√wrapbool输入表示填充方式公式中的 wrap高矩阵场景下为 True 时每 N 行形成新对角线N min(col, row)BOOL-1、2√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----aclOpExecutor**输出返回 op 执行器包含算子计算流程-----第一段接口的入参校验与错误码第一段接口会完成入参校验校验失败时返回aclnnStatus状态码具体错误码定义参见 aclnn 返回码说明返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 fillValue 或 selfRef 是空指针ACLNN_ERR_PARAM_INVALID161002selfRef 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002selfRef 的维度小于等于 1ACLNN_ERR_PARAM_INVALID161002当 selfRef 的维度大于 2 时各维度的长度不相同ACLNN_ERR_PARAM_INVALID161002当 fillValue 不能转换为 FLOAT 时ACLNN_ERR_PARAM_INVALID161002当 fillValue 转换为 selfRef 的数据类型时发生溢出这些校验逻辑在 aclnn_fill_diagonal.cpp 的CheckParams中按「空指针 → 数据类型 → shape → 类型可转换 → 溢出」的顺序依次执行与错误码表格一一对应。第二段接口 aclnnInplaceFillDiagonal 参数参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnInplaceFillDiagonalGetWorkspaceSize 获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream完整性约束与确定性确定性计算aclnnInplaceFillDiagonal默认确定性实现即相同输入在多次执行下结果可复现。完整调用示例可编译运行以下示例代码摘自仓库 examples/test_aclnn_fill_diagonal_v2.cpp完整演示了从环境初始化、张量构造、两段式调用到结果回拷与资源释放的全流程。该示例对 3×3 的 FLOAT 矩阵执行wrapfalse的主对角线填充填充值 11.1执行后主对角线三个元素均变为 11.1。具体编译与运行方式请参考 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_fill_diagonal.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 shape_size 1; for (auto i : shape) { shape_size * i; } return shape_size; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法AscendCL 初始化 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 接口创建 aclTensorND 格式 *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 初始化 int32_t deviceId 0; // 根据自己的实际 device 填写 aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret 0, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入按 API 接口自定义 std::vectorint64_t selfShape {3, 3}; void* selfDeviceAddr nullptr; aclTensor* self nullptr; aclScalar* fillValue nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7, 8}; float value 11.1f; bool wrap false; // 创建 self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建 fillValue aclScalar fillValue aclCreateScalar(value, aclDataType::ACL_FLOAT); CHECK_RET(fillValue ! nullptr, return ret); // 3. 调用 CANN 算子库 API两段式 uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段计算 workspace 大小并完成入参校验 ret aclnnInplaceFillDiagonalGetWorkspaceSize(self, fillValue, wrap, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceFillDiagonalGetWorkspaceSize 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); } // 第二段执行算子计算 ret aclnnInplaceFillDiagonal(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnInplaceFillDiagonal 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 侧并打印 auto size GetShapeSize(selfShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), selfDeviceAddr, size * sizeof(float), 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 和 aclScalar aclDestroyTensor(self); aclDestroyScalar(fillValue); // 7. 释放 device 资源 aclrtFree(selfDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中几个值得留意的实操要点workspaceSize可能为 0此时无需申请 workspace 内存直接传入空指针即可selfRef同时是输入与输出结果直接在原张量上修改回拷时读取的是同一个 device 地址非连续 Tensor 场景下宿主层会通过Contiguous/ViewCopy自动处理内存布局调用方无需手动搬移数据。源码级实现解析算子定义OpDeffill_diagonal_v2_def.cpp 通过OP_ADD(FillDiagonalV2)注册算子关键信息包括输入x、fill_value与输出x均为REQUIRED参数格式限定为 NDwrap属性为OPTIONAL默认值false所有平台的 AICore 配置均开启DynamicCompileStaticFlag(true)、DynamicFormatFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)即算子原生支持动态 shape、动态 rank这也是 README 中「支持非连续 Tensor」与动态场景的基础平台差异体现在fill_value的数据类型集合Atlas A2/A3/950 系列含 BF16Atlas 推理系列与 Kirin 系列不含 BF16。Shape 推导InferShapefill_diagonal_v2_infershape.cpp 的实现非常简单*y_shape *x_shape;即输出 shape 与输入完全一致。这是因为该算子属于原地修改型算子形状不发生变化真正的工作在 tiling 与 kernel 阶段完成。Tiling 策略fill_diagonal_v2_tiling.cpp 是算子性能的关键所在它完成了三件事计算步长与终止位置对二维矩阵step col 1主对角线相邻元素在一维展开下的地址差对更高维等边长 Tensorstep按各维度长度累乘求和得到。end默认等于总长度对于高矩阵且wrapfalse的特殊场景end截断为col * col避免越界写入。多核切分按totalCoreNum来自TilingPrepare4FillDiagonalV2中的GetCoreNumAiv()将总长度均分为blockLength块最后一个核处理lastBlockLength实现多 AIV 核并行。选择稀疏/稠密两条 kernel 路径SetTilingKey4FillDiagonalV2依据数据类型1B/2B/4B/8B 元素宽度映射到不同步长阈值和 wrap 场景决策。当**非 Ascend 950 平台、wrap 为 True、数据量大end 1000000且对角线步长小step 小于按类型设定的阈值**时选择 tiling key1 的稠密Dense路径其余场景走 tiling key0 的稀疏Sparse路径。源码注释明确说明Atlas A2/A3 走 denseAscend 950 因 L2 缓存一致性问题走 sparse。Kernel 实现双路径kernel 入口 fill_diagonal_v2.cpp 根据 tiling key 分派到两个实现fill_diagonal_v2_sparse.hSparse 路径直接在 GM 上按diagIndex step的等差数列逐点写入val不搬数据进 UB。关键细节是它会把读写范围对齐到 128B cache line 边界并在末尾调用DataCacheCleanAndInvalid做缓存一致性处理避免多核写入互相覆盖缓存行。适用于对角线稀疏、直接写 GM 更划算的场景。fill_diagonal_v2_dense.hDense 路径采用经典的「CopyIn → Compute → CopyOut」流水结构BUFFER_NUM 2双缓冲隐藏访存延迟。Compute 阶段将整个 tile 从 GM 搬入 UB 后通过dataLocal(diagIndex - tileStart) val在 UB 内定位对角线元素并原地赋值最后统一写回 GM。适用于对角线条目密集、整块搬运摊销更优的场景。从代码结构看两条路径共享同一套 tiling 数据step、end、blockLength等由编译期 tiling key 在运行时选择兼顾了稀疏写入的低开销与稠密场景的带宽利用。测试验证仓库为该算子提供了完整的 UT 与 ST 覆盖kernel 侧test_fill_diagonal_v2.cpp 使用 gtest 在 CPU 模拟环境跑 kernel并通过 gen_data.py 生成[4,4]、填充值3.14、wraptrue等组合的输入 bin 文件进行对比验证宿主侧test_fill_diagonal_v2_infershape.cpp 与 test_fill_diagonal_v2_tiling.cpp 分别覆盖 shape 推导与 tiling 参数API 侧test_inplace_fill_diagonal.cpp 与 ST 用例 executor_aclnnInplaceFillDiagonal.py 验证两段式接口的端到端行为覆盖了空 tensor、非连续 tensor 等边界场景。小结FillDiagonalV2 是 ops-math 中实现「对角线填充」的标准原地算子README 明确了它在各产品线的支持矩阵、wrap两种填充语义下的精确定位公式、完整参数与约束底层源码则揭示了动态 shape 支持、多核 tiling 切分、稀疏/稠密双 kernel 路径与缓存一致性处理等工程细节。开发者只需按照本文的两段式接口调用示例即可在 Atlas A2/A3/950 推理与训练、Atlas 推理系列及 Kirin X90/9030 平台上完成对角线填充任务若需更深入的接口协议、返回码或编译运行细节可继续查阅 两段式接口说明、aclnn 返回码说明 与 编译与运行样例。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考