3. 编译器前端设计：模型解析与中间表示（IR）、TensorFlow/PyTorch模型导入、ONNX格式解析

好，咱们今天聊聊编译器前端。说白了，就是怎么把训练好的模型“吃进去”，变成编译器能理解的东西。

我刚开始做AI芯片编译器那会儿，最头疼的就是模型解析。你想想看，TensorFlow、PyTorch、ONNX，各家有各家的脾气。模型文件一进来，怎么把它拆开、读懂、再组装成统一的中间表示？这就是前端要干的事。

3.1 中间表示（IR）——编译器的“通用语言”

IR是什么？就是中间表示。它像一座桥，连接着前端和后端。

前端解析完模型，生成IR。后端基于IR做优化，最后生成目标代码。没有IR，编译器就是一盘散沙。

我个人习惯把IR分成几个层次：

• 图级别IR：描述计算图的结构，节点是算子，边是数据依赖。比如TensorFlow的GraphDef、PyTorch的TorchScript。
• 算子级别IR：每个算子内部的计算逻辑。比如卷积的stride、padding这些参数。
• 指令级别IR：接近硬件的指令序列。比如LLVM IR、自家芯片的微码。

嗯，这里要注意。图级别IR和算子级别IR，是编译器前端最关心的。指令级别IR，那是后端的事。

我在项目中遇到过一个问题。当时我们设计IR时，把卷积和全连接层分成了两个不同的算子。结果做算子融合时，发现卷积+BN（批归一化）可以融合，但全连接+BN却不行。为什么？因为全连接层的权重排布和卷积不一样。后来我们统一了IR中“线性变换”的表示，才解决了这个问题。

3.2 TensorFlow模型导入

TensorFlow的模型格式，主要有两种：

• SavedModel：官方推荐格式，包含计算图和权重
• Frozen GraphDef：冻结后的图，权重固化在图中

解析TensorFlow模型，核心是读懂它的GraphDef。GraphDef是一个protobuf序列化的文件，里面定义了所有节点（NodeDef）和它们之间的连接。

每个NodeDef包含：

• name：节点名称
• op：算子类型（如Conv2D、Relu、MatMul）
• input：输入节点列表
• attr：属性字典（如卷积的stride、padding）

举个例子，解析一个卷积层：

// 伪代码示意
def parse_tf_conv(node):
    op_type = node.op  // "Conv2D"
    inputs = node.input  // ["input_tensor", "weight_tensor"]
    strides = node.attr["strides"].list.i
    padding = node.attr["padding"].s
    // 转换成IR中的卷积算子
    ir_node = IRConv2D(strides=strides, padding=padding)
    return ir_node

我曾经踩过一个坑。TensorFlow的NHWC格式和PyTorch的NCHW格式，在解析时很容易搞混。有一次，我解析完模型，发现推理结果全是错的。查了两天，才发现是数据排布没做转换。从那以后，我每次解析完模型，都会先跑一个简单的测试用例，验证结果是否正确。

3.3 PyTorch模型导入

PyTorch的模型导入，比TensorFlow稍微复杂一点。因为PyTorch是动态图，模型结构是在运行时构建的。

常用的方法有两种：

• TorchScript：通过torch.jit.trace或torch.jit.script导出静态图
• ONNX导出：通过torch.onnx.export导出ONNX格式

我个人更推荐用TorchScript。为什么？因为它保留了更多的控制流信息。ONNX导出时，有些动态操作会被展开成静态图，丢失了灵活性。

解析TorchScript模型，核心是读懂它的Graph。TorchScript的Graph是一个有向无环图（DAG），每个节点对应一个操作。

举个例子：

// TorchScript Graph示例
graph(%x : Tensor, %w : Tensor, %b : Tensor):
  %conv = aten::conv2d(%x, %w, %b, %strides, %padding)
  %relu = aten::relu(%conv)
  return (%relu)

解析时，我们需要把aten::conv2d映射到IR中的卷积算子，把aten::relu映射到IR中的激活算子。

嗯，这里有个细节。PyTorch的算子命名和TensorFlow不一样。比如卷积，TensorFlow叫Conv2D，PyTorch叫aten::conv2d。在IR中，我们需要统一命名，方便后续优化。

3.4 ONNX格式解析

ONNX，全称Open Neural Network Exchange。它像一个“翻译官”，让不同框架的模型可以互相转换。

ONNX的模型结构，也是protobuf序列化的。核心是ModelProto和GraphProto。

• ModelProto：包含模型元信息（版本、生产者等）
• GraphProto：包含计算图（节点、边、初始值）

每个节点（NodeProto）包含：

• op_type：算子类型（如Conv、Relu、Gemm）
• input：输入名称列表
• output：输出名称列表
• attribute：属性列表（如卷积的kernel_shape、strides）

解析ONNX模型，核心是遍历GraphProto中的节点，逐个转换成IR中的算子。

举个例子：

// ONNX节点解析伪代码
def parse_onnx_node(node):
    op_type = node.op_type  // "Conv"
    inputs = node.input  // ["input", "weight", "bias"]
    outputs = node.output  // ["conv_output"]
    attrs = {}
    for attr in node.attribute:
        attrs[attr.name] = attr
    // 转换成IR中的卷积算子
    ir_node = IRConv2D(
        kernel_shape=attrs["kernel_shape"].ints,
        strides=attrs["strides"].ints,
        pads=attrs["pads"].ints
    )
    return ir_node

我在项目中遇到过一个问题。ONNX的算子版本经常更新，同一个算子在不同版本中，属性名称可能不一样。比如Conv算子，在opset 10中，pads属性是必须的；在opset 11中，pads变成了可选。如果不做版本兼容，解析就会报错。

3.5 从模型到IR的完整流程

好了，咱们把整个流程串起来。从模型文件到IR，大致分三步：

1. 模型加载：读取模型文件，解析成框架特定的数据结构（如TensorFlow的GraphDef、PyTorch的Graph、ONNX的ModelProto）
2. 算子映射：把框架特定的算子，映射到IR中统一的算子集合
3. 图构建：根据算子之间的连接关系，构建IR图

下面这张图，展示了整个流程：

你看，不管输入是TensorFlow、PyTorch还是ONNX，最终都汇聚到统一的IR。这就是编译器前端的核心价值——屏蔽框架差异，提供统一的优化入口。

好了，关于编译器前端设计，咱们就聊到这儿。记住一句话：前端做得好，后端优化事半功倍；前端做得糙，后面全是坑。

---