第4章 IR图构建

4、IR图构建：从TFLite算子列表构建计算图、节点与边的关系、控制流处理

好，咱们进入第四讲。说实话，IR图构建这块，是模型迁移中最容易出幺蛾子的环节。很多同学觉得不就是把TFLite的算子列表转成一张图嘛，有啥难的？嗯，我当年也是这么想的，直到被一个控制流bug折磨了整整两天。

今天咱们就聊聊，怎么把TFLite那堆扁平的算子列表，变成Fusion Compiler能理解的计算图。说白了，就是给算子们排兵布阵，搞清楚谁是谁的输入，谁是谁的输出。

4.1 从平坦列表到有向图

TFLite的模型结构，其实就是一个大数组。每个算子都带着自己的输入输出tensor索引。你想想看，这就像一堆人站那儿，每个人手里拿着一张纸条，上面写着「我左边是谁，我右边是谁」。我们的任务，就是根据这些纸条，把整个队列还原出来。

我个人习惯，第一步先建一个「tensor到算子」的映射表。为什么？因为TFLite里tensor是全局编号的，但算子只知道自己用哪些tensor，不知道谁生产了这些tensor。我需要在遍历算子的过程中，把每个tensor的生产者算子记录下来。

// 伪代码示意
std::unordered_map<int, TfLiteNode*> tensor_producer;
for (int i = 0; i < interpreter->nodes_size(); ++i) {
    auto* node = interpreter->node(i);
    for (int j = 0; j < node->outputs->size; ++j) {
        int tensor_idx = node->outputs->data[j];
        tensor_producer[tensor_idx] = node;
    }
}

有了这个映射，构建节点和边的关系就简单了。每个算子就是一个节点，它读入的tensor来自哪个算子，那条边就指向谁。我在项目中遇到过一个问题：有些tensor是模型的输入，没有生产者。这时候要特殊处理，把它标记为「图输入节点」。

4.2 节点与边的数据结构设计

Fusion Compiler里，我习惯把节点设计成这样的结构：

struct FusionNode {
    int node_id;                // 节点ID
    std::string op_type;        // 算子类型，比如CONV_2D
    std::vector<int> inputs;    // 输入tensor索引列表
    std::vector<int> outputs;   // 输出tensor索引列表
    std::vector<int> pred_ids;  // 前驱节点ID列表
    std::vector<int> succ_ids;  // 后继节点ID列表
    void* builtin_options;      // 算子参数
};

边呢？我一般不单独建一个Edge类。直接用pred_ids和succ_ids就够了。为什么？因为Fusion Compiler的图优化阶段，经常要增删节点，用邻接表的方式维护起来更灵活。你想想看，如果每条边都是一个独立对象，删一个节点得遍历所有边，多麻烦。

小技巧：构建pred_ids和succ_ids时，我建议用std::vector而不是std::list。虽然list插入删除快，但图优化时经常要随机访问邻居节点，vector的缓存局部性更好。我实测过，在大模型上能快15%左右。

4.3 控制流处理的坑

嗯，这里要注意。TFLite里有些算子不是简单的计算节点，而是控制流节点，比如WHILE、IF。这些家伙会引入子图，让整个图结构变成嵌套的。

我曾经踩过一个坑：处理WHILE循环时，直接把循环体内的算子全部展开到主图中。结果图变得巨大无比，优化器跑了几分钟都没结束。后来我才意识到，对于控制流，应该保留子图结构，而不是暴力展开。

我的做法是这样的：

遇到WHILE算子时，把它当作一个「超级节点」
这个超级节点内部，再递归构建一个子图
子图的输入输出，对应WHILE算子的cond和body子图

void BuildControlFlowNode(TfLiteNode* node, FusionGraph* graph) {
    if (node->builtin_data->type == kTfLiteBuiltinWhile) {
        auto* while_params = reinterpret_cast<TfLiteWhileParams*>(node->builtin_data);
        
        // 创建超级节点
        int super_node_id = graph->AddNode("WHILE");
        
        // 递归构建cond子图
        FusionGraph* cond_subgraph = new FusionGraph();
        BuildSubgraph(while_params->cond_subgraph_index, cond_subgraph);
        
        // 递归构建body子图
        FusionGraph* body_subgraph = new FusionGraph();
        BuildSubgraph(while_params->body_subgraph_index, body_subgraph);
        
        // 把子图挂到超级节点上
        graph->SetSubgraph(super_node_id, cond_subgraph, body_subgraph);
    }
}

注意：IF算子也有类似的问题。它的true分支和false分支是两个独立的子图。千万别把两个分支的节点混在一起，否则优化器会以为所有节点都在一条路径上，做出错误的优化决策。我见过有人因为这个，把永远不会执行的代码也给优化掉了。

4.4 图构建的完整流程

好了，咱们把整个流程串起来。下面这张图展示了从TFLite算子列表到Fusion Compiler计算图的完整过程：

你看，整个流程其实不复杂。但每一步都有细节要注意。我再说几个实战中容易翻车的地方：

多输出算子：比如Split算子，它会产生多个输出tensor。每个输出tensor都要正确映射到生产者，不能漏掉任何一个。
共享tensor：同一个tensor可能被多个算子消费。比如某个中间结果既给Conv用，又给Add用。这时候这个tensor的生产者节点，要有多个后继节点。
图验证：构建完图之后，我习惯跑一遍拓扑排序。如果排不出来，说明图里有环，那肯定是有bug。我遇到过因为控制流处理不当，把循环体的输出连回输入，导致出现环的情况。

核心要点：IR图构建的本质，就是把TFLite的「扁平算子列表」转换成Fusion Compiler的「有向无环图」。控制流节点要特殊处理，保留子图结构。构建完成后，一定要做图验证，确保没有环、没有悬空的tensor。

好了，这一讲就到这里。图构建是后续所有优化的基础，这块做扎实了，后面的算子融合、内存规划才能顺利推进。下一讲咱们聊聊图优化中的算子融合策略，到时候我会分享一些实际项目中的融合经验。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录

4、IR图构建：从TFLite算子列表构建计算图、节点与边的关系、控制流处理

4.1 从平坦列表到有向图

4.2 节点与边的数据结构设计

4.3 控制流处理的坑

4.4 图构建的完整流程