背景
第一步——分析构建Control Flow相关信息
这个过程在代码中是通过AddControlFlow函数实现的,由于改代码深度依赖于Control Flow Op的相关模块,且对于不含有Control Flow Op的Graph几乎什么都没有做,因此我们先忽略这个过程,等到对Control Flow模块做详细解读时再回过头来研究其在Graph Partitioner中的意义。
1 GraphInfo g_info; 2 if (!opts.control_flow_added) { 3 // Add the "code" for distributed execution of control flow. Code is 4 // added only for the frames that are placed on multiple devices. The 5 // new graph is an equivalent transformation of the original graph and 6 // has the property that it can be subsequently partitioned arbitrarily 7 // (down to the level of individual device) for distributed execution. 8 status = AddControlFlow(opts, g, &g_info); 9 if (!status.ok()) return status; 10 }
第二步——构建Op的Input和Output Memory类型信息
在介绍这个过程之前,首先需要明确两种概念,他们是DeviceMemory和HostMemory。前者指的是计算设备的Memory类型,后者指的是CPU的Memory类型,它们在TensorFlow中被定义为Enum类型,代码如下所示。
1 // MemoryType is used to describe whether input or output Tensors of 2 // an OpKernel should reside in "Host memory" (e.g., CPU memory) or 3 // "Device" Memory (CPU memory for CPU devices, GPU memory for GPU 4 // devices).5 enum MemoryType { 6 DEVICE_MEMORY = 0, 7 HOST_MEMORY = 1, 8 };
对Op的Input和Output Memory信息进行检索并构建缓存的函数是BuildMemoryDeviceInfo,该过程构建的信息对后面真正做图切割非常重要。因为TensorFlow的Op在注册时需要不但需要指定其在各个Device上的实现版本(比如CPU版本的Op和GPU版本的Op都是分别注册到系统中的),还需要指出其Input和Output Tensor的类型以及所使用的Memory类型,即使某个Op存在GPU上的实现,它的GPU版本也有可能需要在CPU上读入数据或输出结果。例如,GPU版本的Reshape Op注册代码如下。
1 #define REGISTER_GPU_KERNEL(type) \ 2 REGISTER_KERNEL_BUILDER(Name("Reshape") \ 3 .Device(DEVICE_GPU) \ 4 .HostMemory("shape") \ 5 .TypeConstraint<type>("

