3. LLVM中间表示

LLVM IR（Intermediate Representation）是LLVM编译框架的核心枢纽，它将高级语言（如C/C++、Rust、Swift）与后端机器码生成解耦。对于深度学习编译器而言，理解LLVM IR不仅是编写自定义Pass的前提，更是将TVM的Relay/TIR算子高效映射到CPU/GPU指令的关键桥梁。本章将深入剖析LLVM IR的语法结构、强类型系统、指令集分类以及独特的内存模型，帮助你在阅读和生成IR时建立清晰的认知地图。

3.1 LLVM IR语法结构：从Module到BasicBlock

核心原理：三明治式分层设计

LLVM IR采用类似高级语言的结构化组织方式，但每个层次都严格遵循SSA（静态单赋值）约束。一个完整的IR单元称为Module（模块），它包含全局符号（全局变量、函数声明、函数定义）和模块级元数据。每个函数（Function）内部由若干BasicBlock（基本块）组成，基本块是直线执行的指令序列，且只能通过第一条指令（入口）进入，通过最后一条终结指令（terminator）退出。这种设计使得数据流分析变得简单：每个变量只被赋值一次，且定义点必然支配使用点。

一个典型的IR函数定义如下：

define i32 @add(i32 %a, i32 %b) {
  %sum = add i32 %a, %b
  ret i32 %sum
}

其中define关键字声明函数，i32是返回类型和参数类型，%a、%b是虚拟寄存器（SSA值），add指令执行整数加法，ret返回结果。所有局部变量必须以%开头，全局变量以@开头。

判断方法：识别IR的三种表示形式

LLVM IR有三种等价的表示形式，在调试和优化时需要灵活切换：

在TVM中，我们通常通过tvm.target.codegen接口将TIR节点转换为LLVM IR的内存表示，再通过llvm::raw_fd_ostream输出为文本或位码文件进行调试。

案例拆解：从TVM TIR到LLVM IR的语法映射

假设我们有一个简单的向量加法算子：C[i] = A[i] + B[i]。TVM的TIR代码经过Lowering后，最终生成的LLVM IR核心片段如下：

define void @vec_add(float* %A, float* %B, float* %C, i64 %n) {
entry:
  %cmp = icmp sgt i64 %n, 0
  br i1 %cmp, label %for.body, label %for.end

for.body:
  %i = phi i64 [ 0, %entry ], [ %inc, %for.body ]
  %a_ptr = getelementptr float, float* %A, i64 %i
  %b_ptr = getelementptr float, float* %B, i64 %i
  %c_ptr = getelementptr float, float* %C, i64 %i
  %a_val = load float, float* %a_ptr
  %b_val = load float, float* %b_ptr
  %sum = fadd float %a_val, %b_val
  store float %sum, float* %c_ptr
  %inc = add i64 %i, 1
  %cmp2 = icmp slt i64 %inc, %n
  br i1 %cmp2, label %for.body, label %for.end

for.end:
  ret void
}

注意phi指令的使用：它根据控制流选择%i的初始值0或递增后的值%inc。这是SSA形式处理循环变量的标准方式。TVM在生成IR时会自动插入phi节点，理解这一点对于调试循环优化Pass至关重要。

3.2 类型系统：强类型与向量类型的威力

核心原理：为什么LLVM需要强类型

LLVM IR是强类型语言，每个值（Value）都有明确的类型，且指令的操作数类型必须与指令语义匹配。这种设计带来三大优势：（1）优化Pass无需依赖数据流分析即可判断类型安全；（2）后端指令选择可以直接映射类型到目标机器寄存器类别；（3）调试时能快速发现类型不匹配的错误。对于深度学习编译器，最常用的类型包括：

