CANN ops-math HistogramFixedWidth 算子详解从等宽直方图计算到 NPU 源码实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 开源仓库中的math/histogram_fixed_width算子为对象完整解析 HistogramFixedWidth 算子的功能定义、产品支持范围、参数与约束、图模式调用方式并结合算子原型proto、框架适配TF 插件、形状推导InferShape、Tiling 策略与 Kernel 实现等源码深入剖析其在 NPU 上的底层工作原理。读完本文读者既能独立完成该算子的图模式调用与结果验证也能理解其分核并行 原子累加的实现思路为后续阅读或移植同类统计类算子打下基础。一、功能说明与计算公式HistogramFixedWidth 算子用于计算输入张量 x 的直方图与 TensorFlow 框架中的histogram_fixed_width算子语义兼容参见 算子原型 中 Compatible with the TensorFlow operator HistogramFixedWidth 的注释说明。其核心计算过程为以range [min, max]作为统计上下限在 min 与 max 之间划出nbins 个等宽区间bin逐一统计输入张量 x 中每个元素落入各区间的个数最终输出形状为[nbins]的直方图结果 y。两个边界规则值得特别注意小于 min 的元素会被统计到第一个区间bin 0大于 max 的元素会被统计到最后一个区间bin nbins-1。区间划分可用如下公式描述设区间宽度width (max - min) / nbinsbin_index clamp((x_i - min) / width, 0, nbins - 1)其中clamp将计算结果截断到[0, nbins-1]范围内。这一截断行为同时解决了元素小于 min与元素大于 max两种越界情况与 TensorFlow 的语义保持一致。二、产品支持情况根据 算子 READMEHistogramFixedWidth 算子在不同产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品×从算子侧代码看该支持范围与 算子定义 中的 AICore 配置一致this-AICore().AddConfig(ascend950, aicoreConfig)与this-AICore().AddConfig(ascend350, aicoreConfig)分别对应 Ascend 950 与 A3/A2 系列arch35产品且 ascend950 与 ascend350 目录下均提供了对应的二进制配置文件。尚未为 Atlas 训练系列产品提供配置这与其不支持的状态相互印证。三、参数说明参数名输入/输出/属性描述数据类型数据格式x输入输入张量待统计的数据FLOAT、FLOAT16、INT32、INT64NDrange输入形状为 [2] 的张量包含 [min, max]FLOAT、FLOAT16、INT32、INT64NDnbins输入标量张量直方图区间数量INT32NDy输出直方图结果形状为 [nbins]INT32NDdtype属性预留属性默认值为 3INT64-各参数在源码中均有对应定义可作为参数说明的补充证据输入/输出定义算子定义 中通过Input(x)、Input(range)、Input(nbins)、Output(y)声明了全部参数数据类型集合与上表一致且range与nbins均标记了ValueDepend(OPTIONAL)——意味着编译期需要读取这两个输入的数值而非仅形状这一点在形状推导与 Tiling 阶段有实际体现。dtype 属性Attr(dtype).AttrType(OPTIONAL).Int(3)默认值为 3对应ge::DT_INT32输出 y 的数据类型。该属性目前属于预留属性在 形状推导实现 中会校验其取值必须等于DT_INT32否则报错返回。ND 格式与动态能力所有输入输出均为 ND 格式且算子定义中开启了DynamicCompileStaticFlag、DynamicFormatFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag四项动态能力标志说明该算子支持动态 shape/rank 场景。四、约束说明算子约束如下输入张量 range 的 max 必须大于 min否则报错。该约束在 Tiling 阶段被显式校验在 tiling 实现 的ValidateRange()中先读取 range 的实际数值minVal、maxValReadRangeMinMax()按 x 的四种数据类型分别解析随后检查minVal maxVal时即返回GRAPH_FAILED并输出 The value of max must be greater than min 的报错日志。除此之外源码还揭示了几条 README 之外的隐含约束供调用时注意range 的形状必须为 [2]在 形状推导 中要求 range 的 shape size 等于 2或为动态形状 -1在ValidateRange()中同样检查rangeLength ! HFW_RANGE_LENGTH (2)时报错。range 的数据类型必须与 x 相同ValidateRange()会对比 range 与 x 的 dtype不一致时报 The dtype of range must be the same as ... of x。nbins 必须大于 0形状推导阶段校验nbins 0时报错同时输出 y 的形状大小必须与 nbins 相等。dtype 属性必须等于 INT32见上文第三节。五、调用说明图模式调用调用方式调用样例说明图模式调用test_geir_histogram_fixed_width.cpp通过 算子IR 构图方式调用 HistogramFixedWidth 算子。示例程序 test_geir_histogram_fixed_width.cpp 展示了完整的 GEGraph Engine图模式调用流程核心步骤可归纳如下初始化 GE通过ge::GEInitialize(global_options)初始化 GE全局选项中设置了ge.exec.deviceId: 0与ge.graphRunMode: 1。构造计算图创建ge::Graph随后在CreateOppInGraph()中实例化算子对象op::HistogramFixedWidth(add1)并依次接入三个输入x形状为{9, 9, 10, 10, 9, 8, 8, 0}的 Data 占位输入默认DT_FLOATrange形状为{2}的常量输入通过op::Const构造并SetAttr(value, tensor)写入具体值nbins形状为{1}、DT_INT32的常量输入。示例中 x 使用宏ADD_INPUT生成全 2 的数据range 与 nbins 使用ADD_CONST_INPUT生成常量这从侧面验证了range 与 nbins 需要在编译期获取数值的设计与ValueDepend(OPTIONAL)及InputsDataDependency配置呼应。建立 Session 并运行创建ge::Sessionsession-AddGraph(graph_id, graph, graph_options)将计算图加入会话再通过session-RunGraph(graph_id, input, output)执行。结果落盘与打印运行结束后将输入输出分别写为tc_ge_irrun_test_0008_npu_input_i.bin/..._output_i.bin二进制文件并将输出 y 的每个元素以result[j] is: %d的形式打印到终端便于直接核对直方图结果。开发者可将该示例作为模板修改 x 的形状、range 的取值与 nbins 数量即可快速验证自定义场景下的直方图输出。六、源码级原理剖析6.1 算子原型注册op_graphhistogram_fixed_width_proto.h 通过REG_OP(HistogramFixedWidth)声明算子 IRREG_OP(HistogramFixedWidth) .INPUT(x, TensorType({DT_FLOAT16, DT_FLOAT, DT_INT32, DT_INT64})) .INPUT(range, TensorType({DT_FLOAT16, DT_FLOAT, DT_INT32, DT_INT64})) .INPUT(nbins, TensorType({DT_INT32})) .OUTPUT(y, TensorType({DT_INT32})) .ATTR(dtype, Int, 3) .OP_END_FACTORY_REG(HistogramFixedWidth);该声明精确锁定了四类输入输出张量类型x/range 支持浮点与整型nbins/y 固定为 INT32以及 dtype 属性的默认值是算子对外可见的契约。6.2 框架适配与 TensorFlow 兼容frameworkhistogram_fixed_width_tf_plugin.cpp 负责将该算子挂接到 TensorFlow 框架REGISTER_CUSTOM_OP(HistogramFixedWidth) .FrameworkType(TENSORFLOW) .OriginOpType(HistogramFixedWidth) .ParseParamsByOperatorFn(AutoMappingFnHistogramFixedWidth) .ImplyType(ImplyType::TVM);其中AutoMappingFnHistogramFixedWidth在完成通用参数映射后会将 GE 图上的dtype属性从ge::DataType枚举值转换为整型static_castint(dataType)完成跨框架的参数归一化这是与 TensorFlow 算子兼容的落地点之一。6.3 形状推导InferShapehistogram_fixed_width_infershape.cpp 注册了IMPL_OP_INFERSHAPE(HistogramFixedWidth)关键逻辑包括校验 range 的形状大小为 2 或动态-1校验 dtype 属性值必须为DT_INT32从常量输入 nbins 中读取区间数量若 nbins 非常量无法读取数值则将输出置为未知形状SetUnknownShape(1, *outputShape)并成功返回否则校验nbins 0后将输出形状设置为一维[nbins]通过.InputsDataDependency({INPUT_NBINS})声明输出形状依赖 nbins 的数值这与算子定义中的ValueDepend(OPTIONAL)一致。6.4 Tiling 策略op_hostTiling 由 histogram_fixed_width_tiling_arch35.cpp 完成注册入口为IMPL_OP_OPTILING(HistogramFixedWidth)并声明TilingInputsDataDependency({HFW_INPUT_IDX_RANGE, HFW_INPUT_IDX_NBINS})即编译期必须拿到 range 与 nbins 的具体数值。Tiling 的完整流程为DoTiling() → ParamCheck() GetSocInfo() CalcTiling()ParamCheck校验 x 的 dtype、读取 nbins 数值、校验输出 y 大小与 nbins 一致并调用ValidateRange校验 max min。GetSocInfo通过platform_ascendc::PlatformAscendC获取可用 AIV 核数coreNum_与 UBUnified Buffer大小ubSize_。CalcTiling依据输入规模与UB 容量的关系在三种加载模式间选择见 tilingkey.h 中的枚举定义loadMode取值选择条件从源码逻辑推断对应 kernelUB_FULL0bins 数量小于 UB 可容纳的 INT32 个数直方图可整体驻留 UBhistogram_fixed_width_simt_full_load.hUB_NOT_FULL1bins 较大但输入规模也较大totalLength_ bins_ / 100直方图在 GM 上原子累加histogram_fixed_width_simt_not_full_load.hUB_NOT_FULL_SIMT2bins 大且输入规模相对小采用 GM 原子累加模式histogram_fixed_width_simt_not_full_load_simt.h此外CalcTiling还完成了数据的分核切分将输入 x 按核数均分为formerLength 尾部tailLength将输出 y 的清零工作也按核切分为clearYFactor/clearYTail最终取两者所需核数的最大值作为needCoreNum并SetBlockDim(needCoreNum)、SetScheduleMode(1)。tiling 结果写入 tilingdata.h 定义的HistogramFixedWidthSimtTilingData结构体供 kernel 侧读取。6.5 Kernel 实现op_kernelKernel 入口在 histogram_fixed_width.cpp通过KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY)声明仅使用 AIV 核执行并根据 tiling key 中的 loadMode 模板参数分发到三种 SIMT 实现。以UB_FULL模式histogram_fixed_width_simt_full_load.h为例其执行流程为并行清零 ySimtCleanWsReduce让 THREAD_NUM默认 512FPGA 环境为 128个线程按clearYFactor分片将 GM 上的 y 清零随后SyncAll()保证全局可见每核局部统计每个计算核在 UB 上申请一块bins * sizeof(int32_t)的本地直方图缓冲Duplicate清零后由UbSimtComputeWsReduce对分配给本核的 x 数据做 6 路循环展开的逐元素处理bin 定位DoBin对每个元素 v 依次判断——若 v 为 ±Infinity、NaNv ! v或-Infinity而 min 非-Infinity直接跳过不计数若 min 为-Infinity归入最后一个 bin若 max 为Infinity归入第一个 bin否则将 v 截断clamp到[min, max]后按binIndex (clamped - min) * bins / range整型用乘法避免除法浮点用invRange预取倒数加速计算 bin 下标并再次夹取到[0, bins-1]通过asc_atomic_add对本地直方图对应位置原子加 1原子写回本核统计完成后SetAtomicAddint32_t()DataCopyPad将 UB 直方图以原子累加方式写回 GM 上的 y多核结果在此合并随后SetAtomicNone()复位。这一分核计算局部直方图 GM 原子累加归并的设计是该算子实现多核并行、同时保证结果正确的核心思路也解释了 Tiling 阶段为何需要根据 bins 大小选择不同加载模式。6.6 二进制配置configascend950/histogram_fixed_width_binary.json 为每种输入数据类型float32 / int32 / int64 / float16分别登记了一个二进制内核条目bin_filename 分别为HistogramFixedWidth_float32、HistogramFixedWidth_int32、HistogramFixedWidth_int64、HistogramFixedWidth_float16并统一声明了各输入输出的 ND 格式、required 属性、动态形状[-2]与 dtype 属性默认值 3。该文件是算子在对应昇腾产品上完成内核二进制打包与分发的依据。七、测试用例仓库为该算子提供了配套的测试与规格文件可据此验证行为与校验逻辑test_histogram_fixed_width_infershape.cpp覆盖 InferShape 的形状推导与非法输入如 range 形状错误、nbins 非正、dtype 属性非法的报错路径test_histogram_fixed_width_tiling.cpp覆盖 Tiling 阶段的参数校验含 max min 约束与 tiling 数据生成histogram_fixed_width_spec.py算子规格定义可支撑生成测试用例或作为约束描述。八、总结HistogramFixedWidth 是 ops-math 中一个实现简洁但设计完整的统计类算子对外兼容 TensorFlow 语义支持浮点与整型四类输入内部则通过算子原型定义 → TF 插件适配 → InferShape 形状推导 → 按 UB 容量选择三种加载模式的 Tiling → AIV 多核 SIMT Kernel 原子累加归并的完整链路在 NPU 上完成直方图的高效计算。理解该算子的源码实现不仅有助于正确使用与调试它也为阅读 ops-math 中其他统计/归约类算子提供了可复用的分析范式。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考