CANN ops-math FillV2 算子深度解析：NPU 上全量填充的实现与 aclnn 调用实战-开发者社区

CANN ops-math FillV2 算子深度解析：NPU 上全量填充的实现与 aclnn 调用实战

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

导读

本文围绕 CANN 开源数学算子库 ops-math 中的 FillV2 算子（README）展开，讲解其在 Atlas A2 训练/推理系列产品上实现张量全量填充的功能定义、参数约束、调用方式，并深入 op_host、op_kernel 与 examples 源码，剖析算子注册、shape/数据类型推导、AscendC Kernel 双缓冲流水与多核 tiling 切分的底层实现。读完本文，你将掌握 FillV2 的完整使用方式，并理解一个 AscendC 算子从定义到落盘执行的完整链路。

产品支持情况

产品	是否支持
Atlas A2 训练系列产品 / Atlas A2 推理系列产品	√

该支持关系同时体现在算子定义中：在 fill_v2_def.cpp 中，FillV2 通过this->AICore().AddConfig("ascend910b", aicoreConfig)将算子的 AI Core 配置绑定到ascend910b平台，并在 op_host/ascend910b 目录下提供对应的算子二进制描述文件。

功能说明

FillV2 实现张量的填充功能：将张量中的所有元素填充为指定的统一标量值，属于原地（inplace）操作——输出与输入指向同一块张量。

算子功能：将输入张量 T 的所有元素填充为统一值 v。
计算公式：对于任意索引位置(i, j, ...)，执行fill(T, v)后：

$$\forall (i, j, \dots), \quad T[i, j, \dots] = v$$

从图规约定义看，该算子与 PyTorch 的Tensor.fill_语义对齐。fill_v2_proto.h 中通过REG_OP(FillV2)注册算子原型，并明确注释 "Compatible with the Pytorch operator Fill"，输出x与输入x指向同一张量。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	待进行 fill 算子计算的入参，公式中的 T	FLOAT16、FLOAT、INT16、INT32、BF16	ND
fill_value	输入	用于填充的标量值，公式中的 v	FLOAT16、FLOAT、INT16、INT32、BF16	ND
x	输出	执行 fill 算子计算后的出参，公式中的 T'	FLOAT16、FLOAT、INT16、INT32、BF16	ND

参数定义在源码中与文档一一对应：

fill_v2_def.cpp 中Input("x")、Input("fill_value")、Output("x")三个参数均为REQUIRED，数据类型集合为{ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT16, ge::DT_INT32, ge::DT_BF16}，格式统一为ge::FORMAT_ND，且支持UnknownShapeFormat（未知 shape 场景下仍为 ND）。
fill_value虽然在语义上是标量，但在接口层面以 Tensor 形式传入（示例中 shape 为{1}，见 test_aclnn_fill_v2.cpp）。
数据类型推导保证一致性：fill_value与x必须同类型，输出数据类型等于输入数据类型（见下文"推导逻辑"小节）。

约束说明

无额外约束（文档明确"约束说明：无"）。
从实现细节看，fill_v2_def.cpp 中的DynamicCompileStaticFlag(true)、DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)表明该算子支持动态 shape 与动态 rank；PrecisionReduceFlag(true)允许精度归约优化；同时 fill_v2_binary.json 中每个输入/输出的shape均为[-2]（代表任意 shape 的动态 shape 描述符），因此 FillV2 可处理任意形状的 ND 张量（测试样例中的{8, 1023}只是其中一例）。

调用说明

调用方式	调用样例	说明
aclnn 调用	test_aclnn_fill_v2.cpp	通过 aclnn 两段式接口（`aclnnFillV2GetWorkspaceSize`+`aclnnFillV2`）调用 FillV2 算子
GEIR 调用	test_geir_fill_v2.cpp	通过 Graph Engine IR 构建图并下发的调用方式

说明：原始 README 中"通过 [test_aclnn_s_where] 接口方式调用"的表述属于文档笔误，实际接口名为aclnnFillV2（对应头文件aclnnop/aclnn_fill_v2.h），下文按源码实际接口展开。