• 整数类型：i1（布尔）、i8、i32、i64，用于索引、计数、位运算。
• 浮点类型：float（32位）、double（64位）、half（16位，GPU常用）。
• 向量类型：<4 x float>表示4个float组成的向量，直接映射到SIMD寄存器（如AVX的256位寄存器可容纳8个float）。
• 指针类型：float*、i8*，在opaque pointer（不透明指针）模式下，指针不再携带pointee类型信息，但load/store指令仍需显式指定数据类型。
• 聚合类型：{ float, i32 }结构体、[4 x float]数组，用于描述复杂数据结构。

操作流程：类型转换与检查

在编写TVM的Codegen时，经常需要将TIR的DataType映射为LLVM类型。例如：

// C++伪代码
llvm::Type* getLLVMType(const DataType& dtype) {
  if (dtype.is_float()) {
    switch (dtype.bits()) {
      case 16: return llvm::Type::getHalfTy(ctx);
      case 32: return llvm::Type::getFloatTy(ctx);
      case 64: return llvm::Type::getDoubleTy(ctx);
    }
  } else if (dtype.is_int()) {
    return llvm::Type::getIntNTy(ctx, dtype.bits());
  }
  // 向量类型处理
  if (dtype.lanes() > 1) {
    return llvm::VectorType::get(elemType, dtype.lanes(), false);
  }
}

