编者注

为了推动不同架构 AI 硬件系统的创新和规模化落地，智源研究院联合多家机构打造开源、统一的 AI 系统软件生态 FlagOS。系统软件栈 FlagOS 包括统一 AI 编译器 FlagTree、高性能通用 AI 算子库 FlagGems/FlagAttention、大模型训推一体框架 FlagScale 和统一通信库 FlagCX 等关键技术。目前，FlagTree 项目已于3月份对外开源，社区希望以开放合作的方式，打造一个面向多种硬件架构的开源、统一 AI 编译器，扩展 Triton 语言的上下游生态，赋能大模型技术开发者。

本文作者李嘉楠，来自先进编译实验室，也是 FlagTree 项目的核心开发贡献者。通过文章为大家呈现面向 GPGPU 硬件架构时，FlagTree 编译优化的技术实践方案。

---

正文

FlagTree 整个项目由两部分组成：Triton语言与编译器。首先，从语言角度看，Triton语言提供了线程块级别的编程模型，开发者可以使用Python风格的代码，直接编写并行计算逻辑。可以灵活指定张量的划分方式、访存顺序等，而底层的调度、同步、线程控制逻辑由语言自动封装。但要实现“写得简洁，跑得高效”这一目标，语言只是入口，核心还在 Triton 编译器层。

面向多种芯片的统一 AI 编译器 FlagTree 负责把线程块以上的张量操作自动转换为线程块以下的高性能调度策略。它具备以下几个关键能力：自动进行共享内存管理、支持张量布局变换等深度优化。这一整套自动优化流水线，让开发者“专注业务逻辑，不必过分关注性能细节”。如下图所示，统一 AI 编译器 FlagTree 支持多种国产芯片平台，采用统一的中间表示和优化机制，打通从高层 AI 编程到传统编译链路，构建高效可移植的 AI 编程体系。

 

01 优化遍 Pass 的结构与添加

AI 编译器中的核心是优化遍 Pass，Pass 指对程序的中间表示（IR，Intermediate Representation）进行分析或转换的一个独立阶段。目标是优化或转换程序，使代码能更简洁高效地在目标硬件上正确运行。整体分为两类，一类是分析类 Pass，不修改程序，只做信息收；一类是降级优化类 Pass，通过修改或重写 IR 优化代码。

Pass 的组织方式即为 Pass 管理器，它确定了 Pass 的执行顺序、维护 Pass 之间的依赖关系，并确保每个 Pass 的输出可以作为下一个 Pass 的输入。不同编译器 Pass 管理方式不一样，llvm 编译器中间优化遍统一在一个 pass manager 文件，有固定的优化顺序，主要负责中间代码的优化，本身并不直接负责高层 IR 的降级；tvm 编译器优化顺序不固定，因此各种优化顺序排布构成解空间，因此带来最优解的问题；统一 AI 编译器 FlagTree 的优化遍既包含优化也包含降级，ttir-ttgir-llir 模块清晰，其中 ttir、llir 模块主要负责降级，面向 GPGPU 的优化往往集中在 ttgir 这一部分。Pass管理器源码位置

flagtree /third_party/xxx/backend/compiler.py 

那么如何添加优化遍呢？主要分为以下几个步骤：

(1) 在 Passes.td 中定义Pass的元数据(名称、描述、构造方法等)：

flagtree /include/triton/Dialect/TritonGPU/Transforms/Passes.td

 

 

低版本例如 v2.1 需要在 Passes.h 中定义Pass的接口声明 flagtree/include/triton/Dialect/TritonGPU/Transforms/Passes.h

 

(2) 创建一个优化遍C++文件来实现 Pass 的逻辑，例如 flagtree/lib/Dialect/TritonGPU/Transforms/目录下创建一个PrintAddIOp.cpp

 

(3) 修改 CMakeLists.txt，确保新的 Pass 文件被添加到构建系统中，路径flagtree/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

 

 

(4) 修改triton.cc，用于将 C++ 实现的 Pass 和其他功能暴露给 Python，路径flagtree/python/src/ passes.cc （triton.cc）

 

 

(5) 修改compile.py，将 Pass 加入到实际编译流程中路径是flagtree/python/triton/compile/compiler.py或triton/third_party/xxx/backend/compiler.py

 

 

02 配置策略优化

我们在做 FlagTree 编译优化时，首先要考虑配置优化策略，在面向 GPGPU 进行优化遍解析之前，Kernel 调优中配置参数也会影响整体的执行效率：

• kwargs：用于给内核传递关键字参数，例如可以设置BLOCK_SIZE用于确定块大小

• num_warps：设置编译时内核使用的线程数

• num_stages：设置流水线循环的阶段数，多用于矩阵乘法

• maxnreg：设置单个线程可以使用的最大寄存器数

常用的参数配置策略有三种：

• 运行时传参：用户在程序运行时传递自定义配置参数

• @triton.autotune：使用自动调优方法从给定Config列表中选择kernel的最优配置参数

• @triton.heuristics：使用启发式算法计算Kernel配置参数

其中 @triton.autotune 允许开发者为Kernel中BLOCK_SIZE_M/N/K等关键参数设定候选值或范围，每一组参数配置都会编译生成可执行文件，并通过测试选取执行时间最短的配置作为最佳参数配置。另外我们可以通过添加调试选项TRITON_PRINT_AUTOTUNING=1打印自动调优所花费的时间和最佳配置，cache文件中每种配置都会生成编译后的相关中间文件。

 

 

 

 

