TensorFlow 拆包(四):Device

接上篇:

这篇要分析的是 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 结合的选项。


目前ThreadPoolDeviceFactoryGPUDevice最终都是用到的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
// 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

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


后续:

0%