TensorFlow运行时核心：DirectSession::Run全解析

TensorFlow运行时核心：DirectSession::Run全解析

TensorFlow的本地执行模式通过DirectSession::Run实现高效计算图处理，其核心流程分为三个阶段：图剪枝生成ClientGraph、设备间图分裂、跨设备并发执行。下面结合源码层设计深入剖析各环节实现原理。

---

一、图剪枝：从FullGraph到ClientGraph

目标：根据Session::Run指定的输入（feeds）和输出（fetches），从完整计算图（FullGraph）中提取最小依赖子图（ClientGraph）。
关键步骤：

1. 输入/输出节点重构

   • 输入边处理：为每个feed插入Arg节点（本地模式）或Recv节点（分布式模式），替代原始输入边。
   • 输出边处理：为每个fetch插入RetVal节点（本地）或Send节点（分布式）。

   # 输入边剪枝示例：插入Arg节点替代原始输入
   with tf.device('/cpu:0'):
       input_node = tf.Arg(0, dtype=tf.float32)  # 替代原始输入边
   

2. 反向剪枝算法

   • 以输出节点为起点，反向BFS遍历图，标记所有与feeds/fetches关联的节点。
   • 删除未访问节点（非依赖节点），连接入度为0的节点到Source，出度为0的节点到Sink。

   void PruneForTargets(Graph* g, const std::vector<Node*>& fetch_nodes) {
     std::unordered_set<const Node*> targets(fetch_nodes.begin(), fetch_nodes.end());
     ReverseBFS(g, targets);  // 反向遍历标记依赖节点
     RemoveUnvisitedNodes(g, targets); // 移除孤立节点
   }
   

---

二、图分裂：按设备切分ClientGraph

目标：将ClientGraph按设备约束拆分为多个子图（PartitionGraph），每个子图由独立设备执行。
设备间边处理策略：

1. 同设备子图
   直接保留边，无需额外操作（图13-18）。

2. 跨设备普通边
   插入Send/Recv节点对，通过Rendezvous通信（图12-19）：

   void InsertSendRecv(GraphDef* src_graph, const Edge* edge, GraphDef* dst_graph) {
     NodeDef* send = AddSend(opts, src_graph, edge); // 源设备插入Send节点
     NodeDef* recv = AddRecv(opts, dst_graph, edge); // 目标设备插入Recv节点
   }
   

3. 跨设备控制依赖边
   插入Const哑节点与Identity节点，维护控制流语义（图12-20）。

分裂算法核心：

Status Partition(const PartitionOptions& opts, Graph& client_graph) {
  for (Node* dst : client_graph.op_nodes()) {
    for (Edge* edge : dst->in_edges()) {
      if (edge->src().device() != dst->device()) {
        InsertSendRecv(edge);  // 设备不同则插入通信节点
      }
    }
  }
}

---

三、并发执行：Executor调度与跨设备通信

执行流程：

1. Executor启动
   每个PartitionGraph由独立的Executor调度：

   • 初始化FunctionCallFrame：传递输入数据（feeds → Arg节点）。
   • 构建ExecutorBarrier：同步多个Executor的执行状态。

   Executor::Args args;
   args.call_frame = &call_frame;        // 数据传递接口
   args.rendezvous = run_state.rendez;   // 跨设备通信组件
   item.executor->RunAsync(args, barrier->Get());  // 异步启动Executor
   

2. 跨设备通信：Rendezvous原理
   角色：轻量级消息队列，解耦Send/Recv节点。
   核心机制：

   • 唯一标识：ParsedKey = src_device + dst_device + tensor_id
   • 双缓冲队列：TableType table_ GUARDED_BY(mu_)
   • 异步回调：Recv未收到数据时注册回调，数据到达后触发。

   void IntraProcessRendezvous::RecvAsync(const ParsedKey& key, DoneCallback done) {
     if (auto* item = FindItem(key)) {
       done(Status::OK(), item->args, item->value); // 立即返回缓存数据
     } else {
       Table::InsertCallback(key, std::move(done)); // 注册异步回调
     }
   }
   

3. 结果收集

   • RetVal节点通过FunctionCallFrame.SetRetval()写入结果。
   • 所有Executor完成后，从FunctionCallFrame提取输出。

---

四、Rendezvous的工程价值

1. 设备异构性透明化
   统一处理CPU/GPU/ASIC设备间通信，对OP实现隐藏底层差异。
2. 死锁预防
   超时机制 + 协调器（Coordinator）确保跨设备等待不会阻塞系统。
3. 性能优化
   • 本地设备通信：零拷贝共享内存（IntraProcessRendezvous）。
   • 分布式通信：gRPC流 + RDMA加速（GrpcRemoteRendezvous）。

---

总结

DirectSession::Run的高效性源于三层核心设计：

1. 剪枝：最小化执行图，避免冗余计算。
2. 分裂：通过Send/Recv节点实现设备间解耦。
3. 并发：Executor + Rendezvous构建低延迟通信管道。

正是这些机制使得TensorFlow能在单一进程内高效调度异构设备资源，为分布式扩展奠定坚实基础。下一步可深入探索分布式运行时中GrpcSession与Master-Worker交互协议的设计奥秘。

Reference

1. tensorflow内核剖析