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 算子计算的入参公式中的 TFLOAT16、FLOAT、INT16、INT32、BF16NDfill_value输入用于填充的标量值公式中的 vFLOAT16、FLOAT、INT16、INT32、BF16NDx输出执行 fill 算子计算后的出参公式中的 TFLOAT16、FLOAT、INT16、INT32、BF16ND参数定义在源码中与文档一一对应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 与动态 rankPrecisionReduceFlag(true)允许精度归约优化同时 fill_v2_binary.json 中每个输入/输出的shape均为[-2]代表任意 shape 的动态 shape 描述符因此 FillV2 可处理任意形状的 ND 张量测试样例中的{8, 1023}只是其中一例。调用说明调用方式调用样例说明aclnn 调用test_aclnn_fill_v2.cpp通过 aclnn 两段式接口aclnnFillV2GetWorkspaceSizeaclnnFillV2调用 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 - aclrtCreateStream2. 构造输入与输出 aclTensor以 INT16 为例创建 shape 为{8, 1023}的输入selfRef与 shape 为{1}的填充值valuestd::vectorint64_t selfRefShape {8, 1023}; std::vectorint64_t valueShape {1}; std::vectorint16_t selfRefHostData(8184, 0); std::vectorint16_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. 调用两段式算子 APIuint64_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将结果aclrtMemcpyDEVICE_TO_HOST拷回 host 并打印。6. 释放 aclTensoraclDestroyTensor(selfRef); aclDestroyTensor(value);7. 释放 device 资源依次aclrtFree输入、输出、workspace、aclrtDestroyStream、aclrtResetDevice、aclFinalize。由于aclCreateTensor要求输入按 ND 连续布局样例中手动计算连续 stridesstrides[i] shape[i1] * strides[i1]这也是所有 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.cppInferShapeFillV2直接将输入 shape 拷贝给输出*yShape *xShape与填充不改变形状的语义一致数据类型推导位于 fill_v2_infer.cppInferDataTypeFillV2将输出数据类型设置为与输入一致。这两个推导与 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::FillV2DTYPE_X并执行Init与Process。核心计算在 fill_v2.h 中完成InitL64-L91从fill_value的 GlobalTensor 中取出标量值fillValue_valGm.GetValue(0)根据当前核 IDAscendC::GetBlockIdx()与 tiling 下发的 big-core/small-core 划分计算每个核负责的数据量并设置 GlobalTensor 起始偏移pipe-InitBuffer(inputQueueX, BUFFER_NUM, tileDataNum * sizeof(T))申请双缓冲区BUFFER_NUM 2。ComputeL100-L106调用向量指令AscendC::DuplicateT(xLocal, this-fillValue_, this-processDataNum)将processDataNum个元素全部写为填充值——这是整个算子的核心计算指令由硬件向量单元高效执行。ProcessL116-L130按 tile 数循环最后一个 tile 使用tailDataNum处理尾部数据每个 tile 依次执行CopyIn - Compute - CopyOut。由于 FillV2 不需要搬运输入数据直接在本地缓冲填充即可CopyIn仅负责AllocTensorEnQueCopyOut通过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 时报错返回。校验 dtypeGetShapeAttrsInfo从输入 shape 计算总元素数GetStorageShape().GetShapeSize()并对{DT_FLOAT16, DT_FLOAT, DT_INT16, DT_INT32, DT_BF16}集合外类型直接拒绝。申请 workspaceGetWorkspaceSize按WS_SYS_SIZE16MB 系统库 API workspace申请对应 aclnn 样例中第二段接口前aclrtMalloc的 workspace。切分计算以BLOCK_SIZE 32字节为基本块将输入按inputLengthBytes换算为总块数blocksTotal并以 AIV 核数均分核数超过块数时收敛为块数前tailBlockNum个核为 big-core各多分 1 个块其余为 small-core实现负载均衡每个核内部再按 UB 容量ubSize / BLOCK_SIZE / BUFFER_NUM切成多个 tiletileDataNum保证每个 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_hostshape 推导、tiling、workspace与 op_kernelAscendC 双缓冲流水 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），仅供参考