aclnn 两段式调用流程

test_aclnn_fill_v2.cpp 给出了一个可直接运行的完整样例，整体流程分为 7 步：

1. device / stream 初始化（固定写法）

int32_t deviceId = 0; aclrtStream stream; ret = Init(deviceId, &stream); // 内部依次执行 aclInit -> aclrtSetDevice -> aclrtCreateStream

2. 构造输入与输出 aclTensor

以 INT16 为例，创建 shape 为{8, 1023}的输入selfRef与 shape 为{1}的填充值value：

std::vector<int64_t> selfRefShape = {8, 1023}; std::vector<int64_t> valueShape = {1}; std::vector<int16_t> selfRefHostData(8184, 0); std::vector<int16_t> valueHostData = {1}; // 内部依次调用 aclrtMalloc 申请 device 内存、aclrtMemcpy 拷贝 host->device、 // 计算连续 strides、aclCreateTensor 创建 aclTensor CreateAclTensor(selfRefHostData, selfRefShape, &selfRefDeviceAddr, ACL_INT16, &selfRef); CreateAclTensor(valueHostData, valueShape, &valueDeviceAddr, ACL_INT16, &value);

3. 调用两段式算子 API

uint64_t workspaceSize = 0; aclOpExecutor* executor; // 第一段：获取 workspace 大小，完成算子准备工作 ret = aclnnFillV2GetWorkspaceSize(selfRef, value, &workspaceSize, &executor); // 根据计算出的 workspaceSize 申请 device 内存（可能为 0，需判断） void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段：执行算子 ret = aclnnFillV2(workspaceAddr, workspaceSize, executor, stream);

4. 同步等待任务执行结束

ret = aclrtSynchronizeStream(stream);

5. 获取输出值：由于 FillV2 是原地操作，直接从selfRefDeviceAddr将结果aclrtMemcpy（DEVICE_TO_HOST）拷回 host 并打印。

6. 释放 aclTensor：aclDestroyTensor(selfRef); aclDestroyTensor(value);

7. 释放 device 资源：依次aclrtFree（输入、输出、workspace）、aclrtDestroyStream、aclrtResetDevice、aclFinalize。

由于aclCreateTensor要求输入按 ND 连续布局，样例中手动计算连续 strides（strides[i] = shape[i+1] * strides[i+1]），这也是所有 aclnn 调用构造aclTensor的通用前提。

GEIR 图下发调用

test_geir_fill_v2.cpp 展示了另一条调用路径：通过ge::op::FillV2（由 fill_v2_proto.h 生成的算子类）构建计算图，配合op::Data占位输入与op::Const常量输入生成输入数据，再经 GE 接口完成构图与下发。该方式适用于需要将 FillV2 嵌入更大计算图（如经 GEIR 序列化、优化与整图下发）的场景。

源码级实现剖析

算子原型与推导逻辑

FillV2 的 shape 推导位于 fill_v2_infershape.cpp：InferShapeFillV2直接将输入 shape 拷贝给输出（*yShape = *xShape），与"填充不改变形状"的语义一致；数据类型推导位于 fill_v2_infer.cpp：InferDataTypeFillV2将输出数据类型设置为与输入一致。这两个推导与 fill_v2_def.cpp 中OpDef声明的五类数据类型（FLOAT16/FLOAT/INT16/INT32/BF16）共同构成算子在 GE 侧的完整描述。

Kernel 实现：AscendC 双缓冲 + Duplicate 填充

Kernel 入口 fill_v2.cpp 中通过REGISTER_TILING_DEFAULT与GET_TILING_DATA_WITH_STRUCT读取 host 侧下发的 tiling 参数，随后实例化NsFillV2::FillV2<DTYPE_X>并执行Init与Process。

核心计算在 fill_v2.h 中完成：

