Chenfan Blog

Do cool things that matter.

0%

TensorFlow 拆包（四）：Device

Posted on Edited on In

接上篇：

这篇要分析的是 TensorFlow 中跟计算设备相关的内容。

DeviceFactory::AddDevices()

TensorFlow 在创建 Session 时会首先扫描当前设备中存在哪些可用的计算资源（CPU、GPU）。

Python 层的 Session 构建时通过 tf_session.TF_NewDeprecatedSession() 这个接口调用 C 层的运行时来在 C 层新建 Session。

单节点下创建 DirectSession 由 DirectSessionFactory::NewSession()完成，其中又通过 DeviceFactory::AddDevices()获取当前可用的设备信息。

不同的计算设备分别注册各自的 DeviceFactory，AddDevices()调用不同设备的CreateDevices()，最终将所有的 device 信息传递回上层。

这里device_factories()中是一个静态的 std::unordered_map<string, FactoryItem>，比较奇怪的是我找了半天没找到这个东西是在哪赋值的，gdb 调试进去，明明第一次调用的时候应该是创建一个空的新对象，结果里面居然就有内容了 ：

1
2
3
4
std::unordered_map with 2 elements = {
    ["CPU"] = {factory = std::unique_ptr<tensorflow::DeviceFactory> containing 0x15fcf00, priority = 70},
    ["GPU"] = {factory = std::unique_ptr<tensorflow::DeviceFactory> containing 0x15f2de0, priority = 210}
}

GPU 版编译的 TensorFlow 中，CPU 端对应的是tensorflow::GPUCompatibleCPUDeviceFactory::CreateDevices，GPU 端对应的是 tensorflow::GPUDeviceFactory::CreateDevices。

这两个 DeviceFactory 的定义在 /tensorflow/core/common_runtime/gpu/gpu_device_factory.cc 中，然后在这里面还比较惊喜地发现了两条与前面几篇中注册 Op 时很像的宏：

1
2
REGISTER_LOCAL_DEVICE_FACTORY("GPU", GPUDeviceFactory, 210);
REGISTER_LOCAL_DEVICE_FACTORY("CPU", GPUCompatibleCPUDeviceFactory, 70);

这应该就是用于注册 Device 到前面的std::unordered_map<string, FactoryItem>结构中的内容了，但是关于注册这个过程是在什么时候运行的我还不是很明确，因为 gdb 在这里加了断点却没有进。

在整个 TensorFlow 的源码中，除了上面的两条以外，还可以找到：

1
2
REGISTER_LOCAL_DEVICE_FACTORY("CPU", ThreadPoolDeviceFactory, 60);
REGISTER_LOCAL_DEVICE_FACTORY("SYCL", SYCLDeviceFactory, 200);

两种设备，ThreadPoolDeviceFactory应该是用于纯 CPU 版的 TensorFlow。事实上GPUCompatibleCPUDeviceFactory就是继承的ThreadPoolDeviceFactory，额外加了一些与 GPU 结合的选项。

目前ThreadPoolDeviceFactory和GPUDevice最终都是用到的LocalDevice，配合 Eigen 使用。

AddDevices()得到的 device 列表存在 DeviceMgr 中传入 DirectSession。

Add a new type of device to TF

官方在这块内容基本上没有给什么太详细的说明，不过从前面的分析也可以很容易看出来，要创建一个新类型的设备首先要继承一个新的 tensorflow::Device，以及其生成用的工厂模式tensorflow::DeviceFactory，中间涉及到一些必要的函数需要重载，然后用REGISTER_LOCAL_DEVICE_FACTORY(...)注册即可。

这里有一个用 CPU 改个名字来虚拟新硬件的测试例子。

Allocate nodes with Devices

Placer::Run()

tensorflow::GraphExecutionState::Extend()创建完整的运行图时，用了一个tensorflow::Placer结构来处理图中的运行节点与设备的关联问题。

这里的分配策略非常简单，用一个并查集（！！666！！）来维护所有节点的连通性，然后把连通在一起的节点分到同样的设备上。其中有手动指定的话就按手动指定的来，没有手动指定的则按 device 优先级来，默认 GPU 最高，然后是 OpenCL 的 SYCL，最后才是分配到 CPU 上。

这里有几个额外规则：

  1. source node 和 sink node 必须分到 CPU 上
  2. 对于没有输入，有一个输出的 GeneratorNode，分配到它的目标节点所在的设备上；
  3. 对于直接在原数据上进行操作（例如说 reshape）这样的 MetadataNode，分配到它的源节点所在的设备上。

tensorflow::Partition()

上篇中，在DirectSession::GetOrCreateExecutors()的过程中，用DirectSession::CreateGraphs()根据当前的输入、输出等信息从完整图中创建出了一个当前运行所用的子图。

之后需要用到tensorflow::Partition()来完成需要运行的子图与设备的分配关联：

  • 为图中的每一个节点和边创建内存和设备类型信息
  • 检查每一个节点的输入输出、以及其目标节点的输入信息等
  • 给每一个节点上添加控制边等，如果数据传输不在同一个设备上就添加一对 send/recv 的 node

Devices & Compute

事实上，感觉 DirectSession 中创建的很多结构都像是有一一对应关系的，比如前面Partition()得到的子图、子图所关联的设备、以及子图所构建的 Executor。

在tensorflow/core/common_runtime/device.h中，可以找到 Device 类的计算函数的定义：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// Performs the actual compute function.
//
// Subclasses may override this function if they wish to perform
// some initialization before each compute.
virtual void Compute(OpKernel* op_kernel, OpKernelContext* context) {
  op_kernel->Compute(context);
}

// Asynchronous kernel's compute.
virtual void ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context,
                          AsyncOpKernel::DoneCallback done) {
  op_kernel->ComputeAsync(context, std::move(done));
}

其实就是运行传入的 Op 所注册的 OpKernel 函数。

ThreadPoolDevice 类重载了这两个函数，加上一些额外需要记录的信息，核心部分还是运行 OpKernel。

Summary

这里把前面几篇中的涉及到的内容稍微做一下总结：

后续：

Related Posts