CANN ops-math Real 算子深度解析：复数实部提取的原理、Tiling 分核策略与 ACLNN 调用实践

算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载Real 算子是 CANN ops-math 数学算子库experimental/math/real中负责提取复数张量实部的单目算子对于复数输入输出其实部对于实数输入则执行恒等拷贝。本文围绕该算子的功能语义、参数与数据类型映射、Host 侧 Tiling 多核分核算法、Kernel 侧基于 GatherMask 的复杂双路径实现以及 ACLNN 两段式调用接口展开帮助读者既能在工程中正确调用aclnnReal也能从源码层面理解其性能优化设计。产品支持情况产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√需要说明的是Real 算子的 AI CoreAICore实现目前仅在 ascend910b 上注册这一点在算子定义 real_def.cpp 的this-AICore().AddConfig(ascend910b, aicoreConfig)中有直接体现而 COMPLEX128 类型走 AICPU 实现适用于更多平台详见后文数据类型映射。功能说明算子功能提取复数张量的实部real part。对于实数类型输入输出等于输入恒等操作。计算公式$$ output_i\text{real}(input_i)\begin{cases} \text{Re}(input_i), \text{if } input_i \text{ is complex} \ input_i, \text{if } input_i \text{ is real} \end{cases} $$与 PyTorch 对应torch.real()与 NumPy 对应np.real()从图模式的算子语义看real是一个逐元素elementwise算子其 Shape 推导直接复用基类提供的逐元素推导函数InferShape4Elewise见 real_infershape.cpp输出 shape 与输入保持一致。参数说明参数名输入/输出/属性描述数据类型数据格式input输入待提取实部的输入张量。COMPLEX32, COMPLEX64, COMPLEX128, FLOAT16, FLOATNDoutput输出提取出的实部张量。FLOAT16, FLOAT, DOUBLENDTout属性可选属性指定输出数据类型。默认值为 DT_FLOAT(float32)。Int-在算子定义 real_def.cpp 中input/output均声明为REQUIRED必选Tout属性声明为OPTIONAL可选且默认值ge::DT_FLOAT同时算子声明了DynamicCompileStaticFlag(true)、DynamicFormatFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)即同时支持动态 shape、动态 rank 与动态格式这也是它能广泛接入框架侧动态推导链路的基础。数据类型映射输入类型输出类型Tiling KeyTout值说明COMPLEX128DOUBLE--提取复数实部 (AICPU only)COMPLEX32FLOAT1611提取复数实部COMPLEX64FLOAT20提取复数实部FLOAT16FLOAT1641恒等操作FLOATFLOAT50恒等操作注COMPLEX128 类型仅支持 AICPU 实现ascend910b 的 AICore 不支持该类型。Tout 值为数据类型枚举值DT_FLOAT0, DT_FLOAT161, DT_DOUBLE11可通过 Tout 属性可选指定输出类型。Tiling Key 的语义在源码中也有对应定义RealTilingKey枚举real_tiling.h中TILINGKEY_COMPLEX321、TILINGKEY_COMPLEX642、TILINGKEY_COMPLEX1283、TILINGKEY_FLOAT164、TILINGKEY_FLOAT5Host 侧在 real_tiling.cpp 中根据输入 dtype 通过 switch 语句为每种输入类型设置对应的 tilingKeyKernel 侧再依据 tilingKey 模板实例化不同的输入/输出类型组合。约束说明输入张量的 shape 必须与输出张量的 shape 相同。支持动态 shape 和动态 rank。输入数据不能为空在 Tiling 阶段若totalLength 0会直接报错返回见 real_tiling.cpp。在 ACLNN 接口层aclnn_real.cpp 还会额外校验所有参与张量的维度不超过ACLNN_MAX_SHAPE_RANK8 维、self与out的 shape 必须一致、输入输出 dtype 组合必须落在DTYPE_SUPPORT_LIST白名单内如 COMPLEX64→FLOAT、COMPLEX32→FLOAT16、COMPLEX128→DOUBLE、FLOAT16→FLOAT16、FLOAT→FLOAT。实现说明目录结构Real 算子遵循 CANN 自定义算子开发的典型分层结构仓库中的实际布局如下experimental/math/real/ ├── op_host/ # Host侧实现 │ ├── real_def.cpp # 算子定义 │ ├── real_infershape.cpp # Shape推导 │ └── real_tiling.cpp # Tiling计算实现 ├── op_kernel/ # Kernel侧实现 │ ├── real.cpp # Kernel入口 │ ├── real_kernel.h # Kernel模板类实现 │ └── real_tiling.h # Tiling数据结构定义 ├── op_api/ # API接口 │ ├── aclnn_real.cpp # ACLNN接口实现 │ ├── aclnn_real.h # ACLNN接口声明 │ ├── real.cpp # 算子API实现l0op::Real │ └── real.h # 算子API声明 ├── docs/ │ └── aclnnReal.md # API文档 ├── examples/ │ └── test_aclnn_real.cpp # ACLNN调用示例 └── tests/ut/ # 单元测试 ├── op_api/ │ └── test_aclnn_real.cpp ├── op_host/ │ ├── CMakeLists.txt │ └── test_real_tiling.cpp └── op_kernel/ ├── CMakeLists.txt ├── real_tiling.h ├── test_real.cpp └── real_data/ ├── gen_data.py # 测试数据生成 └── compare_data.py # 结果比对Tiling 参数说明Tiling 阶段产出的RealTilingData结构体定义在 real_tiling.h核心字段含义如下totalUsedCoreNum: 实际使用的总核数tailBlockNum: 大核数量余数 block 数ubPartDataNum: 每次 UB 循环处理的元素数smallCoreDataNum: 小核数据量元素数smallCoreLoopNum: 小核 UB 循环次数smallCoreTailDataNum: 小核最后一次循环的元素数bigCoreDataNum: 大核数据量元素数bigCoreLoopNum: 大核 UB 循环次数bigCoreTailDataNum: 大核最后一次循环的元素数tilingKey: 算子类型标识1complex32, 2complex64, 4float16, 5floatuseNonInplace: 是否使用非 inplace GatherMask 路径0inplace, 1非 inplaceTiling 计算入口 real_tiling.cpp 会从输入 shape 取totalLength校验输入输出 dtype 合法性然后调用CalcRealTilingParam完成核心参数计算平台信息AIV 核数、UB 大小通过platform_ascendc::PlatformAscendC获取编译期也可通过RealCompileInfo含totalCoreNum30、ubSizePlatForm字段在TilingPrepare4Real阶段提前预取。最终在PostTiling中写入 Tiling Data并调用context_-SetBlockDim(totalUsedCoreNum)与context_-SetTilingKey(tilingKey)决定 Kernel 的并发核数与执行分支同时还会预留一块固定大小的 userWorkspaceRESERVED_WORKSPACE16MB。多核处理策略大小核分核分核逻辑参考 Exp 等逐元素算子的通用做法核心思路是把数据按 block 粒度均摊到多个 AI Core 上并通过大核 1 block的方式吸收余数保证各核负载差不超过 1 个 block对齐粒度Complex 类型128B 对齐满足 GatherMask inplace 的 256B 源数据约束Real 类型32B 对齐分核策略按输出数据类型字节数将总数据量对齐到对应 blocktotalBlocks Align(totalLength * dataTypeLength, alignSize) / alignSize若ubPartDataNum totalLength使用 1 核否则coreNum min(totalCoreNum, totalBlocks)everyCoreBlockNum totalBlocks / coreNumtailBlockNum totalBlocks % coreNum大小核分配前tailBlockNum个核为大核数据量 (everyCoreBlockNum 1) * alignSize / dataTypeLength其余核为小核数据量 everyCoreBlockNum * alignSize / dataTypeLength大小核负载差 ≤ 1 个 block偏移计算见 real_kernel.h 中Init的实现大核globalOffset blockIdx * bigCoreDataNum小核globalOffset blockIdx * bigCoreDataNum - (bigCoreDataNum - smallCoreDataNum) * (blockIdx - tailBlockNum)即先按全大核假设计算偏移再回退小核与大核的数据量差Complex 类型输入inputOffset globalOffset * 2每个元素占 2 个 output 元素空间源码中以uint64_t计算防止大张量下 offset 溢出Complex 双路径策略GatherMask inplace 要求count * 2 * sizeof(T) % 256 0tiling 据此选择路径对应useNonInplace标志Inplace 路径useNonInplace0适用多核场景 / 单核且 totalLength 满足 256B 对齐UB 分配inQueue(2x) × 2缓冲 4倍系数GatherMask inplaceGatherMask(src, src, mode1, ...)配合 pipeline prefetch 优化在循环中提前下发下一 tile 的 DMA-in与当前 tile 的 GatherMask 计算重叠见ProcessComplexTiling非 Inplace 路径useNonInplace1适用单核且 totalLength 不满足 256B 对齐如 complex32[4,4]16元素16×2×264 256UB 分配inQueue(2x) outQueue(1x) × 2缓冲 6倍系数GatherMask 非 inplaceGatherMask(dst, src, mode1, maskcount*2, repeatTimes1)从 Kernel 侧实现看RealKernelS, T模板类real_kernel.h的核心思路非常巧妙复数在内存中按实部、虚部交错存储因此提取实部等价于从2N个元素中每隔一个取出一个——这正是向量指令 GatherMask 的典型应用场景ExtractRealPartinplace通过params.repeatTimes count * 2 * sizeof(T) / 256一次处理 256B 对齐的数据块ExtractRealPartNonInplace通过mask count * 2指定抽取的元素个数将结果写入独立的 outQueue。而实数输入的恒等分支ProcessRealIdentityreal_kernel.h则使用TQueBind绑定输入输出队列做纯拷贝copy同样带 prefetch 流水优化。Kernel 入口 real.cpp 按 tilingKey 实例化四组类型组合COMPLEX32_MODEkey1RealKernelint32_t, half即输入按 32bit 复数存储、输出 halfFLOAT16COMPLEX64_MODEkey2RealKernelint64_t, floatFLOAT16_MODEkey4RealKernelhalf, half恒等拷贝FLOAT_MODEkey5RealKernelfloat, float恒等拷贝模板参数S, T中S为输入存储类型、T为输出计算类型if constexpr (IsSameTypeS, T::value)在编译期就区分了恒等拷贝与复数抽取两条路径零运行时开销。调用说明ACLNN API 调用Real 算子对外提供标准的 CANN ACLNN 两段式接口声明见 aclnn_real.h#include aclnnop/aclnn_real.h // 1. 获取workspace大小 aclnnStatus aclnnRealGetWorkspaceSize(const aclTensor* self, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor); // 2. 执行算子 aclnnStatus aclnnReal(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream);注意ACLNN API 使用self和out作为参数名与图模式的input/output不同。第一段接口aclnnRealGetWorkspaceSize内部完成了四类参数校验空指针、dtype 支持范围、输入输出 dtype 匹配、shape 一致性见 aclnn_real.cpp并构造出完整的计算图self → l0op::Contiguous非连续转连续→ l0op::Real → l0op::ViewCopy结果写回 out。其中有两个值得注意的优化分支FLOAT / FLOAT16 输入直接走l0op::ViewCopy透传拷贝不再进入 Real 计算aclnn_real.cpp与实数输入恒等操作的语义完全一致self为空 Tensor 时直接返回workspaceSize 0跳过后续计算。L0 层入口l0op::Realreal.cpp负责推导输出 dtypeCOMPLEX64→FLOAT、COMPLEX32→FLOAT16、COMPLEX128→DOUBLE并依据当前 SoC 版本选择执行后端ASCEND910B / ASCEND910_93 上支持 AICore 加速支持 FLOAT、FLOAT16、COMPLEX32、COMPLEX64其余平台或 COMPLEX128 类型则走 AICPU 实现。完整调用示例以下示例来自仓库 examples/test_aclnn_real.cpp完整演示了初始化 → 构造 Tensor → 两段式调用 → 同步取结果 → 释放资源的标准流程#include aclnnop/aclnn_real.h #include acl/acl.h #include iostream #include vector #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, return ret); return 0; } int main() { // 1. device/stream初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 2. 构造输入与输出FLOAT 输入shape [4,2]恒等输出 std::vectorint64_t selfShape {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {1, -1, -1, -2, 2, -2, -3, 3}; std::vectorfloat outHostData {1, -1, -1, -2, 2, -2, -3, 3}; // 此处通过 CreateAclTensor 辅助函数完成 aclrtMalloc aclrtMemcpy aclCreateTensor // 完整代码见 examples/test_aclnn_real.cpp // 3. 两段式调用先获取 workspace 大小再执行 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnRealGetWorkspaceSize(self, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, return ret); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, return ret); } ret aclnnReal(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 4. 同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, return ret); // 5. 将 device 侧结果拷回 host 侧并打印 // PrintOutResult(outShape, outDeviceAddr); // 6. 释放 aclTensor 与 device 资源 aclDestroyTensor(self); aclDestroyTensor(out); aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例中selfShape {4, 2}、selfHostData {1, -1, -1, -2, 2, -2, -3, 3}由于输入为 FLOAT 类型走恒等路径后输出与输入完全一致若将输入改为 COMPLEX64则需要把实部、虚部交错排列输出将只保留实部数据偶数下标元素。API 的详细参数表与输入输出类型对应关系可进一步查阅 docs/aclnnReal.md。约束补充ACLNN 接口层self 与 out 的 shape 必须相同。支持动态 shape 和动态 rank。输入数据不能为空。确定性计算aclnnReal 默认确定性实现。测试与验证Real 算子提供了三层单元测试见 tests/ut 目录op_kernel 测试test_real.cpp配合real_data/gen_data.py与real_data/compare_data.py生成输入数据并比对 Kernel 输出覆盖复数抽取与实数恒等两类场景op_host 测试test_real_tiling.cpp校验不同 shape / dtype 组合下 Tiling 参数的合理性重点覆盖单核/多核、大小核分核与 256B 对齐约束的边界条件op_api 测试test_aclnn_real.cpp验证两段式 ACLNN 接口在 Device 上的实际执行结果。这类数据生成脚本 Host Tiling 单测 Kernel 仿真 API 端到端的组合正是 CANN 算子开发中推荐的完整验证链路。总结Real 算子是理解 CANN ops-math 库单目逐元素算子工程范式的极佳样例语义上它只做一件简单的事提取复数实部 / 实数恒等但工程实现上融合了动态 shape 推导、基于 dtype 的 tilingKey 分派、大小核负载均衡的多核分核算法以及巧用 GatherMask 向量指令在 inplace / 非 inplace 两条路径之间自适应切换的 Kernel 设计。无论是希望快速在工程中接入复数实部计算还是想借鉴其 Tiling 与 Kernel 优化思路来开发自己的算子都可以从 experimental/math/real 这份实现与文档中找到完整参考。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 数学算子实战Tan 自定义算子的原理、Tiling 与 aclnn 调用全解析CANN ops math 数学算子实战Tan 自定义算子的原理、Tiling 与 aclnn 调用全解析 Tan 算子是 CANN ops math 数学算算子库人工智能CANNCANN ops-math IsClose 算子深度解析原理、aclnn API 调用与 AscendC 实现CANN ops math IsClose 算子深度解析原理、aclnn API 调用与 AscendC 实现 IsClose 是 CANN ops math算子库人工智能CANNCANN ops-nn 算子解析HingeLossGrad 反向梯度算子原理、ACLNN 调用与多核 Tiling 实现CANN ops nn 算子解析HingeLossGrad 反向梯度算子原理、ACLNN 调用与多核 Tiling 实现 导读 Hinge Loss合页损失人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考