算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载TruncateMod 是 CANN ops-math 仓库中实现的逐元素截断取模truncated modulo算子用于在 Atlas A2/A3 训练与推理系列产品上高效计算y x1 - trunc(x1/x2) * x2余数与被除数同号。本文以 experimental/math/truncate_mod/README.md 与配套的 aclnnTruncateMod 接口文档 为主线结合算子定义、Shape 推导、Tiling 与 Kernel 源码及单元测试完整讲解算子语义、参数与约束、aclnn 两段式 API 调用流程并深入剖析其在 NPU 上的核心切分与计算实现帮助读者既会用、又懂原理。算子功能与数学定义TruncateMod 是一个逐元素element-wise二元算子完成截断除法取余计算对输入x1被除数和x2除数逐元素计算余数余数与x1同号。计算公式为对每个位置i先计算截断商向零取整$$ tq_i trunc(x1_i / x2_i) \lfloor \max(x1_i / x2_i, 0) \rfloor \lceil \min(x1_i / x2_i, 0) \rceil $$再计算余数$$ y_i x1_i - tq_i \cdot x2_i $$其中trunc表示向零取整正数向下取整、负数向上取整。由于取余结果符号跟随被除数-7 mod 2 -1而7 mod (-2) 1这与 Python 的%余数符号跟随除数或 C 语言对负数的行为有所区别是使用前需要特别留意的语义差异。产品支持情况根据算子 README 及 aclnn 接口文档TruncateMod 支持以下产品AICore 配置文件对应ascend910b与ascend910_93见 truncate_mod_def.cpp产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√参数说明参数名输入/输出/属性描述数据类型数据格式x1输入公式中的输入 x1被除数FLOAT16、FLOAT、BFLOAT16、INT32、INT8、UINT8NDx2输入公式中的输入 x2除数shape 与 x1 一致FLOAT16、FLOAT、BFLOAT16、INT32、INT8、UINT8NDy输出公式中的输出 y截断除法余数数据类型与输入一致FLOAT16、FLOAT、BFLOAT16、INT32、INT8、UINT8ND上述参数声明可以在算子注册源码 truncate_mod_def.cpp 中得到印证x1、x2、y均为必选REQUIRED参数支持DT_FLOAT16 / DT_FLOAT / DT_BF16 / DT_INT32 / DT_INT8 / DT_UINT8六种数据类型数据格式统一为FORMAT_ND且都设置了.AutoContiguous()。同时Shape 与数据类型推导在 truncate_mod_infershape.cpp 中实现输出y的 shape 直接拷贝自x1输出数据类型也直接取x1的数据类型即输出与输入x1完全同 shape、同 dtype。约束说明x1、x2、y的 shape 需保持一致不支持广播与仓库中大量支持广播的 element-wise 算子不同这是 TruncateMod 的明确限制。x1、x2的数据类型需保持一致输出y与输入同类型。除零为未定义行为当除数x2为 0 时结果未定义使用方需要自行保证输入不含 0 除数。Tiling 源码 truncate_mod_tiling.cpp 也显式校验了 x1 and x2 dtype must be consistent.不满足时直接返回GRAPH_FAILED。aclnn 两段式接口调用TruncateMod 通过 CANN 的 aclnn 算子库接口对外提供能力。按照 CANN 的两段式接口规范必须先调用aclnnTruncateModGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器再调用aclnnTruncateMod执行计算。函数原型aclnnStatus aclnnTruncateModGetWorkspaceSize( const aclTensor *x1, const aclTensor *x2, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnTruncateMod( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)aclnnTruncateModGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorx1输入待进行 TruncateMod 计算的入参公式中的被除数 x1无FLOAT16、FLOAT、BFLOAT16、INT32、INT8、UINT8ND0-8√x2输入待进行 TruncateMod 计算的入参公式中的除数 x2shape 与 x1 相同数据类型与 x1 相同FLOAT16、FLOAT、BFLOAT16、INT32、INT8、UINT8ND0-8√out输出待进行 TruncateMod 计算的出参公式中的输出 yshape 与 x1 相同数据类型与 x1 相同FLOAT16、FLOAT、BFLOAT16、INT32、INT8、UINT8ND0-8√workspaceSize输出返回需要在 Device 侧申请的 workspace 大小-----executor输出返回 op 执行器包含了算子计算流程-----值得注意接口支持0-8 维 shape的 Tensor且支持非连续 Tensor即 stride 非紧凑的视图 Tensor这由算子定义中的.AutoContiguous()配合框架自动落盘/拷贝机制保证。对于TruncateMod这样的 element-wise 算子实际计算不需要额外 workspaceTiling 函数中将 workspace 大小置为 0见 truncate_mod_tiling.cpp但两段式接口的调用流程依然完整保留。返回值与入参校验接口返回aclnnStatus状态码具体含义参见 aclnn 返回码说明。第一段接口会完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 tensor 是空指针。ACLNN_ERR_PARAM_INVALID161002x1、x2 的数据类型和数据格式不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002x1 和 x2 的数据类型不一致。ACLNN_ERR_PARAM_INVALID161002x1 和 out 的数据形状不一致。aclnnTruncateMod 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnTruncateModGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。完整调用示例以下示例代码取自 aclnnTruncateMod 接口文档仓库中 examples/test_aclnn_truncate_mod.cpp 提供了支持多 dtype 转换的完整版本编译与运行流程参见 编译与运行样例#include iostream #include vector #include acl/acl.h #include aclnn_truncate_mod.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } 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); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API手册 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t x1Shape {2, 2}; std::vectorint64_t x2Shape {2, 2}; std::vectorint64_t outShape {2, 2}; void* x1DeviceAddr nullptr; void* x2DeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* x1 nullptr; aclTensor* x2 nullptr; aclTensor* out nullptr; std::vectorfloat x1HostData {7.0f, -7.0f, 5.0f, 1.0f}; std::vectorfloat x2HostData {2.0f, 2.0f, 3.0f, 4.0f}; std::vectorfloat outHostData {0.0f, 0.0f, 0.0f, 0.0f}; // 期望: 7 - trunc(7/2)*2 1, -7 - trunc(-7/2)*2 -1, 5 - trunc(5/3)*3 2, 1 - trunc(1/4)*4 1 // 创建x1 aclTensor ret CreateAclTensor(x1HostData, x1Shape, x1DeviceAddr, aclDataType::ACL_FLOAT, x1); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建x2 aclTensor ret CreateAclTensor(x2HostData, x2Shape, x2DeviceAddr, aclDataType::ACL_FLOAT, x2); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnTruncateMod第一段接口 ret aclnnTruncateModGetWorkspaceSize(x1, x2, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnTruncateModGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnTruncateMod第二段接口 ret aclnnTruncateMod(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnTruncateMod failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy( resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor需要根据具体API的接口定义修改 aclDestroyTensor(x1); aclDestroyTensor(x2); aclDestroyTensor(out); // 7. 释放device资源需要根据具体API的接口定义修改 aclrtFree(x1DeviceAddr); aclrtFree(x2DeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例以{7, -7, 5, 1}对{2, 2, 3, 4}逐元素求余期望结果依次为1, -1, 2, 1直观验证了余数与被除数同号的语义-7 mod 2 -1。仓库中的 test_aclnn_truncate_mod.cpp 还额外演示了 FLOAT16/BF16 等低精度类型在 Host 侧构造数据时的FloatToHalf/FloatToBFloat16位级转换方法以及在workspaceSize 0时使用aclDestroyOpExecutor释放执行器的完整收尾流程可作为多 dtype 调用的参考模板。源码级实现原理Host 侧Tiling 三级切分truncate_mod_tiling.cpp 实现了 element-wise 算子的经典三级切分策略core split核间切分→ UB split片上切分→ tile split瓦片切分并通过TruncateModTilingData见 truncate_mod_tiling_data.h将切分结果下发给 Kernel。核间切分按 32 字节 blockBYTES_PER_BLOCK 32对齐计算总 block 数以每核最少 4096 字节工作量BLOCK_PER_CORE 128为粒度决定实际启用的 AIV 核数coreNum将数据尽量均匀分给每个核不足 block 对齐的尾元素tailElems由最后一个核吸收。UB 切分根据每个元素在 UB 上消耗的空间3 个输入/输出队列 3 个 float 计算缓冲 half 中转缓冲 mask 缓冲计算单次可处理的元素数tileLength当单核数据多于一个 tile 时自动开启双缓冲bufferNum 2隐藏搬运与计算延迟。tile 切分将每个核的数据量拆分为若干完整 epoch 与一个尾 tiletailTileLength末核还要叠加tailElems。值得一提的是Tiling 依据x1的 dtype 选择模板参数 schModetiling key0float16、1float32、2bfloat16、3int32、4int8、5uint8定义见 truncate_mod_tiling_key.h并用GET_TPL_TILING_KEY将运行期 dtype 映射为编译期模板分支实现同一份 Kernel 代码、按 dtype 特化的高效编译。Device 侧Kernel 统一升 float 计算truncate_mod.cpp 是 Kernel 入口按 schMode 将不同 dtype 分派到NsTruncateMod::RunT其中 bfloat16 仅在__CCE_AICORE__ 220的平台上启用即对应 Atlas A2/A3。核心实现类 TruncateModKernel 的关键设计包括统一提升到 float 计算Kernel 头文件注释明确说明所有 dtype 在 UB 上统一提升到 float 计算。对 int8/uint8 因 AICORE 无 int8 与 float 直转指令采用int8 → half → float两级 Cast 中转值域 ≤255half 可精确表示对 int32 使用Castfloat, int32_tfloat16/bfloat16 直接 Cast 到 float。截断取整的向量化实现Compute中用Div求商t x1 / x2后按公式trunc(t) ceil(min(t, 0)) floor(max(t, 0))用Mins/Ceil/Maxs/Floor/Add五个向量指令完成向零取整再Mul回乘x2、Sub得到余数全程无需标量分支且避免了直接 Cast 取整在负数语义上的歧义与 int32 溢出风险。尾块安全搬运GM↔UB 搬运采用DataCopyPad配合DataCopyExtParams指定实际长度尾块非 32B 对齐时也能避免越界读写UB 内计算则按VEC_REPEAT_F32 64元素对齐长度执行保证向量指令吞吐。写回时对 int32/int8/uint8 等整型使用CAST_ROUND舍入模式。零 workspaceelement-wise 算子无需归约与中间结果落盘用户 workspace 大小置 0这也是两段式接口中workspaceSize恒为 0 的原因。测试与验证算子配套了完整的 Host/Device 双层单测。Kernel 侧 test_truncate_mod.cpp 通过tikicpulib在 CPU 仿真环境运行 Kernel并调用数据生成脚本 gen_data.py 生成随机输入与 golden 数据脚本用np.trunc(x1/x2)计算截断商再按y x1 - trunc(x1/x2) * x2生成期望结果覆盖 float16/float32/int32/int8/uint8 五种类型被除数取 -50~50除数取 1~10避免除零Host 侧 test_truncate_mod_tiling.cpp 验证 Tiling 数据的正确性。构建层面CMakeLists.txt 通过add_all_modules_sources(OPTYPE truncate_mod ACLNNTYPE aclnn)将算子纳入仓库统一的编译与 aclnn 接口生成体系。贡献说明贡献者贡献方贡献算子贡献时间贡献内容weixin_45448057个人开发者TruncateMod2026/07/14TruncateMod 算子适配开源仓小结TruncateMod 是一个语义清晰、实现简洁但又不失工程细节的典型 element-wise 算子数学上它定义了一种余数随被除数的截断取模语义对外通过 aclnn 两段式接口暴露 0-8 维、支持非连续 Tensor 的通用调用方式对内则体现了 ops-math 仓库通用的三级 Tiling 切分、dtype 模板特化schMode、UB 统一 float 计算与尾块安全搬运等 AscendC 算子开发范式。无论是直接调用该算子还是将其作为学习 CANN 自定义算子定义—Shape 推导—Tiling—Kernel—单测完整链路的范例本文所梳理的 README、接口文档 与源码文件都值得进一步阅读。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math TruncateDiv 算子实战指南截断除法语义、aclnn 两段式调用与 NPU 源码实现解析CANN ops math TruncateDiv 算子实战指南截断除法语义、aclnn 两段式调用与 NPU 源码实现解析 TruncateDiv 是 CA算子库人工智能CANNCANN ops-math Addcmul 算子全解析两段式 aclnn 接口原理与 NPU 实战调用CANN ops math Addcmul 算子全解析两段式 aclnn 接口原理与 NPU 实战调用 导读 Addcmul 是 CANN ops math算子库人工智能CANNCANN ops-math Digamma 算子全解析数学原理、aclnn 两段式接口与源码实现CANN ops math Digamma 算子全解析数学原理、aclnn 两段式接口与源码实现 Digammadigamma是 CANN ops mat算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考