TensorFlow源码分析（7）：TensorFlow的DeviceBase类

在文章（5）中说过，TensorFlow在定义设备的时候，不是把它们分成CPU、GPU、TPU设备等类别，而是分成本地设备和远程设备这两个类别。这样的分类方式，也有其方便之处。因为有许多的运算，需要CPU和GPU协同工作。如果强硬地把它们放在两个类中，会带来许多麻烦。而分成本地设备和远程设备，实现逻辑就非常清晰明了。本地设备只需调用底层的库函数（如Eigen, CUDA，SYCL, OpenCL等），而远程设备是通过RPC和Proto buffer进行通信。

所以的设备都必须继承自DeviceBase类，它定义了一些基本的数据结构与接口。首先分析DeviceBase类的成员变量。参见tensorflow/core/framework/device_base.h。

class DeviceBase {

private:
  Env* const env_;
  CpuWorkerThreads* cpu_worker_threads_ = nullptr;
  // Set by GPUs as well as by TPU devices.
  GpuDeviceInfo* gpu_device_info_ = nullptr;
  thread::ThreadPool* device_thread_pool_ = nullptr;
  std::vector<Eigen::ThreadPoolDevice*> eigen_cpu_devices_;

• env_对象定义了与操作系统相关的一些方法。
• cpu_worker_threads_对象定义了一个线程池，在底层调用了Eigen的线程池管理机制。默认使用的是文章（6）介绍的SimpleThreadPool。
• gpu_device_info_对象是一个与GPU相关的结构体。参考上面的注释可以得悉，这个结构体其实也可以管理TPU设备。device_thread_pool_对象默认也是SimpleThreadPool类。但是与cpu_worker_threads_不同的是，device_thread_pool_是专门为GPU服务的。
• eigen_cpu_devices_对象是Eigen定义的ThreadPoolDevice类，该类在eigen的thread pool基础上做了封装。它实现了一些简单的内存管理接口；同时，最重要地是实现了一个ParallelFor函数，该函数可以自动将一个任务切分成多个block，以实现高效的并行计算（block的概念类似于GPU的block）。注意cpu_worker_threads_和eigen_cpu_devices_对象底层都依赖于Eigen的线程池，都为了实现CPU的多线程计算。但是cpu_worker_threads_是为TensorFlow实现自己的线程管理机制服务，而eigen_cpu_devices_是直接使用Eigen的线程管理机制。在TensorFlow的代码中，会根据任务的不同，调用不同的线程管理。

DeviceBase除了定义了上述成员变量外，还有一些接口函数。这在介绍DeviceBase子类时再慢慢介绍。