TF代码研读-Device(DeviceContext)

最新推荐文章于 2022-04-06 18:23:28 发布

橘子都吃不起！

最新推荐文章于 2022-04-06 18:23:28 发布

阅读量372

点赞数

分类专栏： Tensorflow源码研读文章标签： tensorflow 深度学习 C++

本文链接：https://blog.csdn.net/weixin_42754237/article/details/113471607

版权

Tensorflow源码研读专栏收录该内容

10 篇文章 10 订阅 ¥29.90 ¥99.00

订阅专栏

DeviceContext是TensorFlow中非CPU设备执行计算的核心类，尤其关注stream管理和张量拷贝。它为GPU、SYCL和XLA设备衍生出特定的上下文实现。在GPU上，DeviceContext表现为GPUDeviceContext，而XLA设备则对应XlaDeviceContext，后者着重实现了张量管理和流管理的定制化功能。在执行计算任务时，GPU等设备不直接调用底层接口，而是通过ThenExecute函数进行间接操作。

摘要由CSDN通过智能技术生成

DeviceContext是非CPU设备执行计算操作的核心类，重点提供了stream虚函数、以及张量拷贝相关函数。并且根据不同设备进一步派生出不同的设备上下文类（GPUDeviceContext、SYCLDeviceContext、XlaDeviceContext）。

与CPU设备不同 GPU设备(SYCL/XLA/RPC也一样)，设备在执行具体计算任务时并不是给Eigen底层计算的，而是由设备底层接口计算的(例如cuda)，当计算较多时，不同节点的与底层设备的交互并不是直接调用相关接口，而是通过传入执行函数的函数指针、执行所需的数据流、以及执行的设备信息到一个ThenExecute函数中，由ThenExecute函数执行具体的计算操作。

//设备的上下文信息类,注意继承自一个通用的计数器类,为一个抽象类
class DeviceContext : public core::RefCounted {
 public:
  ~DeviceContext() override {}
  virtual stream_executor::Stream* stream() const { return nullptr; }//TF流
  //维持流的生命周期
  virtual void MaintainLifetimeOnStream(const Tensor* t,
                                        stream_executor::Stream* stream) const {
  }
  /

了解本专栏