鲲鹏编译器AI场景竞争力

随着搜推业务的持续发展当前推业务中大约50%的模型特别是中小模型仍运行在CPU上需要鲲鹏编译器优化鲲鹏上AI模型的推理性能支撑优化推理性能、增强鲲鹏平台竞争力。随着AI推理性能优化需求的不断增长传统的“手工编写算子通用内存管理”模式已难以为继。首先面对日新月异的应用场景底层算子的手工调优速度远跟不上模型迭代步伐其次面向TensorFlow的XLA图编译器虽具备通用图优化能力但仅支持静态shape。在生产环境中数据动态变化如变长特征会导致频繁的重编译引发严重延迟毛刺峰值延迟可达正常值的100倍严重威胁服务稳定性。最后AI 任务极高的并发吞吐对系统的即时内存分配提出了严苛要求通用内存库在多核架构下的锁竞争和碎片问题已成为制约硬件性能释放的最后一公里。本方案通过鲲鹏编译器与运行时软硬协同双管齐下打造通用加速底座聚焦 AI 编译优化与内存加速两大方向在 AI 编译优化方面首先基于 MLIR 构建面向稀疏计算的专用方言 ATIR通过启发式规则自动识别由 Gather、SegmentReduce、Unique 等算子构成的稀疏 Embedding 子图并将其融合为单一、高效的定制化算子同时支持用户自定义融合模式以满足业务快速迭代需求实现鲲鹏亲和的高性能稀疏算子自动化构建与优化其次面向 TensorFlow XLA 提出动态 shape 性能优化技术通过异步编译、自动填充与运行时 shape 动态替换等手段从根源上解决生产环境中由数据动态变化如变长特征引发的频繁重编译和延迟毛刺问题并引入并行图切分与多内核动态选择机制使 XLA 图编译器的全局优化能力得以在鲲鹏 CPU 上充分发挥。在内存加速引入专为鲲鹏设计的 KQMalloc 内存库通过六级无锁缓存逐级回退、虚拟地址分段设计与 Tag 标签系统实现极致的无加锁并发分配与低碎片内存回收全面构筑鲲鹏 AI 场景的性能优势。整体架构针对鲲鹏搜推业务中 Embedding 计算占比高、稀疏访存瓶颈突出、动态 shape 引发频繁重编译等痛点问题本系统面向鲲鹏 CPU 架构构建了一套 AI 推理优化技术体系聚焦搜索、推荐、广告等稀疏计算密集型场景从图编译优化与运行时内存管理两个维度协同发力。在图编译层面系统采用分层优化架构前端基于 MLIR 构建面向稀疏计算的专用 ATIR 方言在兼容XLA核心算子的基础上扩展SparseFillEmptyRows、SegmentSum、DynamicPartition等稀疏原语实现Gather、SegmentReduce、Unique等Embedding相关子图的自动识别与深度图融合并支持用户自定义融合模式融合后的算子针对鲲鹏 CPU 微架构特性自动生成高度优化的执行代码显著减少中间张量生成、降低全局内存访问次数并提升缓存与带宽利用率。中端针对动态 shape如动态 batch导致的频繁重编译与时延毛刺问题提出无损优化方案通过 TF 侧符号化类型推导对同来源动态维度进行聚类降低子图复杂度。基于规则将动态 shape 自动填充为固定值避免反复重编译同时在 XLA 中保留动态 shape 相对初始 shape 的逻辑关系运行时通过 SessionState 获取实际 shape 进行动态替换从而消除 padding 带来的冗余计算将时延毛刺峰值从千毫秒级降低到十毫秒级。后端引入并行图切分与多内核动态选择机制并行图切分通过识别矩阵乘法、softmax 等高计算开销的可并行操作并扩展融合集群将计算图切分为多个可并行执行的融合集群在保证融合后吞吐提升的同时确保时延不劣化多内核选择器综合 Eigen、OpenBLAS、LibShalom 及面向鲲鹏优化的 KDNN 等算子库根据实际工作负载特征动态匹配最优算子实现充分释放 CPU 加速能力。在内存分配层面系统基于 KQMalloc 内存库采用六级无锁缓存逐级回退架构结合虚拟地址分段编码与 Tag 标签系统通过 Segment-Heap 位运算布局直接从指针地址推导堆索引、偏移及 NUMA 节点信息释放时通过预先写入的 1 字节 Tag 以 O(1) 复杂度确定对象元信息实现高并发场景下的极致低延迟内存分配与释放。上述技术已在典型搜推模型及生产环境中完成验证为鲲鹏 AI 推理提供了从编译优化到运行时内存管理的系统性性能加速能力。图1 鲲鹏AI推理编译加速底座鲲鹏图编译器灵活图算融合在搜索与推荐系统中Embedding 相关计算如 Gather、SegmentReduce因高频调用和随机内存访问通常占端到端推理耗时的 20%–95%在生成式搜推中仍高达约 50%。这类操作具有典型的稀疏性计算密度低、访存开销大、内存访问不规则。现有优化依赖人工编写融合算子泛化性差、维护成本高而通用 AI 编译器面向稠密计算设计难以有效识别和优化稀疏 Embedding 子图导致性能瓶颈长期存在。图2 待融合Embedding算子实例为系统性解决稀疏 Embedding 计算的性能瓶颈我们提出一套面向鲲鹏计算架构深度优化的 Embedding 自动图融合方案。该方案聚焦搜推场景中典型的稀疏访存模式如 ID 查表、分段聚合等通过在编译期自动识别由 Gather、SegmentReduce、Unique 等算子构成的稀疏子图并将其融合为单一、高效的定制化算子从而显著减少中间张量生成、降低全局内存访问次数并提升缓存与带宽利用率。整个流程无需人工参与一方面系统内置一套基于图拓扑结构和算子语义的匹配引擎可自动识别预定义的稀疏融合模板另一方面为满足业务快速迭代和模型定制化需求方案支持用户自定义融合模式——开发者可声明新的融合规则包括输入/输出语义、索引依赖关系及融合后的计算逻辑。这些自定义规则将被统一纳入融合调度框架在编译时与内置模板协同参与子图匹配与优化决策。融合后的算子结合鲲鹏处理器的微架构特性自动生成高度优化的执行代码。图3 融合算子示例关键技术一基于 MLIR 构建专用方言 ATIR实现稀疏计算领域原语表达在 MLIR 框架下我们设计并实现了面向推荐与搜索场景的专用方言 ATIR其核心目标是精准刻画稀疏计算的语义特征并为后续优化提供结构化基础。ATIR 在兼容 XLA 核心算子如 Dot、Broadcast、Reduce的同时重点扩展了搜推模型中高频出现的稀疏原语包括 SparseFillEmptyRows、SegmentSum、DynamicPartition 等并为其定义了严格的类型系统与操作语义。ATIR 作为中间表示层既可由高层框架通过适配器转换生成又能通过 MLIR 的多级 IR 能力向下逐步降级至 CPU 高性能内核。该方言的引入使得稀疏逻辑从“黑盒组合”变为“可分析、可变换、可融合”的结构化表达为自动图融合奠定了语义基础。关键技术二基于启发式规则的自动图融合机制实现端到端性能提升针对 ATIR 中表达的稀疏计算图我们构建了一套基于启发式规则的自动图融合引擎专门用于识别 ModelZoo 中典型模型的关键计算模式。该引擎通过预定义的结构化匹配模板结合算子语义、数据依赖和维度约束在编译期自动定位可融合子图。匹配成功后系统将其替换为一个语义等价但高度优化的手写融合算子该算子针对鲲鹏 CPU 的微架构特性进行深度调优显著减少中间张量分配、合并多次索引遍历、提升缓存命中率。鲲鹏图编译器动态Shape优化XLA的图优化主要面向GPU/NPU构建因底层架构差异无法在CPU上充分发挥效果。因此亟需面向鲲鹏CPU研发适配动态shape的图编译优化技术从根本上消除重编译毛刺进一步发挥XLA图优化能力提供面向CPU架构的图变换、算子融合等优化能力确保模型在CPU上实现稳定、低延迟、高吞吐的推理性能。我们提出的面向TFXLA的动态shape性能优化技术旨在从根本上解决生产环境中因数据动态变化导致的频繁重编译与时延毛刺问题。核心包含三大能力异步编译技术隐藏首次编译开销自动填充将动态shape规整为固定值避免反复重编译运行时shape动态替换恢复实际shape消除填充带来的冗余计算。在此基础上围绕鲲鹏CPU架构特性从三个维度协同优化实现低时延、高吞吐推理1) 针对CPU算力有限、内存访问更快的特点优化图融合策略简化执行路径减少冗余计算2利用CPU上下文切换和数据移动开销低的优势挖掘算子间并行提升并行效率降低时延3)综合Eigen、OpenBLAS、LibShalom及面向鲲鹏优化的KDNN等算子库充分发挥CPU加速能力释放硬件性能。通过TFXLA的动态shape性能优化技术可解决重编译问题预期将时延毛刺峰值从千毫秒级降低到十毫秒级达到生产环境时延要求。CPU侧图优化相关技术在modelzoo benchmark上进行了验证吞吐性能平均提升10%。关键技术一动态shape性能无损优化。TF侧符号化类型推导分析不同来源的动态shape同来源的动态维度标记相同值基于动态维度信息聚类子图降低子图动态shape复杂度。基于指定规则填充动态维度若缓存中存在该对应填充值副本可复用。在XLA子图中将实际动态维度信息保存在SessionState上下文中。XLA中执行类型推导保留动态shape相对其初始动态shape的逻辑关系运行时从SessionState中获取实际动态shape信息并基于逻辑关系进行替换消除padding产生的冗余计算。图4 融合算子示例关键技术二并行图切分优化。搜推广场景时延是一个重要指标而并行化是减少推理延迟的关键优化。如果没有数据依赖可以并行执行多个算子以更快完成。如下图所示时间轴显示了5个矩阵乘法运算符并行执行因此完成它们的延迟接近于只执行一个从而获得了5倍的延迟降低。这样的并行化方案被称为算子间并行即多个算子并行执行。图5 Tensorflow原生执行方式下图左侧展示使能XLA图编译器后所有的matmul操作及其他操作被融合在一起后的执行情况在这种情况下所有矩阵乘法matmul运算符都被融合为单个大运算符标记为XlaRun。但是由于被融合的算子被视为单个操作其内部计算不能被并行调度而是顺序执行多个操作导致实际并行程度下降。我们通过将选定的可并行操作分组到同一个融合集群来优化图切分策略。该过程从识别计算开销较大的可并行操作开始如矩阵乘法和softmax。然后它通过合并额外的可融合操作来扩展融合集群直到遇到另一个计算开销大的操作或者依赖关系上不可并行的操作。通过这种方式将计算图切分为多个可并行的融合集群每个集群都围绕着单个高计算瓶颈的操作及其相关可融合操作。这些集群被编译成可以并行执行融合算子。如下图右侧所示该优化保证使能融合后吞吐提升的情况下时延不劣化。图6 动态优化效果关键技术三多内核选择优化。我们还引入了一个内核选择器它可以在不同的算子实现中进行动态选择从而利用CPU丰富的软件生态系统最大限度地提高性能。许多运算符有多个CPU实现每个都针对特定的工作负载或CPU变体进行了优化。例如Eigen项目提供了一系列广泛的数学运算符包括矩阵乘法并对通用CPU进行了合理的优化。相比之下OpenBLAS提供了专门针对选定CPU架构调整的矩阵乘法实现而LibShalom专注于优化较小的矩阵大小。由于推荐模型中的推理工作负载千差万别很少有一种解决方案适合所有场景。一个工作负载可能使用OpenBLAS的矩阵乘法获得更好的性能而另一个工作负载可能从LibShalom中获益更多。内核选择器将一系列可用库和知识库相结合能够选出最匹配目标工作负载特征的算子实现从而解决了这一挑战。通过内核选择器在modelzoo benchmark和实际模型负载上取得显著提升。鲲鹏内存库分配优化六级无锁缓存架构KQMalloc 的分配与释放路径采用六级缓存逐级回退的设计。核心思想是越靠近线程的缓存锁粒度越细、访问越快越远离线程的缓存容量越大但开销也越大。最常见的小对象分配路径仅访问第一级线程本地缓存全程无锁约10 条指令即可完成一次分配。每一级缓存都针对不同的并发场景做了针对性优化第一级完全无锁线程私有第二级使用轻量 trylock32 路随机选择降低争用第三级使用 CAS 原子交换完全无锁第四级使用单字节锁的 trylock第五级使用原子交换获取 cell 索引第六级使用自旋锁加退避策略。这种逐级回退的设计确保了在高并发场景下大部分请求在前几级就能完成避免了全局锁争用。图7 融合算子示例虚拟地址分段设计KQMalloc 将进程的虚拟地址空间划分为固定大小的段Segment每个段包含 512 个堆Heap每个堆占 128GB2 的 37 次方虚拟地址空间。通过精心设计的位运算布局从任意一个已分配的指针地址可以直接推导出其所属的堆索引、堆内偏移和 NUMA 节点 ID整个过程仅需 2-3 条位运算指令无需任何查表。这种设计的核心优势在于释放路径不需要遍历堆数组来查找指针属于哪个堆也不需要维护复杂的地址映射表。所有信息都编码在指针地址本身之中通过位掩码和移位操作即可提取。图8 融合算子示例Tag标签系统在内存释放时分配器需要知道一个指针指向的对象属于哪个堆、对象大小是多少。传统方法需要遍历堆数组查找或维护复杂的地址映射表开销较大。KQMalloc为每个堆的每个cell预先写入一个 1 字节的tag标签释放时直接读取tag即可O(1) 确定对象的全部元信息。tag的含义随堆类型不同而变化在Slab堆和Block堆中tag.value存储的是sizeclass对象大小分类编号在Fast堆中tag.value存储的是cellnum连续占用的cell数量在 Span 堆中tag.value存储的是alignedsizeclass对齐大小分类在Clean堆中tag.value存储的是cell_num。这种写一次、读多次的模式非常高效tag在对象分配时写入在对象释放时读取。由于分配频率远低于释放频率对象被多次引用后释放tag 的读取次数远大于写入次数将计算开销前置到分配时是合理的。图9 融合算子示例上手教程鲲鹏编译器编译器图算融合加速表1 TensorFlow图融合接口表2 XLA图融合接口算子优化表3 冗余算子优化接口表4 矩阵算子优化接口表5 Softmax算子优化接口常量折叠优化注常量折叠和图优化同时开启时需先进行常量折叠优化。表6 常量折叠模型转换接口表7 常量折叠优化接口说明此部分供ANNC已有特性尝鲜使用更强优化能力敬请期待。鲲鹏内存库加载鲲鹏内存库部署KQMalloc支持多种不同的编译方式1.单独编译只会编译出适配鲲鹏某一种CPU型号的动态库KQMalloc支持HIP08、HIP09、HIP10A、HIP11和HIP12共计五种型号的鲲鹏CPU。# 默认使用gcc编译可以指定鲲鹏CPU型号make clean build targetHIP09若指定CPU型号为HIP09编译完成后则输出文件位于build_HIP09文件夹下。2.全量编译会编译出适配鲲鹏5种型号的动态库。# 默认使用clang编译make clean kunpeng# 同时支持使用gcc编译make clean kunpeng CCgcc编译完成后会在build目录下创建HIP08HIP09HIP10AHIP11和HIP12五个文件夹存放动态库。使用鲲鹏内存库KQMalloc通过优化标准库中的内存相关函数只需使用预加载方式即可生效。请根据所使用的鲲鹏处理器型号(HIP08, HIP09, HIP10A, HIP11, HIP12)选择对应的目录下的文件。其中libkqmalloc.so是该库的单线程版本libkqmallocmt.so是该库的多线程版本libkqmallocmt.debug.so是启用了断言、double free()检查等功能的调试版本libkqmallocmt.profiler.so是启用了内存分析功能的分析版本分析文件会保存在/tmp/kqmalloc进程名CPUcpuid.txt文件中。以鲲鹏HIP12的多线程场景为例export LD_PRELOAD/path to KQMalloc/build/HIP12/lib/libkqmallocmt.so针对鲲鹏搜推业务中Embedding计算占比高、稀疏访存瓶颈突出、动态shape引发频繁重编译等痛点问题本系统面向鲲鹏CPU架构构建了一套 AI 推理全栈优化技术体系聚焦搜索、推荐、广告等稀疏计算密集型场景从图编译优化与运行时内存管理两个维度协同发力加速鲲鹏推理业务性能20%同时鲲鹏图编译器支持动态shape从根本上解决TensorFlow软件栈在生产环境的时延毛刺问题。上述技术已在典型搜推模型及鲲鹏互联网客户的生产环境中完成验证为鲲鹏 AI 推理业务提供了从编译优化到运行时内存管理的系统性性能加速能力。