Blackwell GPU提供LLVM和MLIR支持的相关工作 报告
0. 信息来源
https://llvm.org/devmtg/2025-04/slides/technical_talk/ozen_blackwell.pdf
1. 一段话总结
在2025年4月于柏林举办的EuroLLVM会议上，Durgadoss Ramanathan等人介绍了为NVIDIA Blackwell GPU提供LLVM和MLIR支持的相关工作，包括Blackwell GPU的核心特性（如MMA从Hopper的Warp Group演进为CTA对、新增TMEM及专用指令、TMA新模式等）、编译器 lowering 流程（MLIR→LLVM→PTX）中针对Blackwell的扩展（新增约30个MLIR操作、1000个NVVM内在函数及TMEM的6号地址空间）、内在函数设计优化（通过打包标志和分离标志平衡可读性与数量膨胀问题）、APFloat对FP6/FP4等新类型的支持、NVVM Dialect的支持进展（除MMA外tcgen05操作已支持，14/90个操作待迁移），以及通过NVDSL Python实现128x128x64的GEMM示例，目前Hopper基础支持、tcgen05基础建模等已完成，剩余TMA、MMA等Blackwell更新仍在进行中。
2. 思维导图（mindmap）
## 会议背景
- 会议：EuroLLVM 2025
- 地点：柏林
- 时间：2025年4月15日
- 演讲者：Durgadoss Ramanathan、Guray Ozen、Pradeep Kumar
## Blackwell GPU核心特性
- MMA演进：Ampere(Warp)→Hopper(Warp Group)→Blackwell(CTA对)
- 新增Tensor Memory (TMEM)：per SM、专用分配/释放/加载/存储指令、累加器/操作数存储
- 支持Block-scaled types：暴露为“tcgen05”指令家族
- TMA增强：新模式（im2col_w/scatter/gather）、支持Masked Copy
## 编译器Lowering流程与Blackwell扩展
- 基础流程：MLIR→LLVM→PTX（含Python NVDSL示例）
- MLIR扩展：新增约30个针对Blackwell的操作
- NVVM Dialect扩展：新增约1000个内在函数、6号地址空间（用于TMEM）
- NVPTX Backend：负责将NVVM内在函数转换为PTX
## 内在函数设计与优化
- 内在函数数量：单TMA变体→数百个内在函数、tcgen05-MMA家族→约750个内在函数、pre-Blackwell总计约3000个（MMA/TMA占40%）
- 二进制大小影响：1K/2K/5K/10K内在函数下，总二进制大小分别增加492KB/980KB/2420KB/4852KB（占比0.32%/0.63%/1.56%/3.12%）
- 构建时间影响：10K内在函数时，FrontEnd IR Tablegen时间增66.77%、NVPTX Backend时间增4.2倍（整体构建时间无影响）
- 编译时间影响：10K内在函数时，geo mean为1.17（opt阶段无影响）
- 优化方案：打包标志（PR 96083 - v1，用位域编码修饰符）、分离标志（为不同修饰符设独立标志，提升可读性，比特码大小增最多3%）
## APFloat与NVVM Dialect支持
- APFloat支持：FP6(e2m3、e3m2，PR 94735)、FP4(e2m1，PR 95392)、解决e8m0类型“precision=0”问题（PR 107127）
- NVVM Dialect：除MMA外tcgen05操作已支持、支持inline-ptx与内在函数 lowering、14/90个操作待迁移至内在函数 lowering
## NVDSL Python GEMM示例
- 示例参数：128x128x64 GEMM，启动1个CTA，block=(128,1,1)
- 核心步骤：TMA fence与mbarrier初始化→TMEM分配→TMA加载→等待TMA加载→Tensor Core MMA→等待MMA→TMEM加载到寄存器→寄存器存储到GMEM→TMEM释放
## 后续计划与结论
- 后续计划：支持Imm内在函数参数（flags）美化打印（RFC: discourse 82629）、为内在函数附加地址空间特定RW属性、NVVM Ops的SM版本检查（PR 126886）、完成剩余指令/内在函数支持
- 已完成工作：Hopper基础支持（NVPTX与NVVM Dialect）、tcgen05家族基础建模、APFloat新类型支持
- 进行中工作：Blackwell的TMA、MMA等更新
- NVDSL价值：提供可运行示例、支持基于示例构建自定义编译器
3. 详细总结
一、会议背景
- 会议名称：EuroLLVM 2025
- 举办地点：德国柏林
- 举办时间：2025年4月15日
- 演讲者：Durgadoss Ramanathan、Guray Ozen、Pradeep Kumar
- 核心主题：为NVIDIA Blackwell GPU提供LLVM和MLIR支持的技术方案与进展
二、Blackwell GPU核心特性（AI计算加速增强）
1. MMA（矩阵乘法累加）演进
- 迭代历程：从Ampere的Warp级、Hopper的Warp Group级，演进为Blackwell的一对CTA（线程块） 级，提升并行计算能力。
- 指令家族：新增Block-scaled types支持，暴露为“tcgen05”指令家族，是Blackwell MMA的核心指令集。
2. 新增Tensor Memory（TMEM）
- 部署范围：每个SM（流多处理器）配备专用TMEM。
- 指令支持：提供专用的分配（alloc）、释放（dealloc）、加载（load）、存储（store）指令，确保高效数据操作。
- 数据存储：累加器（Accumulators）和操作数（Operands）直接存储于TMEM，减少数据搬运延迟。
- 地址空间：在编译器中为TMEM分配6号地址空间（Addrspace 6），明确数据存储位置。
3. TMA（Tensor Memory Accelerator）增强
- 新操作模式：支持im2col_w、scatter、gather等新模式，扩展数据处理场景。
- 功能升级：新增Masked Copy（掩码复制） 支持，提升数据复制的灵活性与准确性。
三、编译器Lowering流程与Blackwell扩展
1. 基础Lowering流程
- 整体链路：MLIR → LLVM → PTX，其中Python NVDSL仅用于示例演示，不参与实际编译链路。
- 关键组件：
- MLIR层：包含linalg、vector、nvgpu、自定义Dialect（Your Dialect 1/2）及NVVM Dialect（贴近PTX的低级操作）。
- LLVM层：通过NVVM内在函数（NVVM Intrinsics）承接MLIR操作。
- NVPTX Backend：将NVVM内在函数转换为最终的PTX指令。
2. Blackwell专属扩展
| 编译器组件 | 扩展内容 | 关键数字 |
|---|---|---|
| MLIR | 新增针对Blackwell的操作数 | 约30个 |
| NVVM Dialect | 新增Blackwell相关内在函数 | 约1000个 |
| NVPTX Backend | 新增TMEM对应的地址空间 | 6号地址空间 |
四、内在函数（Intrinsics）设计与性能影响
1. 内在函数数量规模
- 单指令扩展：1条TMA变体指令可扩展为数百个内在函数（受维度、模式、cta_group、缓存提示、多播等因素影响）。
- 指令家族规模：tcgen05-MMA家族 alone可扩展至约750个独立内在函数（受类型、复用提示、元数据缩放/稀疏性、掩码等因素影响）。
- 整体基数：pre-Blackwell阶段，NVVM IR内在函数总计约3000个，其中MMA和TMA相关内在函数占比约40%。
2. 对二进制大小的影响
基于TMA内在函数基础结构的实证观察，不同数量内在函数对二进制大小的影响如下表（单位：KB）：
| 二进制组件 | 1K内在函数 | 2K内在函数 | 5K内在函数 | 10K内在函数 | 与基线相比增长占比 |
|---|---|---|---|---|---|
| opt | 196 | 388 | 952 | 1908 | - |
| llc | 196 | 388 | 952 | 1908 | - |
| llvm-as | 48 | 100 | 256 | 516 | - |
| llvm-dis | 52 | 104 | 260 | 520 | - |
| Total | 492 | 980 | 2420 | 4852 | 0.32%/0.63%/1.56%/3.12% |
3. 对LLVM构建时间的影响
构建Tablegen生成必要头文件的时间变化如下表（单位：ms）：
| 内在函数数量 | FrontEnd IR Tablegen（IntrinsicImpl.inc） | NVPTX Backend（GenDAGISel.inc） | ||
|---|---|---|---|---|
| 时间（ms） | 增长占比 | 时间（ms） | 增长占比 | |
| 基线（Baseline） | 322 | 0.00% | 752 | 0.00% |
| 1K内在函数 | 370 | 14.91% | 919 | 22.21% |
| 2K内在函数 | 396 | 22.98% | 1034 | 37.50% |
| 5K内在函数 | 452 | 40.37% | 1612 | 2.1倍 |
| 10K内在函数 | 537 | 66.77% | 3168 | 4.2倍 |
- 关键说明：上述增长仅针对特定组件，整体LLVM构建时间（make -j8 with Release=1）无影响。
4. 对编译时间的影响
- 测试条件：100个Blackwell内核，使用“time-passes”选项，取20次运行平均值，基于NVCC（cicc）Release构建。
- 测试结果：
- opt阶段：无任何影响。
- llc阶段：不同内在函数数量对应的geo mean如下表：
| 内在函数数量 | 1K内在函数 | 2K内在函数 | 5K内在函数 | 10K内在函数 |
|---|---|---|---|---|
| Geo Mean | 1.01 | 1.04 | 1.08 | 1.17 |
5. 内在函数优化方案
为平衡“内在函数数量膨胀”与“代码可读性”，提出两种优化方案：
（1）打包标志（Packed flags，PR 96083 - v1）
- 核心逻辑：用一个“flag”操作数，通过位域（bit fields） 编码所有修饰符（如缓存提示、多播、加载模式等）。
- 示例结构：
typedef union {int V;struct { unsigned CacheHint: 1;unsigned Multicast: 1;unsigned LoadMode: 3; // 对应CpAsyncBulkTensorLoadModeunsigned NumCTAs: 1;unsigned reserved: 28;} U; } CpAsyncBulkTensorFlags; enum class CpAsyncBulkTensorLoadMode { TILE = 0, IM2COL = 1, IM2COL_W = 2, IM2COL_W_128 = 3 }; - 优势：减少内在函数数量，提升可维护性。
- 不足：当修饰符字段超出位宽时难以处理；内在函数复杂时，flag字段解读困难。
（2）分离标志（Separate flags，PR 96083）
- 核心逻辑：为每一组指令修饰符设置独立的“flag”字段（如flag_mc对应多播、flag_ch对应缓存提示、flag_mode对应加载模式）。
- 示例转换：
cp_async_bulk_tensor_var1/var2/…varN(...)→cp_async_bulk_tensor(i1 %flag_mc, i1 %flag_ch, i32 %flag_mode, ...) - 优势：
- 理论上永不出现位宽不足问题。
- 显著提升IR调试/修改效率（flag值与枚举值直接对应）。
- 不足：部分内核的比特码大小最多增加3%，但可读性提升的收益远大于此损耗。
五、APFloat与NVVM Dialect支持
1. APFloat类型扩展（支持OCP-MX规范新类型）
- 已支持类型：
- FP6类型：包含e2m3、e3m2两种格式，对应PR 94735。
- FP4类型：采用e2m1格式，对应PR 95392。
- 关键突破：解决e8m0类型“precision=0”的处理难题（PR 107127），通过“Float8E8M0FNUExhaustivePair”测试（APFloatTest.cpp）暴露并修复边角案例。
- 集成状态：所有新类型在MLIR中均已完成类型定义支持。
2. NVVM Dialect支持进展
- 核心支持范围：除MMA外，所有tcgen05操作均已在NVVM Dialect中支持，且通过NVPTX内在函数完成lowering。
- 功能特性：
- 同时支持inline-ptx和内在函数lowering两种方式。
- 支持单个Op部分lowering至任一方式（灵活性高）。
- 迁移进展：正积极将现有操作从inline-ptx迁移至内在函数lowering，目前仍有14/90个Op待迁移。
六、NVDSL Python GEMM示例（128x128x64）
1. 示例基础信息
- 计算任务：128x128x64规模的GEMM（矩阵乘法）
- 线程配置：启动1个CTA（grid=(1,1,1)），线程块大小为block=(128,1,1)
- 核心目标：将GEMM计算卸载到Blackwell的Tensor Core，展示NVGPU/NVVM Dialect的实际应用。
2. 核心执行步骤（基于gemm_kernel函数）
| 步骤序号 | 操作内容 | 执行线程条件 |
|---|---|---|
| 1 | TMA fence（确保TMA描述符一致性）、初始化2个mbarrier（同步用） | warpIdx == 0 且 elect_one() |
| 2 | 分配128列的TMEM资源 | warpIdx == 0 |
| 3 | 执行nvvm.fence_mbarrier_init（确保mbarrier初始化顺序）、nvvm.barrier（线程同步） | 所有线程 |
| 4 | 执行TMA加载（将数据加载到TMEM） | warpIdx == 0 且 elect_one() |
| 5 | 等待TMA加载完成 | warpIdx == 0 |
| 6 | 执行Tensor Core MMA（4次循环，使用tcgen05指令）、提交mbarrier到达信号 | warpIdx == 0 且 elect_one() |
| 7 | 等待Tensor Core MMA计算完成 | 所有线程 |
| 8 | 将TMEM中的数据加载到寄存器（使用nvvm.tcgen05_ld，形状为32X32B） | 所有线程 |
| 9 | 将寄存器中的结果存储到全局内存（GMEM） | 所有线程（128次循环） |
| 10 | 释放TMEM资源、放弃TMEM分配许可 | warpIdx == 0 |
3. NVDSL工具价值
- 简化开发：通过装饰器构建MLIR函数、支持JIT编译与执行、算术Dialect操作符重载、NumPy到memref的转换。
- 参考意义：提供可直接运行的Blackwell示例，支持开发者基于示例构建自定义编译器。
七、后续计划与结论
1. 后续开发计划
- 功能优化：支持Imm内在函数参数（flags）的“美化打印”，对应RFC文档（discourse 82629）。
- 内在函数增强：为内在函数附加地址空间特定的RW（读写）属性（基于NVVM Dialect）。
- 版本检查：为NVVM Ops添加SM版本检查逻辑，目前PR 126886 处于审核中。
- 功能补全：继续完成Blackwell剩余指令（如TMA、MMA）和内在函数的支持。
2. 已完成与进行中工作
| 工作类别 | 已完成内容 | 进行中内容 |
|---|---|---|
| 基础支持 | Hopper基线支持（NVPTX与NVVM Dialect均完成） | Blackwell的TMA、MMA等更新 |
| Blackwell核心 | tcgen05家族基础建模、APFloat新类型（FP6、FP4、e8m0）支持 | 剩余指令/内在函数支持 |
| NVDSL示例 | 提供128x128x64 GEMM可运行示例、支持基于示例构建自定义编译器 | - |
3. 致谢
感谢所有审核者（reviewers）和维护者（Maintainers）对该项目的支持。
4. 关键问题
问题1：Blackwell GPU在MMA和TMEM上的核心改进是什么？与Hopper GPU相比有哪些突破？
答案：
Blackwell GPU在MMA和TMEM上的核心改进及与Hopper的突破如下：
- MMA（矩阵乘法累加）改进：MMA的并行计算粒度从Hopper的“Warp Group（线程束组）”演进为Blackwell的“一对CTA（线程块）”，显著提升了大规模矩阵计算的并行效率；同时新增Block-scaled types支持，并将其暴露为“tcgen05”指令家族，成为Blackwell MMA的核心指令集，扩展了数据类型处理能力。
- TMEM（Tensor Memory）新增与优化：Blackwell为每个SM（流多处理器）配备了专用的TMEM，区别于Hopper无此硬件资源的设计；TMEM提供专用的分配（alloc）、释放（dealloc）、加载（load）、存储（store）指令，且累加器（Accumulators）和操作数（Operands）直接存储于TMEM，减少了数据在不同存储层级间的搬运延迟，提升计算效率；在编译器层面，为TMEM分配了6号地址空间，明确其数据存储位置，便于编译器优化。
问题2：在为Blackwell GPU扩展LLVM内在函数时，内在函数数量膨胀带来了哪些性能影响？又通过哪些方案缓解了这一问题？
答案：
一、内在函数数量膨胀的性能影响
- 二进制大小影响：随着内在函数数量增加，二进制各组件大小同步增长，1K/2K/5K/10K内在函数对应的总二进制大小分别比基线增加492KB（0.32%）、980KB（0.63%）、2420KB（1.56%）、4852KB（3.12%），但增长幅度仍在可控范围内。
- 构建时间影响：仅特定组件构建时间增长，10K内在函数时，FrontEnd IR Tablegen时间比基线增加66.77%，NVPTX Backend（GenDAGISel.inc）时间比基线增加4.2倍；整体LLVM构建时间（make -j8 with Release=1）无影响。
- 编译时间影响：opt阶段无影响，llc阶段随内在函数数量增加略有上升，10K内在函数时geo mean为1.17（即编译时间约为基线的1.17倍），对整体编译效率影响较小。
二、缓解方案
- 打包标志（PR 96083 - v1）：用一个“flag”操作数通过位域（bit fields）编码所有修饰符（如缓存提示、多播、加载模式等），减少内在函数数量，提升可维护性；但存在“修饰符字段超位宽难处理”“复杂内在函数flag解读困难”的不足。
- 分离标志（PR 96083）：为每一组指令修饰符设置独立“flag”字段（如flag_mc对应多播、flag_ch对应缓存提示），理论上永不出现位宽不足问题，且flag值与枚举值直接对应，显著提升IR调试/修改效率；仅部分内核比特码大小最多增加3%，但可读性收益远大于此损耗。
问题3：目前NVIDIA Blackwell GPU在LLVM和MLIR中的支持进展如何？已完成哪些关键工作，仍有哪些待推进任务？
答案：
一、已完成的关键支持工作
- 基础架构支持：Hopper GPU的基线支持已全部完成，覆盖NVPTX Backend和NVVM Dialect两个核心模块，为Blackwell支持奠定基础。
- Blackwell核心指令建模：完成“tcgen05”指令家族的基础建模，该家族是Blackwell MMA的核心指令集，支撑矩阵乘法累加等关键AI计算。
- 数据类型支持：APFloat模块已完成Blackwell所需新类型的支持，包括FP6（e2m3、e3m2，PR 94735）、FP4（e2m1，PR 95392），并解决了e8m0类型“precision=0”的处理难题（PR 107127），且所有新类型在MLIR中均完成类型定义。
- NVVM Dialect支持：除MMA外，Blackwell的所有tcgen05操作均已在NVVM Dialect中支持，且支持inline-ptx与内在函数lowering两种方式，仅剩余14/90个Op待迁移至内在函数lowering。
- 示例工具支持：通过NVDSL Python提供了128x128x64规模的GEMM可运行示例，展示了NVGPU/NVVM Dialect的实际应用，支持开发者基于示例构建自定义编译器。
二、仍待推进的任务
- 指令与内在函数补全：继续完成Blackwell GPU剩余TMA（Tensor Memory Accelerator）、MMA（矩阵乘法累加）等指令及对应内在函数的支持，确保所有核心计算功能覆盖。
- 功能优化与增强：支持Imm内在函数参数（flags）的“美化打印”（对应RFC: discourse 82629），提升代码可读性；为内在函数附加地址空间特定的RW（读写）属性，优化数据访问控制。
- 版本检查完善：推进NVVM Ops的SM版本检查功能（PR 126886，目前处于审核中），确保指令仅在支持的SM版本上执行，避免兼容性问题。