通用编译器,例如GCC和Clang/LLVM,凝聚了数十年的研究与工程成果,为使用C++、Java或Fortran等语言编写的传统软件实现了卓越的性能。它们擅长优化标量代码、管理控制流、执行指令调度、寄存器分配以及使用SIMD指令进行向量化循环。然而,部署机器学习模型会引入一套与典型CPU密集型应用截然不同的计算模式和优化需求。这些差异使得编译和运行时优化需要采用专门的方法。机器学习计算的特性机器学习工作负载,尤其是在深度学习中,主要由对大型多维数组(即张量)的操作构成。可以想象,卷积、矩阵乘法、池化和激活函数等操作的数据常以 $N \times C \times H \times W$(批次,通道,高度,宽度)或类似的高维格式组织。尽管这些操作最终会分解为循环和算术运算,但它们的高级结构承载着对有效优化很重要的语义信息。考虑以下特性:高级运算符语义: TensorFlow或PyTorch等框架使用粗粒度运算符(例如Conv2D、MatMul、BatchNorm)表达计算。运算符融合(将Conv + Bias + ReLU组合成单个核)等优化,其核心在于识别和操作计算图中的这些高级运算符。标准编译器在编译过程的早期,当处理像LLVM IR这样的低级中间表示时,通常会丢失这种图结构和运算符语义。强数据局部性和重用: 张量操作展现出复杂的数据访问模式,在缓存和寄存器中具有显著的数据重用潜力。为此进行优化需要复杂的循环变换,例如分块、倾斜和重排序,这些变换常由多面体建模等技术指导,这类技术会明确分析表示张量计算的循环嵌套中的数据依赖关系。通用编译器可能会进行一些循环优化,但通常缺少领域特定的启发式规则或分析能力,无法完全发挥GEMM(通用矩阵乘法)或N维卷积等操作中固有的局部性。硬件特化: 机器学习推理常面向异构硬件。GPU通过数千个核心和NVIDIA的张量核心或AMD的矩阵核心等专用单元提供大规模并行处理能力。像Google TPU或专用NPU这样的定制ASIC采用脉动阵列等架构,这些架构是专为机器学习原语设计的。为这些目标生成最佳代码需要编译器后端知晓它们独特的指令集、内存层次结构(例如GPU共享内存)和执行模型。通用后端通常对这些高度专业化的功能提供不足的支持或抽象。图级别结构: 机器学习模型通常表示为操作的有向无环图(DAG)。通过优化这种图结构本身,可以获得显著的性能提升。这包括消除冗余计算、简化涉及张量操作的代数表达式,以及优化数据布局(例如,在NCHW和NHWC格式之间转换张量)以匹配硬件偏好或提升运算符融合的可能性。通用编译器主要侧重于传统调用图中的过程内优化,或者充其量是过程间优化,而非机器学习模型的数据流图语义。低精度算术的流行: 量化(使用INT8、FP8甚至更低位宽)等技术是加速推理和减少内存占用空间的标准做法。这在管理缩放因子、零点以及为专用低精度指令生成代码方面带来挑战。编译器需要特定的IR支持和优化阶段来有效处理量化算术,并且通常在一个模型中混合精度。这个方面在传统编译中基本不存在。通用编译器为何不足鉴于这些特性,仅依赖通用编译器部署机器学习模型会导致次优性能。主要问题有:语义差异: 当机器学习操作过早地降级到像LLVM IR这样的通用IR时,其高级结构和数学属性就会丢失。编译器看到的是循环和算术,而不是Conv2D或Attention。这种语义信息的丢失使其无法应用强大的、领域特定的变换。想象一下,如果只通过查看数据库执行引擎的低级C代码来优化一个复杂的SQL查询;你将错过高级别的优化机会。抽象不匹配: 标准IR通常缺乏足够的抽象来表示和优化专门的机器学习硬件功能(张量核心、专用内存空间)或机器学习特定的数据类型(带缩放/零点的量化整数)。优化范围有限: 通用编译器在函数或模块内部优化代码。它们通常不是为对表示整个机器学习模型的计算图执行全局的、结构转换的优化而设计的。跨框架级别函数的运算符融合,或通过多个运算符传播的布局转换,超出了它们的常规范围。为不同工作负载调整的启发式规则: 通用编译器中的优化启发式规则是基于数十年处理SPEC基准测试或应用软件等工作负载的经验调整的。这些启发式规则通常与深度学习张量操作的计算密集、内存密集和高度并行的特性不符。专门工具链的必要性为了弥补这一性能差距,该领域已发展出专门的机器学习编译器(例如XLA、TVM、Glow、基于MLIR的编译器)和运行时系统。这些系统采用:多级中间表示: 像MLIR这样的IR允许在多个抽象层次上表示计算,从框架级别操作到硬件特定指令,使得在最合适的层次进行优化成为可能。领域特定优化: 它们实现了针对机器学习量身定制的图级别处理(融合、布局转换)和张量级别优化(多面体转换、专用内核生成)。目标感知代码生成: 它们包含专门设计用于为CPU、GPU和各种AI加速器生成高度优化代码的后端,利用专用指令和内存子系统。运行时集成: 它们与复杂的运行时系统协同工作,处理张量形状变化等动态方面,在异构设备间高效管理内存,并调度异步操作。digraph AbstractionLevels { rankdir=TB; node [shape=box, style=filled, color="#ced4da", fontname="sans-serif"]; edge [fontname="sans-serif", fontsize=10]; Framework [label="机器学习框架操作\n(PyTorch, TF)", fillcolor="#a5d8ff"]; HighLevelIR [label="高级机器学习IR\n(图,张量操作)", fillcolor="#96f2d7"]; MidLevelIR [label="中级IR\n(循环,仿射)", fillcolor="#ffec99"]; LowLevelIR [label="低级IR\n(LLVM IR, SPIR-V)", fillcolor="#ffd8a8"]; MachineCode [label="机器码\n(x86, ARM, PTX, GCN)", fillcolor="#ffc9c9"]; subgraph cluster_ml_compiler { label = "机器学习编译器范围"; style=dashed; bgcolor="#e9ecef"; HighLevelIR; MidLevelIR; } subgraph cluster_general_compiler { label = "通用编译器范围"; style=dashed; bgcolor="#e9ecef"; LowLevelIR; } Framework -> HighLevelIR [label=" 导入/降级 "]; HighLevelIR -> MidLevelIR [label=" 图/张量优化\n 降级 "]; MidLevelIR -> LowLevelIR [label=" 循环/向量优化\n 降级 "]; LowLevelIR -> MachineCode [label=" 代码生成\n (指令选择, 寄存器分配) "]; Framework -> LowLevelIR [label=" 过早降级\n (丢失语义) ", style=dotted, color="#f03e3e", constraint=false]; }机器学习编译中典型的抽象层次。专门的机器学习编译器主要在更高级的IR上操作,保留对领域特定优化重要的语义信息,然后可能利用像LLVM这样的低级编译器进行最终的代码生成。通用编译器通常从较低的抽象层次开始,错失了机器学习图和张量表示中存在的优化可能性。"因此,理解并掌握这些专门的优化技术不仅仅是学术练习;它是实现机器学习部署中具有竞争力的性能的实际需要。本课程提供设计、实现和分析这些先进编译器和运行时策略所需的知识和技能。"