CANN ops-math FillDiagonalV2 算子深度解析：对角线填充的算法原理与 aclnnInplaceFillDiagonal 调用实践

CANN ops-math FillDiagonalV2 算子深度解析：对角线填充的算法原理与 aclnnInplaceFillDiagonal 调用实践

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

FillDiagonalV2 是 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 配置；为ascend310p（Atlas 推理系列）、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 < m，0 <= 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	输入/输出张量	表示输入/输出张量，支持非连续的 Tensor	BFLOAT16、FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL	ND
fillValue	输入属性	表示填充值，数据类型需要是可转换为 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 + 接口文档 + 源码三重印证）：

1. selfRef的维度必须大于 1：即至少是二维 Tensor。在 aclnn_fill_diagonal.cpp 的CheckShapeValid中通过OP_CHECK_MIN_DIM(selfRef, 2, ...)强制校验，维度小于等于 1 会返回ACLNN_ERR_PARAM_INVALID。
2. 当selfRef的维度大于 2 时，各维度的长度必须相同：源码中IsSameDimLength函数逐一比较所有维度长度，不相等则校验失败。
3. fillValue必须能转换为 FLOAT 类型，并且在转换为selfRef的数据类型时不能发生溢出：接口层分别通过CheckCanCastType（能否转 FLOAT）与CheckNotOverflow（转换是否溢出）把关，溢出场景会明确报错（如把 300 填充到 INT8 矩阵会触发溢出检查）。
4. 非连续 Tensor 支持：selfRef支持非连续 Tensor（如切片、转置视图）。宿主侧在 aclnn_fill_diagonal.cpp 中先通过l0op::Contiguous将输入转为连续 Tensor，再执行填充，最后用ViewCopy将结果拷贝回原（可能非连续的）张量。
5. 数据规模提示：接口文档指出，当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)	非连续 Tensor
selfRef（aclTensor*）	输入/输出	表示需要填充的输入、输出 Tensor	selfRef 最大维度不能超过 2，总字节数超过 2GB 会触发执行超时	FLOAT、FLOAT16、DOUBLE、INT32、INT64、INT16、INT8、UINT8、BOOL、COMPLEX64、BFLOAT16	ND	1、2	√
fillValue（aclScalar*）	输入	表示填充值	数据类型需可转换为 FLOAT 且转换到 selfRef 类型时不溢出	FLOAT、FLOAT16、DOUBLE、UINT8、INT8、INT16、INT32、INT64、BOOL	-	-	√
wrap（bool）	输入	表示填充方式，公式中的 wrap	高矩阵场景下为 True 时每 N 行形成新对角线，N = min(col, row)	BOOL	-	1、2	√
workspaceSize（uint64_t*）	输出	返回需要在 Device 侧申请的 workspace 大小	-	-	-	-	-
aclOpExecutor**	输出	返回 op 执行器，包含算子计算流程	-	-	-	-	-

第一段接口的入参校验与错误码

第一段接口会完成入参校验，校验失败时返回aclnnStatus状态码，具体错误码定义参见 aclnn 返回码说明：

返回值	错误码	描述
ACLNN_ERR_PARAM_NULLPTR	161001	传入的 fillValue 或 selfRef 是空指针
ACLNN_ERR_PARAM_INVALID	161002	selfRef 的数据类型不在支持范围之内
ACLNN_ERR_PARAM_INVALID	161002	selfRef 的维度小于等于 1
ACLNN_ERR_PARAM_INVALID	161002	当 selfRef 的维度大于 2 时，各维度的长度不相同
ACLNN_ERR_PARAM_INVALID	161002	当 fillValue 不能转换为 FLOAT 时
ACLNN_ERR_PARAM_INVALID	161002	当 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 矩阵执行wrap=false的主对角线填充（填充值 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::vector<int64_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::vector<T>& hostData, const std::vector<int64_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::vector<int64_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（ND 格式） *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::vector<int64_t> selfShape = {3, 3}; void* selfDeviceAddr = nullptr; aclTensor* self = nullptr; aclScalar* fillValue = nullptr; std::vector<float> 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::vector<float> 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自动处理内存布局，调用方无需手动搬移数据。

源码级实现解析

算子定义（OpDef）

fill_diagonal_v2_def.cpp 通过OP_ADD(FillDiagonalV2)注册算子，关键信息包括：

• 输入x、fill_value与输出x均为REQUIRED参数，格式限定为 ND；
• wrap属性为OPTIONAL，默认值false；
• 所有平台的 AICore 配置均开启DynamicCompileStaticFlag(true)、DynamicFormatFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)，即算子原生支持动态 shape、动态 rank，这也是 README 中「支持非连续 Tensor」与动态场景的基础；
• 平台差异体现在fill_value的数据类型集合：Atlas A2/A3/950 系列含 BF16，Atlas 推理系列与 Kirin 系列不含 BF16。

Shape 推导（InferShape）

fill_diagonal_v2_infershape.cpp 的实现非常简单：*y_shape = *x_shape;，即输出 shape 与输入完全一致。这是因为该算子属于原地修改型算子，形状不发生变化，真正的工作在 tiling 与 kernel 阶段完成。

Tiling 策略

fill_diagonal_v2_tiling.cpp 是算子性能的关键所在，它完成了三件事：

1. 计算步长与终止位置：对二维矩阵，step = col + 1（主对角线相邻元素在一维展开下的地址差）；对更高维等边长 Tensor，step按各维度长度累乘求和得到。end默认等于总长度；对于高矩阵且wrap=false的特殊场景，end截断为col * col，避免越界写入。

2. 多核切分：按totalCoreNum（来自TilingPrepare4FillDiagonalV2中的GetCoreNumAiv()）将总长度均分为blockLength块，最后一个核处理lastBlockLength，实现多 AIV 核并行。

3. 选择稀疏/稠密两条 kernel 路径：SetTilingKey4FillDiagonalV2依据数据类型（1B/2B/4B/8B 元素宽度映射到不同步长阈值）和 wrap 场景决策。当**非 Ascend 950 平台、wrap 为 True、数据量大（end > 1000000）且对角线步长小（step 小于按类型设定的阈值）**时，选择 tiling key=1 的稠密（Dense）路径；其余场景走 tiling key=0 的稀疏（Sparse）路径。源码注释明确说明：Atlas A2/A3 走 dense，Ascend 950 因 L2 缓存一致性问题走 sparse。

Kernel 实现（双路径）

kernel 入口 fill_diagonal_v2.cpp 根据 tiling key 分派到两个实现：

• fill_diagonal_v2_sparse.h（Sparse 路径）：直接在 GM 上按diagIndex += step的等差数列逐点写入val，不搬数据进 UB。关键细节是它会把读写范围对齐到 128B cache line 边界，并在末尾调用DataCacheCleanAndInvalid做缓存一致性处理，避免多核写入互相覆盖缓存行。适用于对角线稀疏、直接写 GM 更划算的场景。

• fill_diagonal_v2_dense.h（Dense 路径）：采用经典的「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、wrap=true等组合的输入 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