匠心时刻丨鲲鹏编译器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 融合算子示例

XLA的图优化主要面向GPU/NPU构建，因底层架构差异无法在CPU上充分发挥效果。因此，亟需面向鲲鹏CPU研发适配动态shape的图编译优化技术，从根本上消除重编译毛刺，进一步发挥XLA图优化能力，提供面向CPU架构的图变换、算子融合等优化能力，确保模型在CPU上实现稳定、低延迟、高吞吐的推理性能。

我们提出的面向TF+XLA的动态shape性能优化技术，旨在从根本上解决生产环境中因数据动态变化导致的频繁重编译与时延毛刺问题。核心包含三大能力：异步编译技术隐藏首次编译开销；自动填充将动态shape规整为固定值，避免反复重编译；运行时shape动态替换恢复实际shape，消除填充带来的冗余计算。

在此基础上，围绕鲲鹏CPU架构特性，从三个维度协同优化，实现低时延、高吞吐推理；1) 针对CPU算力有限、内存访问更快的特点，优化图融合策略，简化执行路径，减少冗余计算；2）利用CPU上下文切换和数据移动开销低的优势，挖掘算子间并行，提升并行效率，降低时延；3)综合Eigen、OpenBLAS、LibShalom及面向鲲鹏优化的KDNN等算子库，充分发挥CPU加速能力，释放硬件性能。

通过TF+XLA的动态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 条指令即可完成一次分配。

每一级缓存都针对不同的并发场景做了针对性优化：第一级完全无锁（线程私有），第二级使用轻量 trylock（32 路随机选择降低争用），第三级使用 CAS 原子交换（完全无锁），第四级使用单字节锁的 trylock，第五级使用原子交换获取 cell 索引，第六级使用自旋锁加退避策略。这种逐级回退的设计确保了在高并发场景下，大部分请求在前几级就能完成，避免了全局锁争用。

图7 融合算子示例

KQMalloc 将进程的虚拟地址空间划分为固定大小的段（Segment），每个段包含 512 个堆（Heap），每个堆占 128GB（2 的 37 次方）虚拟地址空间。通过精心设计的位运算布局，从任意一个已分配的指针地址可以直接推导出其所属的堆索引、堆内偏移和 NUMA 节点 ID，整个过程仅需 2-3 条位运算指令，无需任何查表。

这种设计的核心优势在于：释放路径不需要遍历堆数组来查找指针属于哪个堆，也不需要维护复杂的地址映射表。所有信息都编码在指针地址本身之中，通过位掩码和移位操作即可提取。

图8 融合算子示例

在内存释放时，分配器需要知道一个指针指向的对象属于哪个堆、对象大小是多少。传统方法需要遍历堆数组查找或维护复杂的地址映射表，开销较大。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 target=HIP09

若指定CPU型号为HIP09，编译完成后，则输出文件位于build_HIP09文件夹下。

2.全量编译：会编译出适配鲲鹏5种型号的动态库。

# 默认使用clang编译make clean kunpeng# 同时支持使用gcc编译make clean kunpeng CC=gcc

编译完成后，会在build目录下创建HIP08，HIP09，HIP10A，HIP11和HIP12五个文件夹存放动态库。

使用鲲鹏内存库

KQMalloc通过优化标准库中的内存相关函数，只需使用预加载方式即可生效。

请根据所使用的鲲鹏处理器型号(HIP08, HIP09, HIP10A, HIP11, HIP12)选择对应的目录下的文件。

其中libkqmalloc.so是该库的单线程版本；libkqmallocmt.so是该库的多线程版本；libkqmallocmt.debug.so是启用了断言、double free()检查等功能的调试版本；libkqmallocmt.profiler.so是启用了内存分析功能的分析版本，分析文件会保存在/tmp/kqmalloc<进程名>CPU<cpuid>.txt文件中。

以鲲鹏HIP12的多线程场景为例：

export LD_PRELOAD=/<path to KQMalloc>/build/HIP12/lib/libkqmallocmt.so