CANN ops-math FloorMod 算子全解析：向下取整取模的数学语义、aclnn 调用与 AscendC 内核实现

CANN ops-math FloorMod 算子全解析向下取整取模的数学语义、aclnn 调用与 AscendC 内核实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-mathFloorMod 是 CANN ops-math 数学算子库中负责逐元素向下取整取模floor modulo的基础算子它计算y x1 - floor(x1 / x2) * x2结果的符号与除数x2保持一致是深度学习中归一化、周期性 padding、索引取整等场景的高频基础运算。本文以 FloorMod 算子 README 为骨架结合仓库内的算子定义、tiling、内核与测试源码完整讲解其数学语义、参数与约束、aclnn 两段式调用方式并深入其 host 侧注册与 device 侧内核实现原理帮助读者既能在 CANN 环境下正确调用该算子也能理解其底层工作机制。一、算子功能与数学定义1.1 功能说明FloorMod 对输入张量进行逐元素element-wise取模运算其核心语义是结果的符号与除数x2一致即采用向下取整floor方式的模运算这与 C/C 中截断式取余truncated remainder结果符号与被除数一致行为不同与 Python 的%运算符语义一致。1.2 计算公式算子按照如下公式计算每个输出元素$$ y_i x1_i - \lfloor \frac{x1_i}{x2_i} \rfloor \times x2_i $$其中x1为被除数dividendx2为除数divisory为取模结果remainder。由公式可见y的符号完全由除数x2决定。例如5 % -3 -1因为floor(5 / -3) -25 - (-2) × (-3) -1而-5 % 3 1。这一点在 aclnnRemainderTensorTensor 接口文档 与算子 README 中均被明确强调也是与 TruncateMod截断取模类算子最本质的差异。二、产品支持情况根据 FloorMod 算子 README 的产品支持矩阵产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√需要说明的是该支持矩阵是当前仓库 README 声明的范围从算子注册代码看内核的 TilingKey 在编译期以模板方式实例化浮点与整型多种组合并在 floor_mod_def.cpp 中通过this-AICore().AddConfig(ascend910b)声明了 AICore 配置具体可用硬件范围请以实际部署环境为准。三、参数说明FloorMod 算子共包含两个输入与一个输出全部为 ND 格式支持的数据类型完全一致。详见下表来自 README 参数说明参数名输入/输出/属性描述数据类型数据格式x1输入待进行 FloorMod 计算的被除数入参FLOAT、FLOAT16、BFLOAT16、INT32NDx2输入待进行 FloorMod 计算的除数入参FLOAT、FLOAT16、BFLOAT16、INT32NDy输出FloorMod 计算的输出结果FLOAT、FLOAT16、BFLOAT16、INT32NDx1被除数参与取模计算的分子张量。x2除数参与取模计算的分母张量其符号决定输出结果的符号。y输出计算得到的结果张量shape 与广播后的输入一致。上述参数定义与 op_host/floor_mod_def.cpp 中的算子信息注册完全对应Input(x1)、Input(x2)、Output(y)的数据类型集合均为ge::DT_BF16 / ge::DT_FLOAT16 / ge::DT_FLOAT / ge::DT_INT32数据格式均为ge::FORMAT_ND且都声明为REQUIRED必选参数。四、约束说明使用 FloorMod 算子时需满足以下约束广播规则x1和x2的 shape 需要满足广播Broadcast规则输出y的 shape 为两者广播后的 shape。类型一致性x1与x2的数据类型需要一致或者满足隐式类型转换规则。除此之外从 aclnnRemainderTensorTensor 接口文档 的说明可以补充两点细节约束维度上限输入与输出 Tensor 的维度不能超过 8 维空 Tensorself或other为空 Tensor 时接口直接返回空 Tensor不执行实际计算。五、调用方式aclnn 两段式接口FloorMod 算子通过 aclnn 接口对外提供能力。仓库中实际暴露的 aclnn 接口名为aclnnRemainderaclnn_remainder.h中声明底层在图编译阶段展开为 FloorMod 算子执行。该接口采用 CANN aclnn两段式调用模式第一段aclnnRemainderTensorTensorGetWorkspaceSize—— 根据具体计算流程计算 workspace 大小并返回包含算子计算流程的executor执行器第二段aclnnRemainderTensorTensor—— 传入第一段申请的 workspace 与 executor在指定 stream 上执行计算。两段式接口的函数原型如下详见 aclnn_remainder.haclnnStatus aclnnRemainderTensorTensorGetWorkspaceSize( const aclTensor* self, const aclTensor* other, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor); aclnnStatus aclnnRemainderTensorTensor( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream);5.1 接口变体从 aclnn_remainder.h 的声明可见aclnnRemainder 一族接口提供了丰富的重载变体覆盖不同类型输入组合接口变体输入组合说明aclnnRemainderTensorTensorTensor ÷ Tensor最通用的张量-张量取模支持广播aclnnRemainderTensorScalarTensor ÷ Scalar张量除以标量Scalar 在 host 侧传入aclnnRemainderScalarTensorScalar ÷ Tensor标量除以张量aclnnInplaceRemainderTensorTensorTensor ÷ Tensor原地结果直接写回selfRef无需额外输出aclnnInplaceRemainderTensorScalarTensor ÷ Scalar原地原地张量-标量取模5.2 第一段接口的参数说明以aclnnRemainderTensorTensorGetWorkspaceSize为例详细表格见 aclnnRemainderTensorTensor 文档参数名输入/输出描述数据类型数据格式维度(shape)非连续Tensorself输入被除数入参数据类型需与 other 满足推导规则FLOAT、FLOAT16、BFLOAT16、INT32ND0-8√other输入除数入参数据类型需与 self 满足推导规则shape 需与 self 满足 broadcast 关系FLOAT、FLOAT16、BFLOAT16、INT32ND0-8√out输出出参shape 需是 self 与 other broadcast 后的 shapeFLOAT、FLOAT16、BFLOAT16、INT32ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小----executor输出返回 op 执行器包含算子计算流程----该接口支持非连续 Tensor输入内部会自动做连续性处理见下文源码分析数据格式仅支持 ND。5.3 返回码与报错场景两段式接口均返回aclnnStatus状态码。第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针ACLNN_ERR_PARAM_INVALID161002self、other 和 out 的数据类型和数据格式不在支持范围之内或三者数据维度超过 8 维或 self 和 other 无法做 broadcast以及 broadcast 后的 shape 与 out 不一致六、完整调用示例仓库在 examples/test_aclnn_floor_mod.cpp 中提供了可直接参考的完整可运行示例覆盖了环境初始化、Tensor 创建、两段式计算、结果回读与资源释放的全流程。核心调用逻辑如下#include iostream #include vector #include acl/acl.h #include aclnn_remainder.h // 1. 环境初始化aclInit - aclrtSetDevice - aclrtCreateStream int Init(int32_t deviceId, aclrtStream* stream) { ... } // 2. 构造 aclTensorhost 数据经 aclrtMalloc aclrtMemcpy 拷贝到 device再 aclCreateTensor 创建 template typename T int CreateAclTensor( const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); ... ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); ... *tensor aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } // 3. 两段式执行计算 int Compute(aclrtStream stream, aclTensor* self, aclTensor* other, aclTensor* out) { uint64_t workspaceSize 0; aclOpExecutor* executor; // 第一段获取 workspace 大小与执行器 auto ret aclnnRemainderTensorTensorGetWorkspaceSize(self, other, out, workspaceSize, executor); if (ret ! ACL_SUCCESS) { ... return ret; } // 按需申请 workspace void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); ... } // 第二段执行计算 ret aclnnRemainderTensorTensor(workspaceAddr, workspaceSize, executor, stream); ... ret aclrtSynchronizeStream(stream); ... if (workspaceSize 0) { aclrtFree(workspaceAddr); } return 0; }示例中使用的测试数据selfShape otherShape outShape {4, 4}同时覆盖了正负被除数与正负除数的组合std::vectorfloat selfData {5.5, -11.51, 36.23, 7, -10, -8, -15, -7, 10, 8, 15, 7, -10, -8, -15, -7}; std::vectorfloat otherData {2, 3, -24.1, 2, 3, 5, 4, 2, -3, -5, -4, -2, -3, -5, -4, -2};计算完成后通过aclrtMemcpyACL_MEMCPY_DEVICE_TO_HOST将结果回读并逐元素打印最后依次aclDestroyTensor、aclrtFree、aclrtDestroyStream、aclrtResetDevice、aclFinalize释放全部资源。读者可将该文件作为模板替换输入数据与 shape 即可快速验证不同场景下的算子行为。七、源码级实现纵深FloorMod 在仓库中遵循 CANN 算子标准的四层结构组织op_api对外接口层、op_host算子定义与 host 侧逻辑、op_kernelAscendC 内核、tests单元测试。各文件分布于 experimental/math/floor_mod 目录下。7.1 Host 侧算子注册与 shape 推导算子定义floor_mod_def.cpp通过OpDef注册算子名FloorMod声明x1、x2两个必选输入与y一个必选输出并在同一位置为每个参数并列声明了支持的数据类型集合与 ND 格式同时在编译期声明了 AICore 配置ascend910b最终通过OP_ADD(FloorMod)写入算子信息库。shape 推导floor_mod_infershape.cpp逻辑非常直接读取输入x1的 shape 并直接赋值给输出y*yShape *xShape即输出 shape 与输入一致广播后的 shape 语义由上层 aclnn 接口负责统一。7.2 Host 侧Tiling 策略Tiling数据切分是决定算子在大规模数据下性能的关键。从 floor_mod_tiling.cpp 可以看到其核心设计每核最小数据量MINIMUM_ELEMENT_PER_CORE 1024即元素总数除以 1024 得到基础核数再受硬件totalCoreNumAIV 核数上限约束数据块对齐DATA_BLOCK 64所有切分粒度均按 64 对齐UB 空间划分按不同数据类型使用不同的 UB 除数UB_DIVIDER_FP32 50、UB_DIVIDER_FP16 46、UB_DIVIDER_INT32 66并预留RESERVERD_UB_SIZE 1024字节通过usableUbSize (ubSize - 1024 - sizeof(FloorModTilingData)) / ubDivider计算单核实际可用容量尾块处理通过tailDataCoreNum、lastCoreDataCount两个字段描述尾块数据在多核间的分布避免核间负载不均Workspace固定申请32 * 1024 * 102432MB的 workspaceWORK_SPACE_SIZE通过tilingContext-SetBlockDim(tilingData-needCoreNum)设置实际并行核数并以 TilingKey 区分不同数据类型的内核模板实例。op_kernel/floor_mod.cpp 中的内核入口正是按 TilingKey 编译期分派REGISTER_TILING_DEFAULT(FloorModTilingData)注册 tiling 数据GET_TILING_DATA_WITH_STRUCT解析 tiling随后实例化对应数据类型的FloorModKernelImpl。7.3 Device 侧AscendC 内核实现内核主体在 op_kernel/floor_mod.h 中采用经典的CopyIn → Compute → CopyOut流水线结构QUEUE_DEPTH 2双缓冲按数据类型分为三条计算路径浮点路径FLOAT16 / BFLOAT16 / FLOAT核心是ComputeFPCore函数floor_mod.hDiv求商 →Floor向下取整 →Mul乘回除数 →Sub求余即直接实现文档公式除数 ±Inf 特判|x2| Inf时结果直接取x1Select(resRem, MaskTensor, x1Float, resRem, ...)被除数 ±Inf 特判|x1| Inf时结果置为 NaN符号修正当resRem * x2 0时将结果调整为resRem x2保证结果符号与除数一致即 floor 语义。对于 FLOAT16 / BFLOAT16输入会先Cast到 FLOAT32 精度计算再Cast回原类型以规避低精度下的舍入误差bfloat16 使用CAST_RINT舍入模式。INT32 路径内核实现了高位拆分split高精度算法ComputeInt32由于 INT32 范围超过 FLOAT32 可精确表示的整数范围FLOAT32 尾数 24 位直接转浮点计算会丢失精度。内核将x1拆分为高 8 位ShiftRight(x1, 24)与低 24 位两部分分别求商取整再合并同时在HIGH_PRECISION编译宏下使用FP32MaxValid 2^24与INT32MaxValid 2^24作为分界常量参与修正从而在整型域内保持精确结果。7.4 aclnn 接口层参数校验与执行流程op_api/aclnn_remainder.cpp 完整实现了 aclnnRemainder 各变体其核心流程为入参校验CheckParamsTensorTensor等空指针检查、数据类型推导PromoteType合法性检查、支持类型列表检查、维度 ≤ 8 检查、self/other/out三者 broadcast 关系检查空 Tensor 短路self-IsEmpty() || other-IsEmpty()时直接返回workspace 为 0输入预处理InitializeTensor对非连续 Tensor 调用l0op::Contiguous转为连续0 维 Tensor 通过BroadcastTo升为 1 维再按 promote 类型Cast广播对齐BroadcastTensor统一ReFormat为 ND 格式再BroadcastTo到输出 shape核心计算调用底层l0op::FloorMod声明于 op_api/floor_mod.h经RemainderMainProcess完成取模、结果Cast回输出类型、0 维场景SqueezeNd还原、最后ViewCopy写回输出执行第二段接口统一通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)下发执行。不同 SoC 版本支持的数据类型集合存在差异见GetDtypeSupportList例如 ASCEND910B 平台额外支持 BF16该函数同时通过IsRegBase()判断是否走 RegBase 快速路径省去中间 broadcast 张量构造直接l0op::FloorModCastViewCopy。八、测试验证仓库为 FloorMod 提供了完整的三层单元测试可用于验证上述语义与实现op_api 层test_remainder_tensor_tensor.cpp 通过OP_API_UT(aclnnRemainderTensorTensor, ...)覆盖各类 dtype 组合INT32/INT64/FLOAT16/FLOAT/DOUBLE/BF16、广播 shape、0-8 维场景以及非法参数test_run_invalid断言返回ACLNN_ERR_PARAM_INVALID并对输出设置精度校验Precision(0.00001, 0.00001)。op_host 层test_floor_mod_tiling.cpp 以128×64的 FLOAT16 输入为例断言 tiling 产生的 TilingKey 为1315860、tiling 数据为34359742592 8192 1024 0 1024、workspace 为32MB可精确回归 host 侧切分逻辑。op_kernel 层tests/ut/op_kernel/ 下配套了 gen_data.py 与 compare_data.py 数据生成与比对脚本以及test_floor_mod.cpp内核测试用于校验 device 侧计算结果。模块构建入口见 floor_mod/CMakeLists.txt通过add_all_modules_sources(OPTYPE floor_mod ACLNNTYPE aclnn_exclude)纳入仓库构建体系。九、贡献信息FloorMod 算子是社区贡献算子信息记录于 README 贡献说明贡献者贡献方贡献算子贡献时间贡献内容Tream个人开发者FloorMod2025/12/17FloorMod 算子适配开源仓结语FloorMod 虽然是一个语义简单的逐元素算子但在 CANN ops-math 仓库中体现了完整的算子工程化链路从数学语义与 dtype 约束的定义、aclnn 两段式接口的多变体封装到 host 侧按数据类型差异化切分的 tiling 策略再到 device 侧针对 INT32 精度与浮点 Inf/NaN 边界的内核优化。读者无论是需要在 NPU 上直接调用该算子参考 调用示例还是希望理解 CANN 算子的标准实现范式本文所梳理的文档与源码脉络都可以作为快速入手的索引。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考