03 访存合并优化

访存合并是针对全局内存访问的内存优化技术，通过对线程束同一时刻发出的多条访存请求进行合并，使得处理器可以在最少的访存周期内完成多个数据的存取工作。当 GPU 需要访问内存数据时，会以缓存行为单位进行数据的读取和存储，常见大小是64或128字节，当程序读取内存时，不是一字节一字节读，而是按“缓存行”整体搬一块数据进缓存。内存事务是 GPU 实际向全局内存发起的一次数据传输操作，每个事务会搬运一个或多个缓存行，事务数量越少，访存越高效。

编译器中的访存合并优化遍通过对块布局信息进行优化实现核函数的合并访存和向量化访存：计算所有op轴信息，为计算新的layout做准备；计算访存相关op经coalesce后的layout；对访存相关op及前后进行layout改写。

在ttir到ttgir的降级过程中，首先会添加kernel的layout信息，默认是sizeperthread=1，threadperwarp=32，warppercta=warp数， 轴分析会计算出各个op的divisibility、contiguity、constancy、constvalue四个属性值。访存合并优化遍会遍历中间表示中所有的访存操作，并综合地址对齐、访存连续、硬件支持的向量化访存指令位宽大小等因素计算出最佳的sizePerThread属性值等block相关信息。针对编译器无法高效合并访存的三种场景，设计实现基于访存合并的全局内存优化方法：

 

• 轴连续性优化通过添加额外轴属性信息并优化轴属性值计算公式，保证连续张量经四则运算后连续性的正确推导；

• 通过对TritonGPU中间表示的线程访存布局计算方式进行重构，实现算子的高效访存；

• 通过在抽象语法树向Triton中间表示降级时添加多种轴属性值种类，实现了算子连续访存下多个维度张量大小的向量化访存。

 

04 特定算子优化

以上优化是基于现有优化的不足实现完善，我们也可以基于特定算子的缺陷以及性能瓶颈出发，实现特定算子的编译优化。例如在 Triton GPU方言转换为 LLVM 方言的过程中，Triton 编译器提供了模块化的转换框架，Triton 支持多种关键算子的代码生成，包括归约算子、布局转换算子、逐元素操作算子以及加载和存储算子等。归约算子是通过对张量沿特定维度进行求和、求最大值最小值等操作，将多维数据压缩为低维或标量形式，从而简化后续的处理和分析。归约算子GPU并行方案有树形并行归约、Warp Shuffle、块并行归约等。

归约算子根据归约轴维度的特性可分为连续轴归约和非连续轴归约。编译器中非连续轴归约算子从TritonGPU IR到LLVM IR的降级采用的是树形归约算法，一些shape下存在较大的内存访问开销，因此可以通过增加实现方法以提升整体计算效率：添加Warp Shuffle归约方式的同时，采用数据打包。

05 张量核心优化

除了以上两种优化角度，我们也可以从硬件特征出发发挥特定硬件的计算能力。张量核心（Tensor Core）是现代GPU中专为深度学习设计的专用计算单元，能够以极高吞吐量执行矩阵乘法与累加（MMA）操作。张量核心支持INT8、FP4、FP6、FP8、FP16、BF16、TF32等多种数据精度格式，在执行小规模矩阵计算（如16×16）时可显著提升算力与能效。   在深度学习训练与推理中，张量核心广泛用于卷积与全连接层的高效计算。张量核心的计算瓶颈确实往往来自于访存速度不足，而不是算力本身。因此我们采用以下优化方法实现数据的快速准备：

• 分块配置搜索策略：在矩阵乘法执行前，通过分析张量核心、操作数形状、原子MMA指令和共享内存的特性信息，构造多种分块配置，共同组成分块配置搜索空间，构造搜索算法替代手动枚举的分块配置。

• 定义使用链拆分以及循环分块优化：定义使用链拆分可以调整访存和计算指令的顺序，提升访存和计算的并行程度。循环分块将一个分块计算过程拆分为多个子过程，每个子过程独立申请线程资源进行计算，以提升线程资源的利用率。

   

 

06 代数变化优化

还有一种优化是通用类优化，在基础编译中、AI编译中、手工优化中都会存在，但因为作用阶段不一样优化效果也不一样。这里我们给出一个代数变换优化，代数变换是对程序进行各种数学等价的变换，使得从变换后的程序出发，能生成更高效的目标代码。常用变换：简化表达式、重排运算顺序、运算类型转换、交换律与结合律等。  

以Bert模型经过Pytorch Inductor后端生成的triton kernel 为例，其中就存在代数变换优化的机会。在编译器中实现这样的通用优化，除了之前添加优化遍的整个过程，更是要在实现中增加以下三个功能：

 

 

首先根据定义使用链定位修改节点的位置，在相应的位置上添加新范式的节点，最后删除旧的节点更新整体的定义使用关系以维护整体计算逻辑。

以上就是分享的整体内容，包括编译的结构分析、优化遍的添加，另外从编译优化缺陷、算子性能瓶颈、硬件特征以及通用优化四个角度分析实现优化的思路。