FmodScalar / FmodTensor 算子设计文档【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions需求背景需求来源通过社区任务完成昇腾开源算子贡献的需求。参考昇腾 CANN 内置 TBEMod算子实现在 Atlas A2 / Atlas A3 系列产品上使用 Ascend C 实现功能一致的aclnnFmodScalar与aclnnFmodTensor算子。算子交付目标包括算子工程代码、README、自验证报告、设计文档以及多组 aclnn 调用测试代码。背景介绍Fmod算子用于执行逐元素取模计算。该算子的语义是基于向 0 截断除法的取余而不是基于向下取整的取余。参考 TBE 源码中的注释truncate_mod(x, y) x - truncate_div(x, y) * y其中x对应self/input_xy对应other/input_ytruncate_div(x, y)表示x / y后向 0 截断得到的商输出结果与x同号或为 0。计算公式为out self - trunc(self / other) * other当self / other大于等于 0 时trunc(self / other)等价于floor(self / other)当self / other小于 0 时trunc(self / other)等价于ceil(self / other)。因此 TBE 实现中通过vmin、vmax、floor、ceil组合实现向 0 截断。本任务包含两个 aclnn 接口aclnnFmodScalarself为输入 Tensorother为 Scalarout的 shape 与self一致。aclnnFmodTensorself为输入 Tensorother为输入 Tensorout的 shape 与self一致other需要满足与self的合法广播关系或在同 shape 场景下逐元素计算。算子实现信息TBE 参考实现kernel 实现/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/impl/ops_legacy/dynamic/mod.py算子原型/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_proto/inc/ops_proto_math.h算子信息库/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/config/ascend910b/aic-ascend910b-ops-info-legacy.json开源仓路径开源仓https://gitcode.com/cann/ops-math参考样例https://gitcode.com/cann/ops-math/tree/master/math/mod目标合入路径experimental/math目标硬件与语言适配硬件Atlas A2 训练系列产品 / Atlas A3 系列产品开发语言Ascend C支持格式ND支持 dtypeBFLOAT16、FLOAT16、FLOAT32、INT16TBE 算子现状分析TBE 算子支持方式参考mod.pyTBE 版本通过如下方式实现动态 shape、广播、融合和 bfloat16 支持register_operator_compute(Mod, op_modedynamic, support_fusionTrue, support_bfp16True) def mod_compute(input_x, input_y, output_z, kernel_namemod): ...TBE 入口函数通过ins classify([input_x, input_y], OpPatternMode.ELEWISE_WITH_BROADCAST)对输入进行ELEWISE_WITH_BROADCAST分类再对每一组分类结果创建 TVM placeholder、构造计算图并自动调度。TBE 源码支持的 dtype 为check_list (float16, float32, int8, uint8, int32, bfloat16)本任务的目标 dtype 为BFLOAT16、FLOAT16、FLOAT32、INT16因此 Ascend C 设计以任务书要求为准同时参考 TBE 的核心计算语义与升精度策略。TBE 参数校验与 shape 处理TBE 入口函数mod(input_x, input_y, output_z, kernel_namemod)的主要处理流程如下通过para_check.check_op_params检查两个输入、一个输出以及kernel_name。通过shape_util.compare_tensor_dict_key(input_x, input_y, dtype)要求两个输入 dtype 一致。通过para_check.check_dtype(input_dtype, check_list, param_nameinput_x)检查输入 dtype 是否在支持列表中。通过classify([input_x, input_y], OpPatternMode.ELEWISE_WITH_BROADCAST)对动态 shape 与广播场景进行分类。对每组分类结果调用shape_util.variable_shape([input_x_dict, input_y_dict])获取动态 shape调用shape_util.refine_shapes_for_broadcast(shape_x, shape_y)对广播 shape 进行规整创建data_x、data_y两个 placeholder调用mod_compute构造计算图调用tbe.auto_schedule(res)自动调度最终通过tbe.build(schedule, config)编译生成算子。TBE 核心计算逻辑TBEmod_compute的核心逻辑如下读取输入 shape 与 dtype。若输入 dtype 不是float32且平台支持float32的vdiv则将input_x与input_y升精度到float32。若两个输入 shape 不一致则调用shape_util.broadcast_shapes获取广播后的 shape并对两个输入执行tbe.broadcast。计算除法data_div tbe.vdiv(input_x, input_y)构造 0 常量向量data_zero tbe.broadcast(tvm.const(0, float32), shape_broadcast, float32)将data_div拆分为非正部分与非负部分data_div_min tbe.vmin(data_div, data_zero) data_div_max tbe.vmax(data_div, data_zero)对非负部分取floor对非正部分取ceildata_div_max_floor tbe.floor(data_div_max, float32) data_div_min_ceil tbe.ceil(data_div_min, float32)两部分相加得到向 0 截断后的商data_div_res tbe.vadd(data_div_max_floor, data_div_min_ceil)将截断商转换为input_y的 dtype执行乘法和减法data_div_res tbe.cast_to(data_div_res, input_y.dtype.lower()) data_mul tbe.vmul(data_div_res, input_y) res tbe.vsub(input_x, data_mul)若前面做过升精度则将结果转回原始 dtype。TBE 计算语义说明TBE 的核心公式可等价表示为q_pos floor(max(self / other, 0)) q_neg ceil(min(self / other, 0)) q q_pos q_neg out self - q * other该实现等价于out self - trunc(self / other) * other其中trunc表示向 0 截断。TBE 算子实现流程图需求详细设计使能方式通过 aclnn 接口调用本算子aclnnFmodScalar用于 Tensor 与 Scalar 的逐元素 fmod 计算。aclnnFmodTensor用于 Tensor 与 Tensor 的逐元素 fmod 计算。Tensor 版本需要支持同 shape 与合法广播场景。算子原型aclnnFmodScalar参数名类别描述数据类型数据格式shape非连续 Tensorself输入 Tensor待进行 Fmod 的输入 TensorBFLOAT16、FLOAT16、FLOAT32、INT16ND0~8 维√other输入 Scalar待进行 Fmod 的 ScalarBFLOAT16、FLOAT16、FLOAT32、INT16---out输出 TensorFmod 输出结果shape 与 self 一致BFLOAT16、FLOAT16、FLOAT32、INT16ND与 self 一致按接口要求aclnnFmodTensor参数名类别描述数据类型数据格式shape非连续 Tensorself输入 Tensor待进行 Fmod 的左操作数BFLOAT16、FLOAT16、FLOAT32、INT16ND0~8 维√other输入 Tensor待进行 Fmod 的右操作数BFLOAT16、FLOAT16、FLOAT32、INT16ND0~8 维√out输出 TensorFmod 输出结果shape 与 self 一致BFLOAT16、FLOAT16、FLOAT32、INT16ND与 self 一致√功能约束与语义约束输出 shapeFmodScalarout.shape self.shape。FmodTensorout.shape self.shapeother需要能够广播到self.shape或与self.shape完全一致。dtype 支持支持BFLOAT16、FLOAT16、FLOAT32、INT16。self、other、out的 dtype 需要满足 aclnn 原型侧的类型推导与可转换规则。kernel 实现中优先按输出 dtype 写回。数据格式支持 ND。shape 范围支持 0~8 维。需要覆盖 0 维标量 Tensor、空 Tensor、小 shape、大 shape、广播 shape、非连续 Tensor 等泛化场景。除 0 行为本设计不在 kernel 内额外定义除 0 特殊规则行为以 CANN/TBE 内置算子实际行为为准。测试阶段需要覆盖other 0场景并与 TBE 结果保持一致。负数行为必须按向 0 截断取余实现保证结果与 TBEMod语义一致。例如fmod(5.3, 2.0) 1.3fmod(-5.3, 2.0) -1.3fmod(5.3, -2.0) 1.3fmod(-5.3, -2.0) -1.3需求总体设计总体实现思路本算子属于 elementwise 二元算子整体实现分为 host 侧 tiling 与 kernel 侧计算两部分host 侧解析输入 shape、stride、dtype、数据量、广播关系判断是否为连续 Tensor根据 shape、广播、非连续情况选择tilingKey计算分核策略、单核处理长度、尾块长度、UB 分块大小将必要的 tiling 参数写入 tiling data。kernel 侧根据tilingKey选择连续、广播、非连续或 scalar 路径按 block 维度分配 GM 数据范围将输入数据搬入 UB对当前 tile 执行 fmod 计算将结果搬出到 GM。host 侧设计1. 输入参数解析host 侧需要解析如下信息self的 shape、stride、storage offset、dtype、元素个数other的 shape、stride、storage offset、dtype或 scalar 值out的 shape、stride、storage offset、dtypeself与other的广播关系self、other、out是否为连续 Tensor当前 dtype 对应的元素字节数总元素数totalLength NumElements(self)。2. 分核策略优先采用满核参与计算策略。设可用 AI Core 数为coreNum总元素数为totalLength。当totalLength coreNum时按如下方式分配blockLength ceil(totalLength / coreNum) usedCoreNum coreNum当totalLength coreNum时blockLength 1 usedCoreNum totalLength每个核根据blockIdx计算自己的起始逻辑下标start blockIdx * blockLength end min(start blockLength, totalLength)对于空 Tensorhost 侧可直接返回或生成totalLength 0的 tiling由 kernel 快速退出。3. UB 分块策略Fmod 计算至少需要以下 UB bufferxLocalself tileyLocalother tile 或 scalar broadcast 后 tiletmpDiv除法中间结果tmpMin/tmpMax截断商计算中间结果或复用临时 buffertmpFloor/tmpCeilfloor / ceil 结果或复用临时 bufferqLocal截断商outLocal输出 tile。为降低 UB 占用建议进行 buffer 复用tmpDiv保存x / y。tmpMax可在计算floor(max(div, 0))后复用。tmpMin可在计算ceil(min(div, 0))后复用。qLocal tmpMax tmpMin。outLocal xLocal - qLocal * yLocal。对于FLOAT16、BFLOAT16、INT16若升精度到FLOAT32计算则 UB 中间 buffer 需要按float规划tile 长度需要相应减小。4. tilingKey 规划建议使用如下tilingKeytilingKey场景说明0Scalar 连续 self/outFmodScalar高性能路径other 为 scalarself/out 连续1Tensor 同 shape 连续输入输出FmodTensor高性能路径self/other/out 连续且 shape 一致2Tensor broadcast 连续输入输出other 可广播到 self逻辑连续访问3Scalar 非连续 self 或 outscalar 场景的通用 stride 路径4Tensor 非连续输入或输出tensor 场景的通用 stride / broadcast 路径首版实现若需要控制风险可优先完成tilingKey 0/1/2再补齐非连续路径。但任务验收要求覆盖非连续 Tensor因此最终版本必须实现tilingKey 3/4或通过框架侧连续化机制保证非连续场景正确。5. 非连续 Tensor 处理策略对于非连续 Tensor逻辑下标到物理 GM offset 的转换如下physical_offset storage_offset sum(index[d] * stride[d])其中index[d]是逻辑多维下标stride[d]是对应维度的物理步长storage_offset是 Tensor 起始偏移。host 侧需要将 0~8 维的 shape 与 stride 写入 tiling data。kernel 侧根据线性逻辑下标还原多维下标再计算self、other、out的物理地址。对于广播维度if other_shape[d] 1 and self_shape[d] ! 1: other_index[d] 0 else: other_index[d] self_index[d]kernel 侧设计连续 Tensor 高性能路径连续路径中每个核处理一段连续逻辑区间Scalar 路径Scalar 路径中other是 aclScalar。host 侧需要将 scalar 值写入 tiling data 或通过 kernel 参数传入。kernel 侧将 scalar 转换为CalcT后在 UB 内 broadcast 为向量yLocal broadcast(otherScalar) outLocal fmod(xLocal, yLocal)Scalar 路径不需要处理other的 GM 读入因此性能通常优于 Tensor 路径。Tensor broadcast 路径Tensor broadcast 路径需要根据逻辑下标计算other的读取位置。对于连续但 broadcast 的场景可按以下方式优化若other是 0 维或元素数为 1可退化为 Scalar-like 路径。若 broadcast 只发生在外层维度内层连续维与self一致可按内层连续块批量搬运。若 broadcast 发生在内层维度需要按维度切分避免逐元素 scalar 搬运导致性能下降。对一般 0~8 维 broadcast可实现通用 index 映射路径作为保底正确性路径。非连续通用路径非连续路径优先保证正确性为提升性能可将连续子段识别为小块批量搬运无法形成连续块时使用 DataCopyPad / Gather-like 标量读写或循环处理。Ascend C 实现流程图Ascend C 实现与 TBE 实现差异点TBE 通过classify(..., ELEWISE_WITH_BROADCAST)与auto_schedule自动处理广播与调度Ascend C 需要显式实现 tiling、分核、UB 分配、搬运和广播索引映射。TBE 中tbe.broadcast是逻辑表达式图操作Ascend C 中需要根据不同场景选择 scalar broadcast、连续搬运、按维映射或非连续通用路径。TBE 根据平台能力在计算前对非float32类型升精度到float32Ascend C 需要通过模板参数与 Cast 指令显式实现对应策略。TBE 的floor/ceil由 DSL 表达式描述Ascend C 需要确认对应 Vector API 能力并在不支持某 dtype 的情况下通过float32中间类型实现。TBE 的非连续 Tensor 通常由上层图编译和 Tensor 描述协同处理Ascend C 实现需要明确处理 stride、storage offset 与输出非连续写回。精度设计精度目标精度要求满足 AscendOpTest 工具默认阈值并与内置 TBE 算子结果对齐。性能设计性能目标所有核参与计算场景下性能不低于原 TBE 算子的 95%。小 shape 如果无法达标10us 以下场景允许通过性能仿真图和分析结论说明 Ascend C 实现与 TBE 一致或优于 TBE。支持硬件产品是否支持Atlas A2 训练系列产品√Atlas A3 系列产品√README 使用示例建议README 中建议提供如下内容算子功能说明计算公式支持 dtype / shape / formataclnnFmodScalar调用示例aclnnFmodTensor调用示例编译方法运行测试方法精度和性能对比说明。【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考