类型不匹配时，LLVM会抛出Assertion错误。例如对float类型执行add指令（整数加法）会报错：Assertion `SrcTy->isIntegerTy()' failed。

案例拆解：向量化类型在深度学习中的实战

假设我们要对8个float数据执行ReLU激活函数。标量IR需要8条fcmp和8条select指令，而向量化后只需2条指令：

; 标量版本（部分）
%cmp0 = fcmp ogt float %x0, 0.0
%res0 = select i1 %cmp0, float %x0, float 0.0
; ... 重复8次

; 向量版本
%vec_x = load <8 x float>, <8 x float>* %ptr
%zero = fadd <8 x float> %vec_x, <zeroinitializer> ; 构造全0向量
%cmp = fcmp ogt <8 x float> %vec_x, %zero
%res = select <8 x i1> %cmp, <8 x float> %vec_x, %zero
store <8 x float> %res, <8 x float>* %ptr

TVM的ramp节点和broadcast节点在Lowering到LLVM IR时，会生成对应的向量类型和向量指令。在x86-64平台上，LLVM会自动将<8 x float>映射到YMM寄存器（AVX指令集），实现单指令多数据（SIMD）加速。

3.3 指令集分类：终结指令、算术逻辑指令与内存指令

核心原理：指令的五大类别

LLVM IR指令集约30余种，按功能可分为五类。理解分类有助于快速定位优化点：

操作流程：GEP指令的地址计算

getelementptr（GEP）是LLVM中最复杂也最强大的指令，它用于计算聚合类型（数组、结构体）元素的地址，但不访问内存。其语法为：

%ptr = getelementptr [4 x float], [4 x float]* %base, i64 0, i64 %idx

第一个索引（0）表示从%base指向的数组起始位置偏移0个数组元素，第二个索引（%idx）表示在数组内偏移%idx个float元素。对于多维张量，GEP可以链式使用：

; 访问 A[i][j]，其中A是 [N][M] x float 的二维数组
%row_ptr = getelementptr [N x [M x float]], [N x [M x float]]* %A, i64 0, i64 %i
%elem_ptr = getelementptr [M x float], [M x float]* %row_ptr, i64 0, i64 %j

TVM的BufferLoadNode在Lowering时，会根据张量的shape和strides生成对应的GEP链。理解GEP的语义对于手动调优内存访问模式（如将NCHW转换为NHWC）至关重要。

案例拆解：融合BiasAdd的卷积算子IR优化

假设我们有一个卷积算子，输出为out[n][oc][h][w] = conv(...) + bias[oc]。未优化时，TVM会生成两个独立的循环：先计算卷积结果存入临时缓冲区，再加载临时缓冲区并加上bias。优化后的IR将bias加载融合到卷积计算的内层循环中：

; 优化前（伪IR）
%conv_val = load float, float* %conv_out_ptr
%bias_val = load float, float* %bias_ptr
%sum = fadd float %conv_val, %bias_val
store float %sum, float* %out_ptr

; 优化后（将bias加载提升到循环外）
%bias_val = load float, float* %bias_ptr  ; 循环不变代码外提
; 内层循环中：
%conv_val = load float, float* %conv_out_ptr
%sum = fadd float %conv_val, %bias_val
store float %sum, float* %out_ptr

这个优化在LLVM中称为LICM（Loop Invariant Code Motion）。TVM的tvm.te.stage可以通过compute_inline或compute_at调度原语，在TIR层面提前完成类似优化，从而减少LLVM Pass的负担。

3.4 内存模型：alloca、load/store与别名分析

核心原理：SSA值 vs 内存对象

LLVM IR中，SSA值（如%sum）是编译期虚拟寄存器，不占用内存地址；而内存对象（通过alloca或全局变量定义）才有地址。这种区分使得优化Pass可以自由重排SSA值的计算顺序，但访问内存对象时必须遵守严格的别名规则。内存模型的核心包括：

• alloca：在栈上分配内存，返回指向该内存的指针。常用于局部数组、结构体或需要取地址的变量。
• load/store：显式地从内存读取/写入值。每次load都产生一个新的SSA值，store不产生值。
• 别名规则：如果两个指针可能指向同一内存区域，则它们互为别名。LLVM的noalias关键字可以声明指针不别名，从而允许更激进的优化（如重排load/store顺序）。

判断方法：识别内存访问的优化机会

在深度学习算子中，内存访问往往是性能瓶颈。以下信号提示存在优化空间：

1. 大量load指令重复加载同一地址：可通过gvn（全局值编号）Pass消除。
2. store后立即load同一地址：可通过dse（死存储消除）Pass删除冗余store。
3. GEP计算链过长：可通过strength-reduce（强度削减）将乘法运算替换为加法。
4. 循环内存在不变量load：应通过LICM提升到循环外。

案例拆解：内存布局转换（NCHW→NHWC）的IR级实现

假设输入张量布局为NCHW（通道在前），我们需要转换为NHWC（通道在后）以匹配cuDNN的期望格式。在LLVM IR层面，这涉及多维索引的重映射：

; 原始NCHW索引：n*C*H*W + c*H*W + h*W + w
; 目标NHWC索引：n*H*W*C + h*W*C + w*C + c

; 计算NHWC地址的IR片段
%nhwc_idx = add i64 %n_base, %h_base
%nhwc_idx2 = add i64 %nhwc_idx, %w_base
%nhwc_idx3 = add i64 %nhwc_idx2, %c
%src_ptr = getelementptr float, float* %input, i64 %nhwc_idx3

TVM的tir.transform.LayoutRewritePass会自动生成这类索引计算。但手动编写LLVM IR时，需要注意mul和add指令的溢出风险——当张量维度较大时（如N=256, C=1024, H=56, W=56），索引值可能超过32位整数范围，必须使用i64类型。

总结复盘：内存模型的三个关键认知

第一，SSA值的高效性建立在“不可变”基础上，任何需要修改的变量都必须通过alloca+load/store模拟，这会降低优化效果。因此，在TVM中应尽量使用函数式表达（te.compute），避免在TIR中直接使用tir.var赋值。第二，noalias是深度学习编译器最有力的优化武器——TVM的tir.Buffer默认声明noalias，因为不同张量（如输入和输出）在逻辑上不重叠。如果自定义算子时错误地共享了缓冲区，会导致LLVM优化出错误结果。第三，内存层级（L1/L2缓存、HBM）对性能的影响远大于指令选择，LLVM IR层面的load/store顺序调整（如循环分块tiling）往往比指令级优化带来更大的加速比。