Init（L64-L91）：从fill_value的 GlobalTensor 中取出标量值fillValue_（valGm.GetValue(0)）；根据当前核 ID（AscendC::GetBlockIdx()）与 tiling 下发的 big-core/small-core 划分，计算每个核负责的数据量并设置 GlobalTensor 起始偏移；pipe->InitBuffer(inputQueueX, BUFFER_NUM, tileDataNum * sizeof(T))申请双缓冲区（BUFFER_NUM = 2）。
Compute（L100-L106）：调用向量指令AscendC::Duplicate<T>(xLocal, this->fillValue_, this->processDataNum)，将processDataNum个元素全部写为填充值——这是整个算子的核心计算指令，由硬件向量单元高效执行。
Process（L116-L130）：按 tile 数循环，最后一个 tile 使用tailDataNum处理尾部数据，每个 tile 依次执行CopyIn -> Compute -> CopyOut。由于 FillV2 不需要搬运输入数据（直接在本地缓冲填充即可），CopyIn仅负责AllocTensor+EnQue，CopyOut通过AscendC::DataCopy将填充结果写回 Global Memory 中该核负责的偏移处（inputGMX[progress * tileDataNum]）。双缓冲配合 queue 机制实现取数与回写的流水重叠。

Tiling 切分：按 AIV 核数与 UB 容量均衡

Tiling 逻辑在 fill_v2_tiling.cpp 中：

获取平台信息：GetPlatformInfo通过PlatformAscendC读取 AIV 核数（GetCoreNumAiv）与 Unified Buffer 容量（GetCoreMemSize(UB)），核数为 0 或 UB 为 0 时报错返回。
校验 dtype：GetShapeAttrsInfo从输入 shape 计算总元素数（GetStorageShape().GetShapeSize()），并对{DT_FLOAT16, DT_FLOAT, DT_INT16, DT_INT32, DT_BF16}集合外类型直接拒绝。
申请 workspace：GetWorkspaceSize按WS_SYS_SIZE（16MB）+ 系统库 API workspace申请，对应 aclnn 样例中第二段接口前aclrtMalloc的 workspace。
切分计算：
- 以BLOCK_SIZE = 32字节为基本块，将输入按inputLengthBytes换算为总块数blocksTotal，并以 AIV 核数均分（核数超过块数时收敛为块数）；
- 前tailBlockNum个核为 big-core（各多分 1 个块），其余为 small-core，实现负载均衡；
- 每个核内部再按 UB 容量（ubSize / BLOCK_SIZE / BUFFER_NUM）切成多个 tile，tileDataNum保证每个 tile 至少 1 个元素，尾部不足整块的数据由tailDataNum承接。
写回 tiling 数据：将smallCoreDataNum、bigCoreDataNum、finalBigTileNum、finalSmallTileNum、tileDataNum、smallTailDataNum、bigTailDataNum、tailBlockNum写入 fill_v2_tiling_data.h 定义的FillV2TilingData结构体，并通过context->SetBlockDim(finalCoreNum)下发核数。

该 tiling 结构由 host 与 kernel 共享头文件保持一致，体现了"host 算切分、kernel 照切分执行"的 AscendC 标准开发范式。

算子二进制与平台适配

fill_v2_binary.json 为五种数据类型（float32、float16、bfloat16、int16、int32）各生成一个二进制条目（FillV2Float32、FillV2Float16、FillV2Bfloat16、FillV2Int16、FillV2Int32），每个条目声明输入输出均为 ND 格式、paramType: required、动态 shape（[-2]）。配合同目录下的fill_v2_simplified_key.ini，构成算子在下发时按 dtype 选择对应 kernel 二进制的完整适配描述。

小结

FillV2 是 ops-math 中一个实现简洁但工程链路完整的典型算子：文档层面明确了 Atlas A2 系列产品的支持范围、五类数据类型与 ND 格式约束；代码层面由 op_graph（原型与数据类型推导）、op_host（shape 推导、tiling、workspace）与 op_kernel（AscendC 双缓冲流水 + Duplicate 指令）三部分协同完成。若需在自有工程中快速接入，可直接参考 test_aclnn_fill_v2.cpp 的两段式调用模板；若需理解其性能设计，可从 fill_v2_tiling.cpp 的 big-core/small-core 均衡切分与BUFFER_NUM = 2的双缓冲机制入手。

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

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考