第三章：前端解析与算子识别

好，咱们进入第三章。这一章讲的是编译器前端最核心的活儿——怎么把各种模型文件吃进来，再转化成咱们内部能理解的东西。

说白了，NPU编译器要面对的第一个问题就是：模型格式五花八门，怎么统一处理？

我个人习惯把这一章分成三块：模型格式解析、算子映射与注册、自定义算子处理。咱们一个一个来。

3.1 模型格式解析

现在主流的深度学习框架，无非就是ONNX、TensorFlow、PyTorch这三家。每家都有自己的序列化格式，解析方式也完全不同。

3.1.1 ONNX 格式解析

ONNX 是我个人最喜欢的中间格式。它设计得比较干净，用 Protobuf 序列化，结构清晰。

解析 ONNX 模型，核心就两步：

1. 读取 .onnx 文件，反序列化成 ModelProto 对象
2. 遍历 GraphProto 中的 NodeProto，提取算子信息

代码大概长这样：

// 伪代码，实际用 C++ 或 Python 都行
ModelProto model;
ReadProtoFromBinaryFile("model.onnx", &model);

GraphProto graph = model.graph();
for (const auto& node : graph.node()) {
    // node.op_type() 就是算子类型，比如 "Conv", "Relu"
    // node.input() 和 node.output() 是张量名
    // node.attribute() 是算子属性
    std::string op_type = node.op_type();
    // 然后根据 op_type 做映射
}

3.1.2 TensorFlow 格式解析

TensorFlow 的格式就复杂一些。它有两种主流格式：

• SavedModel：目录结构，包含 saved_model.pb 和变量文件
• Frozen GraphDef：单个 .pb 文件，所有权重都固化在图中

解析 TF 模型，我建议直接用它的 C++ API 或者 Python API。自己手写解析器？嗯，别找罪受。

// 用 TensorFlow C API 加载
TF_Graph* graph = TF_NewGraph();
TF_Session* session = ...;
// 读取 .pb 文件到 buffer
TF_ImportGraphDefOptions* opts = TF_NewImportGraphDefOptions();
TF_GraphImportGraphDef(graph, graph_def, opts, status);
// 然后遍历 graph 中的 operation

3.1.3 PyTorch 格式解析

PyTorch 的 torch.jit.trace 或 torch.jit.script 导出的是 TorchScript，用 zip 格式打包，里面是 pickle 序列化的 Module 对象。

解析 TorchScript，我推荐用 PyTorch 自带的 torch.jit.load()，然后遍历 graph：

import torch

model = torch.jit.load("model.pt")
graph = model.graph

# 遍历所有节点
for node in graph.nodes():
    kind = node.kind()  # 比如 "aten::conv2d"
    inputs = list(node.inputs())
    outputs = list(node.outputs())
    # 根据 kind 做映射

这里有个坑：PyTorch 的算子名字是 aten::xxx 这种格式，跟 ONNX 的命名完全不同。映射的时候要小心。

3.2 算子映射与注册机制

模型解析完了，接下来就是算子映射。说白了，就是把框架里的算子，映射到咱们 NPU 编译器内部的 IR 上。

我习惯用一张映射表来管理：

框架算子	内部 IR 算子	备注
ONNX: Conv	Conv2D	直接映射
TF: Conv2D	Conv2D	直接映射
PyTorch: aten::conv2d	Conv2D	需要处理参数顺序
ONNX: BatchNormalization	BatchNorm	融合到 Conv 中
TF: FusedBatchNorm	BatchNorm	注意 training 标志

映射的注册机制，我推荐用工厂模式加宏定义：

// 注册宏
#define REGISTER_OP_MAPPING(framework_op, internal_op) \
    static OpMapperRegister __mapper_##framework_op(framework_op, internal_op);

// 使用
REGISTER_OP_MAPPING("Conv", "Conv2D");
REGISTER_OP_MAPPING("Relu", "Relu");

// 实际映射时查表
std::string MapOp(const std::string& framework_op) {
    auto it = mapper_map.find(framework_op);
    if (it != mapper_map.end()) {
        return it->second;
    }
    // 没找到？走自定义算子处理
    return HandleCustomOp(framework_op);
}

3.3 自定义算子的处理

这是最头疼的部分。你想想看，用户可能写了一个 MyFancyOp，ONNX 里根本没有对应定义。怎么办？

我一般分三步走：

1. 识别：在映射表里找不到的算子，标记为“自定义”
2. 降级：看能不能用基础算子组合实现。比如 MyFancyOp 可能等于 Conv + Relu + Add
3. 扩展：如果降级不了，就提供扩展接口，让用户自己写 kernel

降级的代码示例：

// 假设用户自定义了一个 "MyFancyOp"
// 我们尝试用基础算子组合
if (op_type == "MyFancyOp") {
    // 拆成三个基础算子
    auto conv = ir_builder->CreateOp("Conv2D", ...);
    auto relu = ir_builder->CreateOp("Relu", ...);
    auto add = ir_builder->CreateOp("Add", ...);
    
    // 连接起来
    conv->SetInput(0, input);
    relu->SetInput(0, conv->GetOutput(0));
    add->SetInput(0, relu->GetOutput(0));
    add->SetInput(1, bias);
    
    return add->GetOutput(0);
}

如果降级搞不定，那就得开放扩展接口了。我一般提供一个 CustomOpHandler 基类：

class CustomOpHandler {
public:
    virtual bool CanHandle(const std::string& op_type) = 0;
    virtual std::vector<IRNode*> Handle(const OpDef& op) = 0;
};

// 用户继承这个类，注册自己的处理器
class MyCustomHandler : public CustomOpHandler {
    bool CanHandle(const std::string& op_type) override {
        return op_type == "MyFancyOp";
    }
    std::vector<IRNode*> Handle(const OpDef& op) override {
        // 用户自己实现
    }
};

嗯，这里要注意：自定义算子的性能通常不如原生算子。因为编译器没法针对它做深度优化。所以我的建议是——尽量用标准算子，实在不行再走自定义。

小结

这一章的内容，说白了就是三件事：

• 解析：把 ONNX、TF、PyTorch 的模型文件读进来，提取算子图
• 映射：把框架算子转成内部 IR 算子，注册到映射表里
• 自定义：处理那些映射表里没有的算子，能降级就降级，不能降级就开放扩展

我个人觉得，前端解析这部分，最考验的是对各种框架细节的了解。你想想看，ONNX 的 padding 模式有 3 种，TF 有 2 种，PyTorch 又有自己的方式。映射的时候稍不注意，精度就差了那么一点点。

下一章咱们讲中端优化，到时候会用到今天映射出来的 IR 图。好，今天就到这儿。