FEATURED · 精选文章

CANN ops-math 算子解析:DenseBincount 稠密分箱计数算子实现与 GEIR 图模式调用指南

发布时间 / 2026/9/20 13:11:28
来源 / 创域科博编辑部
栏目 / 资讯中心
CANN ops-math 算子解析:DenseBincount 稠密分箱计数算子实现与 GEIR 图模式调用指南 算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载本文以 CANN ops-math 开源仓库中的math/dense_bincount算子为对象系统讲解 DenseBincount 在 NPU 上实现按 bin 下标统计出现次数或权重和的数学语义、参数契约与约束边界并结合源码给出 GEIR 图模式的完整调用示例与底层 tiling/内核实现原理。读完本文你将掌握该算子的计算公式、参数配置、产品支持矩阵、调用方式以及从算子注册、shape 推导到调度切分的完整实现链路。产品支持情况DenseBincount 算子在不同硬件产品上的支持情况如下取自 README产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3训练系列产品/Atlas A3推理系列产品√Atlas A2训练系列产品/Atlas A2推理系列产品√Atlas 200I/500 A2推理产品×Atlas推理系列产品×Atlas训练系列产品√从实现结构看该算子的 host 侧 tiling 与 kernel 侧实现均位于arch35目录dense_bincount_tiling.cpp、dense_bincount.cpp并且 算子定义 中通过OpAICoreConfig显式添加了ascend950配置这与文档中Ascend 950 系列 √的支持情况相互印证。功能说明算子功能DenseBincount 对一维或二维整数输入逐行统计每个 bin 的出现次数或权重和。计算公式设 bin 数量为M size[0]。对于一维输入输出形状为[M]对于二维输入输出形状为[N, M]其中N input.shape[0]。对有效的行号r和 bin 下标m非二值输出为$$ output[r,m] \sum_{i \in row(r)} \mathbf{1}(input_im) \times \begin{cases} weights_i, numel(weights)0 \ 1, numel(weights)0 \end{cases} $$binary_output为true时输出为$$ output[r,m] \mathbf{1}(\exists i \in row(r), input_im) $$边界语义输入值大于或等于M时忽略一维输入中的负值忽略二维输入中的负值按 TensorFlow DenseBincount 语义折返到前一行折返后行号或 bin 下标仍越界时忽略binary_output为true时不使用weights的数值。算子 IR 注释dense_bincount_proto.h也明确标注了Third-party framework compatibility: TensorFlow DenseBincount说明该算子在语义上对齐 TensorFlow 的 DenseBincount 行为。参数说明参数名输入/输出/属性是否必选描述数据类型数据格式input输入是一维或二维 bin 下标张量。INT32、INT64NDsize输入是包含一个非负常量值的一维张量size[0]表示 bin 数量。数据类型必须与input一致。INT32、INT64NDweights输入是权重张量。该输入不可省略无权重时传零元素张量表示按次数统计。非空时元素数量必须与input相同并按展平下标一一对应。FLOAT32NDbinary_output属性否是否仅输出 bin 命中标记默认为false。Bool-output输出是统计结果。一维输入对应形状[size[0]]二维输入对应形状[input.shape[0], size[0]]。FLOAT32ND与源码的对应关系IR 层注册dense_bincount_proto.h 中通过REG_OP(DenseBincount)声明三个输入与一个输出其中weights在 IR 层放开到{DT_INT32, DT_INT64, DT_FLOAT, DT_DOUBLE}、binary_output属性默认值为false算子定义层收敛dense_bincount_def.cpp 将weights与output收敛为ge::DT_FLOATFLOAT32、格式全部限定为FORMAT_ND并通过.AutoContiguous()要求输入连续排布同时设置了DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)等动态编译标志输出类型固定dense_bincount_graph_infer.cpp 的InferDataTypeDenseBincount无条件将输出数据类型置为ge::DT_FLOAT与文档中 output 的 FLOAT32 约束一致size 的常量依赖size在 def 中声明为ValueDepend(OPTIONAL)infershape 通过InputsDataDependency({1})声明对输入 1size的数据依赖tiling 阶段通过Ops::Base::GetConstInt读取其常量值。约束说明input的 rank 只能为 1 或 2size必须是仅含一个非负常量值的一维 ND 张量且数据类型必须与input相同weights不可省略。无权重时传入零元素张量非空时元素数量必须与input相同Ascend 950PR/Ascend 950DTweights支持任意零元素 shape非空时为 FLOAT32 类型 ND 张量支持input任意维度大小为 0 或size[0]为 0。input的 shape 为[0]时输出 shape 为[size[0]]shape 为[N, 0]时输出 shape 为[N, size[0]]shape 为[0, N]或[0, 0]时输出 shape 为[0, size[0]]size[0]为 0 时输出最后一维大小为 0输出元素数rows * size[0]不能超过 INT64 索引可表示的 FLOAT32 元素数量上限且输出张量所需存储必须能由运行环境分配一维输入的rows取 1Atlas A3 训练系列产品/Atlas A3 推理系列产品、Atlas A2 训练系列产品/Atlas A2 推理系列产品沿用既有实现的空 Tensor 处理边界weights的 rank 不能超过 2size[0]、输入维度大小和输入元素数量必须在 INT32 范围内本算子仅支持 GE 图模式和 TensorFlow 图解析通路调用不提供公开的 aclnn 接口。源码中的约束落实rank 校验infershape 实现 中input-GetDimNum() 1 || input-GetDimNum() 2直接判定非法未知 rankIsUnknownRank时先透传输出再返回非负常量校验infershape 对size读取失败时输出形状相应维度置-1待后续推导size 0时直接报错size must be a non-negative constant溢出保护infershape 与 tiling 中都定义了MAX_OUTPUT_ELEMENTS std::numeric_limitsint64_t::max() / sizeof(float)当rows * size MAX_OUTPUT_ELEMENTS / rows时拒绝构图/编译weights 元素数校验tiling 的SetDenseBincountTilingData中要求weightsNum 0零元素张量或weightsNum numValues否则报错dtype/rank/shape 契约tiling 的ValidateDenseBincountInputs统一校验 rank ∈ [1,2]、size为一维且元素数为 1、size与input数据类型一致、weights为 FLOAT32上述校验逻辑均有对应的 host 侧单元测试覆盖见 test_dense_bincount_infershape.cpp其中包含一维未知 shape、常量size1的二维输入、未知 rank-2、空输入[0]等用例。调用说明调用方式调用样例说明图模式调用test_geir_dense_bincount.cpp通过算子IR构图调用 DenseBincount 算子。GEIR 图模式调用示例解析test_geir_dense_bincount.cpp 给出了完整的图模式调用流程核心步骤包括初始化 GE通过ge::GEInitialize设置全局选项如{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}构图创建op::DenseBincount(dense_bincount1)示例中inputshape{4}、DT_INT32占位 Data 节点set_attr_index(0)sizeshape{1}、DT_INT32的常量节点值1即M size[0] 1weightsshape{4}、DT_FLOAT的输入节点outputshape{1}、DT_FLOAT建会话并运行session-AddGraph(graph_id, graph, graph_options)后session-RunGraph(graph_id, input, output)完成执行结果校验示例对输入 4 个全 0 的 bin 下标、单位权重做验证期望输出元素数为 1 且值为4.0f4 个元素全部落入 bin 0权重各为 1非二值模式求和为 4并通过WriteDataToFile将输入输出落盘为.bin文件用于进一步比对。该示例同时体现了size作为常量输入通过ADD_CONST_INPUT宏构造、而input/weights作为动态输入通过ADD_INPUT宏构造的典型 GEIR 构图模式。TensorFlow 图解析通路framework/dense_bincount_tf_plugin.cpp 通过REGISTER_CUSTOM_OP(DenseBincount)注册 TensorFlow 自定义算子映射.FrameworkType(TENSORFLOW)、.OriginOpType(DenseBincount)、参数解析采用AutoMappingByOpFnImplyType为TVM。这解释了文档仅支持 GE 图模式和 TensorFlow 图解析通路调用的底层实现依据——用户既可以在 GEIR 图中直接构造算子也可以通过 TensorFlow 图经图解析通路降级到该算子执行。底层调度与内核实现原理Tiling 调度策略tiling 实现 负责在构图阶段确定运行参数主要决策点包括tiling 数据结构dense_bincount_tiling_data.h 定义了numValues输入元素数、inputRows、inputCols、size、usedCoreNum实际使用核数、privateHistElems私有直方图元素数等字段核数估算GetRequiredCoreNum依据工作量与PER_CORE_MIN_ELEMENTS 1024估算并行核数保证每核负载不低于阈值同时以平台GetCoreNumAiv()返回的 AIV 核数为上限UB 内存预算从 UB 总大小中扣除DCACHE_SIZE 128KB的 SIMT DCache 预留后得到可用的 local memoryFitsInLocalMemory按 32 字节对齐检查输出直方图elements * sizeof(float)能否装入私有直方图决策GetPrivateHistElems在输出直方图或二值模式下行级直方图可装入 local memory 且写回开销可接受privateWriteElements / PRIVATE_WRITEBACK_FACTOR numValues时为每核分配私有直方图以减少写回竞争useRowPrivate在二维 二值 多行 行数不超过核数时启用行级私有直方图TilingKey 编码mode (is1D 2) | (binary 1) | hasWeights将输入维度1D/2D、是否二值输出、是否带权重三个编译期分支编码进 tiling key实现不同场景的内核实例化。内核入口与分支解耦kernel 实现 是典型的 AscendC 内核模板constexpr bool IS_1D schMode 4; constexpr bool BINARY_OUTPUT (schMode 0x2) ! 0; constexpr bool HAS_WEIGHTS (schMode 0x1) ! 0; NsDenseBincount::DenseBincountRegbaseDTYPE_INPUT, DTYPE_SIZE, IS_1D, BINARY_OUTPUT, HAS_WEIGHTS op(...);即通过编译期模板参数把一维/二维、二值/加权、有权重/无权重组合展开为 8 种实例避免运行时分支kernel 侧仅需执行InitProcess完成统计写入。总结DenseBincount 是 CANN ops-math 中语义上对齐 TensorFlow DenseBincount 的稠密分箱计数算子支持 1D/2D 整数输入、可选权重加权与二值输出标记仅通过 GEIR 图模式与 TensorFlow 图解析通路调用。从实现链路看其约束rank、size 常量、weights 数量、溢出上限在 infershape 与 tiling 阶段双重校验调度侧通过核数估算 UB 预算 私有直方图进行优化内核侧以 tiling key 驱动的模板分支保证执行效率。相关源码与测试可分别在 math/dense_bincount 目录下进一步查阅。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 算子 Identity在 NPU 上实现张量透传的 GEIR 图模式调用指南CANN ops math 算子 Identity在 NPU 上实现张量透传的 GEIR 图模式调用指南 Identity 是 CANN ops math 算算子库人工智能CANNCANN ops-math CaseCondition 算子详解LU 分解 case 编号判定与 GEIR 图模式调用实战CANN ops math CaseCondition 算子详解LU 分解 case 编号判定与 GEIR 图模式调用实战 CaseCondition 是 C算子库人工智能CANNTradingAgents-CN Dataflows 模块全面优化实践从激进重构到务实演进的目录治理、职责分层与向后兼容重构指南TradingAgents CN Dataflows 模块全面优化实践从激进重构到务实演进的目录治理、职责分层与向后兼容重构指南 本篇技术指南围绕 Tradi算子库人工智能CANN上一篇LeetCode-Solutions-in-Good-Style字符串算法完整指南10个高效学习技巧下一篇清华大学uCore OS Labs从零构建操作系统的终极指南创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