前端解析与算子识别：ONNX模型解析、TensorFlow GraphDef解析、PyTorch TorchScript解析、算子模式匹配算法
各位同学，咱们今天聊聊编译器前端最核心的活儿——怎么把不同框架的模型吃进来，再认出里面的算子。说实话，这块我踩过的坑比代码行数还多。你想想看，ONNX、TensorFlow、PyTorch，三家各玩各的，格式不一样，命名不一样，连图结构都不同。编译器要做的，就是把它们统一成自己的中间表示。
ONNX模型解析：最省心的那个
我个人习惯先从ONNX说起。为什么？因为ONNX本身就是为交换而生的。它的设计目标就是「一次导出，到处运行」。ONNX模型用Protobuf序列化，解析起来其实挺规整的。
ONNX的核心结构就三样：GraphProto、NodeProto、TensorProto。GraphProto里有一堆NodeProto，每个NodeProto代表一个算子。每个算子有输入、输出、属性。就这么简单。
关键点：ONNX的算子命名是标准化的。比如Conv、Relu、Gemm。你只要维护一张映射表，就能把ONNX算子映射到你的IR上。但要注意，ONNX的算子版本号会变。我遇到过ONNX opset 9和opset 11的Conv行为不一样，padding规则改了。嗯，这里要小心。
// 伪代码：ONNX模型解析核心逻辑
ModelProto model;
ParseFromFile("model.onnx", &model);
GraphProto graph = model.graph();
for (NodeProto node : graph.node()) {
// 识别算子类型
string op_type = node.op_type();
// 提取输入输出
vector<string> inputs = node.input();
vector<string> outputs = node.output();
// 提取属性
for (AttributeProto attr : node.attribute()) {
if (attr.name() == "kernel_shape") {
// 处理卷积核大小
}
}
// 映射到内部IR
IRNode ir_node = CreateIRNode(op_type, inputs, outputs, attrs);
graph_builder->AddNode(ir_node);
}
TensorFlow GraphDef解析：最折腾的那个
说实话，TensorFlow的GraphDef是我最头疼的。为什么？因为TF的图结构太灵活了。它不像ONNX那样规规矩矩，而是允许各种奇奇怪怪的控制流、变量作用域、甚至自定义算子。
TF的GraphDef也是Protobuf格式，但它的NodeDef里有个op字段，这个字段的值五花八门。比如卷积叫Conv2D，但有时候又会出现FusedConv2D。还有Identity节点，看着没用，但删掉就会出问题。
避坑指南：我曾经在解析TF模型时，遇到一个Switch节点，死活找不到对应的ONNX算子。后来才发现，这是TF控制流的一部分，需要特殊处理。如果你直接忽略，模型推理结果就是错的。
TF的另一个特点是节点命名规则。每个节点有名字，但名字里可能包含作用域信息。比如conv1/kernel、conv1/bias。解析时你得把这些作用域拆开，才能正确重建图结构。
// 伪代码：TensorFlow GraphDef解析
GraphDef graph_def;
ReadBinaryProto(env, "model.pb", &graph_def);
for (NodeDef node : graph_def.node()) {
string op = node.op();
string name = node.name();
// 处理特殊算子
if (op == "Conv2D") {
// 提取卷积参数
auto strides = node.attr().at("strides");
auto padding = node.attr().at("padding");
} else if (op == "FusedBatchNorm") {
// 融合BN需要拆开处理
} else if (op == "Switch" || op == "Merge") {
// 控制流节点，特殊标记
}
// 处理输入输出
for (string input : node.input()) {
// 注意：TF的输入可能带端口号，比如 "conv1:0"
string node_name = ParseInputName(input);
int port = ParseInputPort(input);
}
}
PyTorch TorchScript解析：最年轻的那个
TorchScript是PyTorch的中间表示。它跟ONNX和TF都不一样。TorchScript不是Protobuf，而是用JIT编译器生成的IR。它的结构更像一个控制流图，有Graph、Block、Value、Node这些概念。
TorchScript的算子叫prim::或aten::开头。比如aten::conv2d、aten::relu。但有个坑：PyTorch的算子参数是位置化的，不像ONNX那样有名字。你得按顺序解析参数。
个人经验：我建议解析TorchScript时，先把它转成ONNX再处理。因为TorchScript的图结构太动态了，很多控制流是if和loop，直接解析容易出错。但如果你非要硬啃，记得处理prim::Constant节点，它代表常量，经常被忽略。
// 伪代码：TorchScript解析
auto graph = module.get_method("forward").graph();
for (Node* node : graph->nodes()) {
string kind = node->kind().toQualString();
if (kind == "aten::conv2d") {
// 按位置提取参数
Value* input = node->input(0);
Value* weight = node->input(1);
Value* bias = node->input(2);
Value* stride = node->input(3);
Value* padding = node->input(4);
} else if (kind == "prim::If") {
// 处理条件分支
Block* true_block = node->blocks()[0];
Block* false_block = node->blocks()[1];
}
}
算子模式匹配算法：最核心的那个
好了，模型解析完了，图也建好了。接下来要干嘛？算子模式匹配。说白了，就是找图中的特定子图结构，然后替换成更高效的算子。
举个例子：Conv + BN + Relu这种模式，在很多NPU上可以融合成一个算子。但你怎么在图中找到这种模式？
常用的方法有两种：
- 基于拓扑排序的暴力匹配：遍历所有节点，对每个节点检查它的前驱和后继，看是否匹配目标模式。简单粗暴，但效率低。
- 基于DAG的图匹配：把目标模式也建成一个子图，然后用子图同构算法去匹配。效率高，但实现复杂。
核心思路：我个人习惯用拓扑排序+模式缓存。先对图做拓扑排序，然后从后往前扫描。每遇到一个可能的模式起点，就尝试匹配。匹配成功后，把匹配结果缓存起来，避免重复计算。这样性能能提升不少。
// 伪代码：算子模式匹配
// 目标模式：Conv + BN + Relu
Pattern pattern;
pattern.AddNode("conv", "Conv");
pattern.AddNode("bn", "BatchNorm");
pattern.AddNode("relu", "Relu");
pattern.AddEdge("conv", "bn");
pattern.AddEdge("bn", "relu");
// 遍历图进行匹配
for (Node* node : graph->TopologicalOrder()) {
if (node->op_type() == "Conv") {
// 尝试匹配模式
MatchResult result = pattern.Match(node);
if (result.success) {
// 替换为融合算子
FusedNode fused = CreateFusedConvBnRelu(result);
graph->ReplaceSubgraph(result.nodes, fused);
}
}
}
知识体系总览
为了让你更直观地理解整个流程，我画了一张图。这张图展示了从模型解析到算子识别的完整链路。
总结一下
前端解析这块，说白了就是三件事：读进来、认出来、统一起来。ONNX最省心，TF最折腾，TorchScript最年轻。但不管哪种格式，最终都要落到你的IR上。算子模式匹配则是锦上添花，能帮你把多个小算子合成一个大算子，提升NPU的执行效率。
嗯，今天就先聊到这儿。下一章咱们会深入IR的设计，讲讲怎么设计一个既灵活又高效的中间表示。到时候我会拿我实际项目中的IR设计案例来拆解，保证干货满满。